Replace torch einsum with opt_einsum

Differential Revision: D37128344

fbshipit-source-id: ff1105b150dfa023b57a7eec96fb5fa8ae58f1c9

Author: sayanghosh

opacus/grad_sample/conv.py

@@ -19,6 +19,7 @@
 import torch
 import torch.nn as nn
 from opacus.utils.tensor_utils import unfold2d, unfold3d
+from opt_einsum import contract
 
 from .utils import register_grad_sampler
 

opacus/grad_sample/dp_rnn.py

@@ -19,6 +19,7 @@
 import torch
 import torch.nn as nn
 from opacus.layers.dp_rnn import RNNLinear
+from opt_einsum import contract
 
 from .utils import register_grad_sampler
 
@@ -40,8 +41,8 @@ def compute_rnn_linear_grad_sample(
     """
     ret = {}
     if layer.weight.requires_grad:
-        gs = torch.einsum("n...i,n...j->nij", backprops, activations)
+        gs = contract("n...i,n...j->nij", backprops, activations)
         ret[layer.weight] = gs
     if layer.bias is not None and layer.bias.requires_grad:
-        ret[layer.bias] = torch.einsum("n...k->nk", backprops)
+        ret[layer.bias] = contract("n...k->nk", backprops)
     return ret

opacus/grad_sample/group_norm.py

@@ -19,6 +19,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from opt_einsum import contract
 
 from .utils import register_grad_sampler
 
@@ -40,7 +41,7 @@ def compute_group_norm_grad_sample(
     ret = {}
     if layer.weight.requires_grad:
         gs = F.group_norm(activations, layer.num_groups, eps=layer.eps) * backprops
-        ret[layer.weight] = torch.einsum("ni...->ni", gs)
+        ret[layer.weight] = contract("ni...->ni", gs)
     if layer.bias is not None and layer.bias.requires_grad:
-        ret[layer.bias] = torch.einsum("ni...->ni", backprops)
+        ret[layer.bias] = contract("ni...->ni", backprops)
     return ret

opacus/grad_sample/instance_norm.py

@@ -18,6 +18,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from opt_einsum import contract
 
 from .utils import register_grad_sampler
 
@@ -49,7 +50,7 @@ def compute_instance_norm_grad_sample(
     ret = {}
     if layer.weight.requires_grad:
         gs = F.instance_norm(activations, eps=layer.eps) * backprops
-        ret[layer.weight] = torch.einsum("ni...->ni", gs)
+        ret[layer.weight] = contract("ni...->ni", gs)
     if layer.bias is not None and layer.bias.requires_grad:
-        ret[layer.bias] = torch.einsum("ni...->ni", backprops)
+        ret[layer.bias] = contract("ni...->ni", backprops)
     return ret

opacus/grad_sample/linear.py

@@ -17,6 +17,7 @@
 
 import torch
 import torch.nn as nn
+from opt_einsum import contract
 
 from .utils import register_grad_sampler
 
@@ -35,8 +36,8 @@ def compute_linear_grad_sample(
     """
     ret = {}
     if layer.weight.requires_grad:
-        gs = torch.einsum("n...i,n...j->nij", backprops, activations)
+        gs = contract("n...i,n...j->nij", backprops, activations)
         ret[layer.weight] = gs
     if layer.bias is not None and layer.bias.requires_grad:
-        ret[layer.bias] = torch.einsum("n...k->nk", backprops)
+        ret[layer.bias] = contract("n...k->nk", backprops)
     return ret

opacus/optimizers/adaclipoptimizer.py

@@ -18,17 +18,17 @@
 from typing import Callable, Optional
 
 import torch
+from opt_einsum import contract
 from torch.optim import Optimizer
 
 from .optimizer import (
-    DPOptimizer,
     _check_processed_flag,
     _generate_noise,
     _get_flat_grad_sample,
     _mark_as_processed,
+    DPOptimizer,
 )
 
-
 logger = logging.getLogger(__name__)
 
 
@@ -108,7 +108,7 @@ def clip_and_accumulate(self):
             _check_processed_flag(p.grad_sample)
 
             grad_sample = _get_flat_grad_sample(p)
-            grad = torch.einsum("i,i...", per_sample_clip_factor, grad_sample)
+            grad = contract("i,i...", per_sample_clip_factor, grad_sample)
 
             if p.summed_grad is not None:
                 p.summed_grad += grad

opacus/optimizers/ddp_perlayeroptimizer.py

@@ -18,17 +18,18 @@
 from typing import Callable, List, Optional
 
 import torch
+from opt_einsum import contract
 from torch import nn
 from torch.optim import Optimizer
 
-from .optimizer import DPOptimizer, _generate_noise
+from .optimizer import _generate_noise, DPOptimizer
 
 
 def _clip_and_accumulate_parameter(p: nn.Parameter, max_grad_norm: float):
     per_sample_norms = p.grad_sample.view(len(p.grad_sample), -1).norm(2, dim=-1)
     per_sample_clip_factor = (max_grad_norm / (per_sample_norms + 1e-6)).clamp(max=1.0)
 
-    grad = torch.einsum("i,i...", per_sample_clip_factor, p.grad_sample)
+    grad = contract("i,i...", per_sample_clip_factor, p.grad_sample)
     if p.summed_grad is not None:
         p.summed_grad += grad
     else:

opacus/optimizers/optimizer.py

@@ -19,10 +19,10 @@
 
 import torch
 from opacus.optimizers.utils import params
+from opt_einsum import contract
 from torch import nn
 from torch.optim import Optimizer
 
-
 logger = logging.getLogger(__name__)
 
 
@@ -404,7 +404,7 @@ def clip_and_accumulate(self):
             _check_processed_flag(p.grad_sample)
 
             grad_sample = _get_flat_grad_sample(p)
-            grad = torch.einsum("i,i...", per_sample_clip_factor, grad_sample)
+            grad = contract("i,i...", per_sample_clip_factor, grad_sample)
 
             if p.summed_grad is not None:
                 p.summed_grad += grad

opacus/optimizers/perlayeroptimizer.py

@@ -18,9 +18,10 @@
 
 import torch
 from opacus.optimizers.utils import params
+from opt_einsum import contract
 from torch.optim import Optimizer
 
-from .optimizer import DPOptimizer, _check_processed_flag, _mark_as_processed
+from .optimizer import _check_processed_flag, _mark_as_processed, DPOptimizer
 
 
 class DPPerLayerOptimizer(DPOptimizer):
@@ -61,7 +62,7 @@ def clip_and_accumulate(self):
             per_sample_clip_factor = (max_grad_norm / (per_sample_norms + 1e-6)).clamp(
                 max=1.0
             )
-            grad = torch.einsum("i,i...", per_sample_clip_factor, p.grad_sample)
+            grad = contract("i,i...", per_sample_clip_factor, p.grad_sample)
 
             if p.summed_grad is not None:
                 p.summed_grad += grad

opacus/tests/privacy_engine_test.py

@@ -29,9 +29,10 @@
 from opacus import PrivacyEngine
 from opacus.layers.dp_multihead_attention import DPMultiheadAttention
 from opacus.optimizers.optimizer import _generate_noise
-from opacus.scheduler import StepNoise, _NoiseScheduler
+from opacus.scheduler import _NoiseScheduler, StepNoise
 from opacus.utils.module_utils import are_state_dict_equal
 from opacus.validators.errors import UnsupportedModuleError
+from opt_einsum import contract
 from torch.utils.data import DataLoader, Dataset, TensorDataset
 from torchvision import models, transforms
 from torchvision.datasets import FakeData
@@ -47,7 +48,7 @@ def get_grad_sample_aggregated(tensor: torch.Tensor, loss_type: str = "mean"):
     if loss_type not in ("sum", "mean"):
         raise ValueError(f"loss_type = {loss_type}. Only 'sum' and 'mean' supported")
 
-    grad_sample_aggregated = torch.einsum("i...->...", tensor.grad_sample)
+    grad_sample_aggregated = contract("i...->...", tensor.grad_sample)
     if loss_type == "mean":
         b_sz = tensor.grad_sample.shape[0]
         grad_sample_aggregated /= b_sz