Merge pull request #26 from lxuechen/v0.2.2

V0.2.2; fixes #25

Author: lxuechen

private_transformers/__init__.py

@@ -1,4 +1,4 @@
 from .privacy_engine import PrivacyEngine
 from .transformers_support import freeze_isolated_params_for_vit
 
-__version__ = '0.2.1'
+__version__ = '0.2.2'

private_transformers/autograd_grad_sample.py

@@ -52,20 +52,6 @@ def requires_grad(module: nn.Module, recurse: bool = False) -> bool:
     return requires_grad
 
 
-def get_layer_type(layer: nn.Module) -> str:
-    """
-    Returns the name of the type of the given layer.
-
-    Args:
-        layer: The module corresponding to the layer whose type
-            is being queried.
-
-    Returns:
-        Name of the class of the layer
-    """
-    return layer.__class__.__name__
-
-
 def add_hooks(model: nn.Module, loss_reduction: str = "mean"):
     r"""
     Adds hooks to model to save activations and backprop values.
@@ -86,7 +72,7 @@ def add_hooks(model: nn.Module, loss_reduction: str = "mean"):
 
     handles = []
     for name, layer in model.named_modules():
-        if get_layer_type(layer) in _supported_layers_grad_samplers.keys():
+        if type(layer) in _supported_layers_grad_samplers:
             # Check if the layer has trainable parameters.
             is_trainable = False
             for p in layer.parameters(recurse=False):

private_transformers/privacy_engine.py

@@ -248,6 +248,11 @@ def step(
         # This option was included to help with another spectrum analysis project.
         callback: Optional[Callable] = None,
     ):
+        if loss.dim() != 1:
+            raise ValueError(
+                f"Expected `loss` to be the per-example loss 1-D tensor, but got a tensor with dims={loss.dim()}."
+            )
+
         if self.clipping_mode == ClippingMode.ghost:
             if callback is not None:
                 raise ValueError("Ghost clipping does not support `callback` in `optimizer.step`.")
@@ -359,11 +364,6 @@ def _ghost_virtual_step(self, loss: torch.Tensor):
     @torch.enable_grad()
     def _double_backward(self, loss: torch.Tensor):
         """Given per-example losses, backward twice to accumulate summed clipped gradients in `.grad`."""
-        if loss.dim() != 1:
-            raise ValueError(
-                f"Expected `loss` to be the per-example loss 1-D tensor, but got a tensor with dims={loss.dim()}."
-            )
-
         first_loss = loss.sum()
         first_loss.backward(retain_graph=True)
 
@@ -437,9 +437,6 @@ def _accumulate_summed_grad(self, loss, scale):
 
         Removes `.grad_sample` and `.grad` for each variable that requires grad at the end.
         """
-        if loss.dim() != 1:
-            raise ValueError(f"Expected `loss` to be a the per-example loss 1-D tensor.")
-
         with torch.enable_grad():
             loss.sum(dim=0).backward()
 
@@ -466,12 +463,8 @@ def _accumulate_summed_grad(self, loss, scale):
             for tensor in norm_sample:
                 shapes[tensor.size()] += 1
 
-            major_shape = None
-            major_count = 0
-            for shape, count in shapes.items():
-                if count > major_count:
-                    major_shape = shape
-            del shape, count
+            # Get the shape that most tensors have.
+            major_shape, major_count = max(shapes.items(), key=lambda x: x[1])
 
             # Check which tensors don't have the major shape!
             extra_msg = f" \n*** Major shape: {major_shape}"

private_transformers/supported_layers_grad_samplers.py

@@ -16,6 +16,7 @@
 
 import numpy as np
 import torch
+import transformers.pytorch_utils
 from opt_einsum import contract
 from torch import nn
 from torch.functional import F
@@ -313,11 +314,11 @@ def _compute_conv2d_grad_sample(layer: nn.Conv2d, activations: Tuple[torch.Tenso
 
 
 _supported_layers_grad_samplers = {
-    "Embedding": _compute_embedding_grad_sample,
-    "Linear": _compute_linear_grad_sample,
-    "Conv2d": _compute_conv2d_grad_sample,  # nn.Conv2d.
-    "LayerNorm": _compute_layer_norm_grad_sample,
-    "Conv1D": _custom_compute_conv1d_grad_sample,  # HuggingFace Open-AI GPT-2.
-    "T5LayerNorm": _compute_t5_layer_norm_grad_sample,
-    "OPTLearnedPositionalEmbedding": _compute_opt_learned_positional_embedding_grad_sample,
+    nn.Embedding: _compute_embedding_grad_sample,
+    nn.Linear: _compute_linear_grad_sample,
+    nn.Conv2d: _compute_conv2d_grad_sample,
+    nn.LayerNorm: _compute_layer_norm_grad_sample,
+    transformers.pytorch_utils.Conv1D: _custom_compute_conv1d_grad_sample,
+    transformers.models.t5.modeling_t5.T5LayerNorm: _compute_t5_layer_norm_grad_sample,
+    OPTLearnedPositionalEmbedding: _compute_opt_learned_positional_embedding_grad_sample,
 }

setup.py

@@ -34,6 +34,7 @@
         "jupyter",
         "ml-swissknife",
         "opt_einsum",
+        "pytest"
     ],
     python_requires='~=3.8',
     classifiers=[