diff --git a/opacus/optimizers/adaclipoptimizer.py b/opacus/optimizers/adaclipoptimizer.py index 7144f06b..4145c20e 100644 --- a/opacus/optimizers/adaclipoptimizer.py +++ b/opacus/optimizers/adaclipoptimizer.py @@ -43,16 +43,13 @@ def __init__( optimizer: Optimizer, *, noise_multiplier: float, - target_unclipped_quantile: float, - clipbound_learning_rate: float, - max_clipbound: float, - min_clipbound: float, - unclipped_num_std: float, max_grad_norm: float, expected_batch_size: Optional[int], loss_reduction: str = "mean", generator=None, secure_mode: bool = False, + normalize_clipping: bool = False, + optim_args: dict = None, ): super().__init__( optimizer, @@ -62,7 +59,13 @@ def __init__( loss_reduction=loss_reduction, generator=generator, secure_mode=secure_mode, + normalize_clipping=normalize_clipping, ) + target_unclipped_quantile = optim_args.get('target_unclipped_quantile', 0.0) + clipbound_learning_rate = optim_args.get('clipbound_learning_rate', 1.0) + max_clipbound = optim_args.get('max_clipbound', torch.inf) + min_clipbound = optim_args.get('min_clipbound', -torch.inf) + unclipped_num_std = optim_args.get('unclipped_num_std') assert ( max_clipbound > min_clipbound ), "max_clipbound must be larger than min_clipbound." diff --git a/opacus/privacy_engine.py b/opacus/privacy_engine.py index de210b6e..92a87a40 100644 --- a/opacus/privacy_engine.py +++ b/opacus/privacy_engine.py @@ -18,6 +18,7 @@ from typing import IO, Any, BinaryIO, Dict, List, Optional, Tuple, Union import torch +from torch import distributed as dist from opacus.accountants import create_accountant from opacus.accountants.utils import get_noise_multiplier from opacus.data_loader import DPDataLoader, switch_generator @@ -111,6 +112,7 @@ def _prepare_optimizer( noise_generator=None, grad_sample_mode="hooks", normalize_clipping: bool = False, + optim_args: dict = None, **kwargs, ) -> DPOptimizer: if isinstance(optimizer, DPOptimizer): @@ -294,6 +296,7 @@ def make_private( grad_sample_mode: str = "hooks", normalize_clipping: bool = False, total_steps: int = None, + optim_args: dict = None, **kwargs, ) -> Tuple[GradSampleModule, DPOptimizer, DataLoader]: """ @@ -375,7 +378,7 @@ def make_private( "Module parameters are different than optimizer Parameters" ) - distributed = isinstance(module, (DPDDP, DDP)) + distributed = dist.get_world_size() > 1 module = self._prepare_model( module, @@ -427,6 +430,7 @@ def make_private( clipping=clipping, grad_sample_mode=grad_sample_mode, normalize_clipping=normalize_clipping, + optim_args=optim_args, **kwargs, ) @@ -454,6 +458,7 @@ def make_private_with_epsilon( grad_sample_mode: str = "hooks", normalize_clipping: bool = False, total_steps: int = None, + optim_args: dict = None, **kwargs, ): """ @@ -569,6 +574,7 @@ def make_private_with_epsilon( clipping=clipping, normalize_clipping=normalize_clipping, total_steps=total_steps, + optim_args=optim_args, ) def get_epsilon(self, delta):