Incorporate code review feedback on PR #694

Author: pagarwl

opacus/grad_sample/embedding.py

@@ -15,10 +15,10 @@
 
 from typing import Dict, List
 
-from opacus.grad_sample import embedding_norm_sample
 import torch
 import torch.nn as nn
 
+from opacus.grad_sample import embedding_norm_sample
 from .utils import register_grad_sampler, register_norm_sampler
 
 

opacus/grad_sample/embedding_norm_sample.py

@@ -46,12 +46,10 @@ def compute_embedding_norm_sample(
       activations:  [tensor([[1, 1],
           [2, 0],
           [2, 0]])]
-      backprops:  tensor([[0.2000],
-          [0.2000],
-          [0.3000],
-          [0.1000],
-          [0.3000],
-          [0.1000]])
+      backprops:  tensor([[[0.2], [0.2]],
+          [[0.3], [0.1]],
+          [[0.3], [0.1]]])
+      backprops.shape:  torch.Size([3, 2, 1])
 
     Intermediate values:
       input_ids:  tensor([[1, 1],

opacus/tests/grad_sample_module_fast_gradient_clipping_test.py

@@ -351,7 +351,7 @@ def test_norm_calculation(self):
         diff = flat_norms_normal - flat_norms_gc
 
         logging.info(f"Diff = {diff}")
-        msg = "Fail: Gradients from vanilla DP-SGD and from fast gradient clipping are different"
+        msg = "Fail: Gradient norms from vanilla DP-SGD and from fast gradient clipping are different"
         assert torch.allclose(flat_norms_normal, flat_norms_gc, atol=1e-3), msg
 
     def test_gradient_calculation(self):

opacus/tests/grad_samples/embedding_norm_sample_test.py

@@ -15,9 +15,9 @@
 
 import unittest
 
-from opacus.grad_sample import embedding_norm_sample
 import torch
 import torch.nn as nn
+from opacus.grad_sample import embedding_norm_sample
 
 
 class TestComputeEmbeddingNormSample(unittest.TestCase):
@@ -36,15 +36,11 @@ def test_compute_embedding_norm_sample(self):
         # Example input ids (activations). Shape: [3, 2]
         input_ids = torch.tensor([[1, 1], [2, 0], [2, 0]], dtype=torch.long)
 
-        # Example gradients with respect to the embedding output (backprops).
-        # Shape: [6, 1]
-        grad_values = torch.tensor(
-            [[0.2], [0.2], [0.3], [0.1], [0.3], [0.1]], dtype=torch.float32
+        # Example backprops. Shape: [3, 2, 1]
+        backprops = torch.tensor(
+            [[[0.2], [0.2]], [[0.3], [0.1]], [[0.3], [0.1]]], dtype=torch.float32
         )
 
-        # Simulate backprop through embedding layer
-        backprops = grad_values
-
         # Wrap input_ids in a list as expected by the norm sample function
         activations = [input_ids]
 
@@ -70,17 +66,17 @@ def test_compute_embedding_norm_sample_with_non_one_embedding_dim(self):
 
         # Manually set weights for the embedding layer for testing
         embedding_layer.weight = nn.Parameter(
-            torch.tensor([[0.1], [0.2], [0.3]], dtype=torch.float32)
+            torch.tensor([[0.1, 0.1], [0.2, 0.2], [0.3, 0.3]], dtype=torch.float32)
         )
 
         # Example input ids (activations). Shape: [6, 1, 1].
         input_ids = torch.tensor(
             [[[1]], [[1]], [[2]], [[0]], [[2]], [[0]]], dtype=torch.long
         )
 
-        # Example gradients per input id, with embedding_dim=2.
+        # Example backprops per input id, with embedding_dim=2.
         # Shape: [6, 1, 1, 2]
-        grad_values = torch.tensor(
+        backprops = torch.tensor(
             [
                 [[[0.2, 0.2]]],
                 [[[0.2, 0.2]]],
@@ -92,9 +88,6 @@ def test_compute_embedding_norm_sample_with_non_one_embedding_dim(self):
             dtype=torch.float32,
         )
 
-        # Simulate backprop through embedding layer
-        backprops = grad_values
-
         # Wrap input_ids in a list as expected by the grad norm function
         activations = [input_ids]
 
@@ -211,7 +204,6 @@ def test_compute_embedding_norm_sample_with_extra_activations_per_example(self):
         expected_norms = torch.tensor(
             [0.0150, 0.0071, 0.0005, 0.0081, 0.0039], dtype=torch.float32
         )
-        print("expected_norms: ", expected_norms)
         computed_norms = result[embedding_layer.weight]
 
         # Verify the computed norms match the expected norms