update all tutorials

Author: Igor Shilov

opacus/tests/batch_memory_manager_test.py

@@ -62,10 +62,19 @@ def test_basic(
             max_grad_norm=1.0,
             poisson_sampling=False,
         )
+        max_physical_batch_size = 3
         with BatchMemoryManager(
-            data_loader=data_loader, max_physical_batch_size=3, optimizer=optimizer
+            data_loader=data_loader,
+            max_physical_batch_size=max_physical_batch_size,
+            optimizer=optimizer,
         ) as new_data_loader:
-            self.assertEqual(len(data_loader), len(new_data_loader))
+            self.assertEqual(
+                len(data_loader), len(data_loader.dataset) // self.batch_size
+            )
+            self.assertEqual(
+                len(new_data_loader),
+                len(data_loader.dataset) // max_physical_batch_size,
+            )
             weights_before = torch.clone(model._module.fc.weight)
             for i, (x, y) in enumerate(new_data_loader):
                 self.assertTrue(x.shape[0] <= 3)

opacus/utils/batch_memory_manager.py

@@ -3,8 +3,12 @@
 
 import numpy as np
 from opacus.optimizers import DPOptimizer
-from torch.utils.data import DataLoader, Sampler, BatchSampler
-from opacus.utils.uniform_sampler import UniformWithReplacementSampler, DistributedUniformWithReplacementSampler
+from opacus.utils.uniform_sampler import (
+    DistributedUniformWithReplacementSampler,
+    UniformWithReplacementSampler,
+)
+from torch.utils.data import BatchSampler, DataLoader, Sampler
+
 
 class BatchSplittingSampler(Sampler[List[int]]):
     def __init__(
@@ -27,13 +31,18 @@ def __iter__(self):
 
     def __len__(self):
         if isinstance(self.sampler, BatchSampler):
-            return int(len(self.sampler) * (self.sampler.batch_size / self.max_batch_size))
-        elif isinstance(self.sampler, UniformWithReplacementSampler) or isinstance(self.sampler, DistributedUniformWithReplacementSampler):
+            return int(
+                len(self.sampler) * (self.sampler.batch_size / self.max_batch_size)
+            )
+        elif isinstance(self.sampler, UniformWithReplacementSampler) or isinstance(
+            self.sampler, DistributedUniformWithReplacementSampler
+        ):
             expected_batch_size = self.sampler.sample_rate * self.sampler.num_samples
             return int(len(self.sampler) * (expected_batch_size / self.max_batch_size))
 
         return len(self.sampler)
 
+
 def wrap_data_loader(data_loader, max_batch_size: int, optimizer: DPOptimizer):
     return DataLoader(
         dataset=data_loader.dataset,

opacus/validators/batch_norm.py

@@ -2,14 +2,14 @@
 # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
 
 import logging
-import sys
 from typing import Union
 
 import torch.nn as nn
 
 from .errors import ShouldReplaceModuleError, UnsupportableModuleError
 from .utils import register_module_fixer, register_module_validator
 
+
 logger = logging.getLogger(__name__)
 
 BATCHNORM = Union[nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d, nn.SyncBatchNorm]

opacus/validators/module_validator.py

@@ -2,7 +2,6 @@
 # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
 
 import logging
-import sys
 from typing import List
 
 import torch.nn as nn
@@ -13,8 +12,10 @@
     UnsupportedModuleError,
 )
 
+
 logger = logging.getLogger(__name__)
 
+
 class ModuleValidator:
     """
     Encapsulates all the validation logic required by Opacus.
@@ -49,7 +50,7 @@ def validate(
         # 2. validate that all trainable modules are supported by GradSampleModule.
         errors.extend(GradSampleModule.validate(module=module, raise_if_error=False))
         # 3. perform module specific validations.
-        #TODO: use module name here - it's useful part of error message
+        # TODO: use module name here - it's useful part of error message
         for _, sub_module in module.named_modules():
             if type(sub_module) in ModuleValidator.VALIDATORS:
                 sub_module_validator = ModuleValidator.VALIDATORS[type(sub_module)]