Add **kwargs to all optimizer classes (#710)

Summary:
Pull Request resolved: #710

Purpose: To enable creating custom PrivacyEngines that extend the PrivacyEngine class and take in additional parameters.

Fix prior diff: D67456352

Reviewed By: HuanyuZhang

Differential Revision: D67953655

fbshipit-source-id: 70aef7571e012a370d6a0fd04948eccee06c9a0d

opacus/optimizers/adaclipoptimizer.py

@@ -53,6 +53,7 @@ def __init__(
         loss_reduction: str = "mean",
         generator=None,
         secure_mode: bool = False,
+        **kwargs,
     ):
         super().__init__(
             optimizer,

opacus/optimizers/ddp_perlayeroptimizer.py

@@ -48,6 +48,7 @@ def __init__(
         loss_reduction: str = "mean",
         generator=None,
         secure_mode: bool = False,
+        **kwargs,
     ):
         self.rank = torch.distributed.get_rank()
         self.world_size = torch.distributed.get_world_size()
@@ -79,6 +80,7 @@ def __init__(
         loss_reduction: str = "mean",
         generator=None,
         secure_mode: bool = False,
+        **kwargs,
     ):
         self.rank = torch.distributed.get_rank()
         self.world_size = torch.distributed.get_world_size()

opacus/optimizers/ddpoptimizer.py

@@ -38,6 +38,7 @@ def __init__(
         loss_reduction: str = "mean",
         generator=None,
         secure_mode: bool = False,
+        **kwargs,
     ):
         super().__init__(
             optimizer,
@@ -47,6 +48,7 @@ def __init__(
             loss_reduction=loss_reduction,
             generator=generator,
             secure_mode=secure_mode,
+            **kwargs,
         )
         self.rank = torch.distributed.get_rank()
         self.world_size = torch.distributed.get_world_size()

opacus/optimizers/ddpoptimizer_fast_gradient_clipping.py

@@ -38,6 +38,7 @@ def __init__(
         loss_reduction: str = "mean",
         generator=None,
         secure_mode: bool = False,
+        **kwargs,
     ):
         super().__init__(
             optimizer,
@@ -47,6 +48,7 @@ def __init__(
             loss_reduction=loss_reduction,
             generator=generator,
             secure_mode=secure_mode,
+            **kwargs,
         )
         self.rank = torch.distributed.get_rank()
         self.world_size = torch.distributed.get_world_size()

opacus/optimizers/optimizer.py

@@ -205,6 +205,7 @@ def __init__(
         loss_reduction: str = "mean",
         generator=None,
         secure_mode: bool = False,
+        **kwargs,
     ):
         """
 

opacus/optimizers/optimizer_fast_gradient_clipping.py

@@ -63,6 +63,7 @@ def __init__(
         loss_reduction: str = "mean",
         generator=None,
         secure_mode: bool = False,
+        **kwargs,
     ):
         """
 
@@ -91,6 +92,7 @@ def __init__(
             loss_reduction=loss_reduction,
             generator=generator,
             secure_mode=secure_mode,
+            **kwargs,
         )
 
     @property

opacus/optimizers/perlayeroptimizer.py

@@ -39,6 +39,7 @@ def __init__(
         loss_reduction: str = "mean",
         generator=None,
         secure_mode: bool = False,
+        **kwargs,
     ):
         assert len(max_grad_norm) == len(params(optimizer))
         self.max_grad_norms = max_grad_norm
@@ -51,6 +52,7 @@ def __init__(
             loss_reduction=loss_reduction,
             generator=generator,
             secure_mode=secure_mode,
+            **kwargs,
         )
 
     def clip_and_accumulate(self):

opacus/privacy_engine.py

@@ -136,6 +136,7 @@ def _prepare_optimizer(
             loss_reduction=loss_reduction,
             generator=generator,
             secure_mode=self.secure_mode,
+            **kwargs,
         )
 
     def _prepare_data_loader(