Replace torch einsum with opt_einsum (#440)

Summary:
Pull Request resolved: #440

We are using optimized einsums in place of Pytorch einsums. As per https://optimized-einsum.readthedocs.io/en/stable/ opt einsums are faster and our results on Opacus benchmarking also corroborate it.

Differential Revision: D37128344

fbshipit-source-id: 891c1cc3e1348a4965a068d6fd1375eb584805b9

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
 
@@ -70,7 +71,7 @@ def compute_conv_grad_sample(
     ret = {}
     if layer.weight.requires_grad:
         # n=batch_sz; o=num_out_channels; p=(num_in_channels/groups)*kernel_sz
-        grad_sample = torch.einsum("noq,npq->nop", backprops, activations)
+        grad_sample = contract("noq,npq->nop", backprops, activations)
         # rearrange the above tensor and extract diagonals.
         grad_sample = grad_sample.view(
             n,
@@ -80,7 +81,7 @@ def compute_conv_grad_sample(
             int(layer.in_channels / layer.groups),
             np.prod(layer.kernel_size),
         )
-        grad_sample = torch.einsum("ngrg...->ngr...", grad_sample).contiguous()
+        grad_sample = contract("ngrg...->ngr...", grad_sample).contiguous()
         shape = [n] + list(layer.weight.shape)
         ret[layer.weight] = grad_sample.view(shape)
 

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,6 +18,7 @@
 from typing import Callable, Optional
 
 import torch
+from opt_einsum import contract
 from torch.optim import Optimizer
 
 from .optimizer import (
@@ -108,7 +109,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,6 +18,7 @@
 from typing import Callable, List, Optional
 
 import torch
+from opt_einsum import contract
 from torch import nn
 from torch.optim import Optimizer
 
@@ -28,7 +29,7 @@ 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

@@ -17,8 +17,10 @@
 import logging
 from typing import Callable, List, Optional, Union
 
-import torch
 from opacus.optimizers.utils import params
+
+import torch
+from opt_einsum import contract
 from torch import nn
 from torch.optim import Optimizer
 
@@ -404,7 +406,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

@@ -16,8 +16,10 @@
 
 from typing import List, Optional
 
-import torch
 from opacus.optimizers.utils import params
+
+import torch
+from opt_einsum import contract
 from torch.optim import Optimizer
 
 from .optimizer import DPOptimizer, _check_processed_flag, _mark_as_processed
@@ -61,7 +63,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

@@ -33,6 +33,7 @@
 from opacus.utils.module_utils import are_state_dict_equal
 from opacus.validators.errors import UnsupportedModuleError
 from opacus.validators.module_validator import ModuleValidator
+from opt_einsum import contract
 from torch.utils.data import DataLoader, Dataset, TensorDataset
 from torchvision import models, transforms
 from torchvision.datasets import FakeData
@@ -48,7 +49,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