add clipping to intermediate server WiP

privacy_engine_xl.py

@@ -1,5 +1,7 @@
 import torch
 import opacus
+from opacus.per_sample_gradient_clip import PerSampleGradientClipper
+from opacus.utils import clipping
 from typing import List, Union
 import os
 
@@ -79,6 +81,81 @@ def apply_noise(weights, batch_size, max_norm, noise_multiplier, noise_type, dev
             noise /= batch_size
         p += noise
 
+# Server side Clipping
+def setup_clipper(model, max_grad_norm, batch_first=False, loss_reduction="mean"):
+
+    norm_clipper = (
+        # pyre-fixme[6]: Expected `float` for 1st param but got
+        #  `Union[List[float], float]`.
+        clipping.ConstantFlatClipper(max_grad_norm)
+        if not isinstance(max_grad_norm, list)
+        # pyre-fixme[6]: Expected `List[float]` for 1st param but got
+        #  `Union[List[float], float]`.
+        else clipping.ConstantPerLayerClipper(max_grad_norm)
+    )
+
+    # experimental clipping from opacus
+    # if self.misc_settings.get("experimental", False):
+    #     norm_clipper = clipping._Dynamic_Clipper_(
+    #         # pyre-fixme[6]: Expected `List[float]` for 1st param but got
+    #         #  `List[Union[List[float], float]]`.
+    #         [max_grad_norm],
+    #         self.misc_settings.get("clip_per_layer", False),
+    #         self.misc_settings.get(
+    #             "clipping_method", clipping.ClippingMethod.STATIC
+    #         ),
+    #         self.misc_settings.get("clipping_ratio", 0.0),
+    #         self.misc_settings.get("clipping_momentum", 0.0),
+    #     )
+
+    clipper = PerSampleGradientClipper(
+        model,
+        norm_clipper,
+        batch_first,
+        loss_reduction,
+    )
+
+    return clipper
+
+# Server side Noise
+def apply_noise_clipping(model, batch_size, max_norm, noise_multiplier, noise_type, device, loss_reduction="mean"):
+    """
+    A function for applying noise and clipping to gradients of a model that utilizes the generate_noise function above.
+
+    @param model
+        The model on which's gradients to apply the noise and clipping.
+    @param batch_size
+        Batch size used for averaging.
+    @param max_norm
+        The maximum norm of the per-sample gradients. Any gradient with norm
+        higher than this will be clipped to this value.
+    @param noise_multiplier
+        The ratio of the standard deviation of the Gaussian noise to
+        the L2-sensitivity of the function to which the noise is added
+    @param noise_type
+        Sets the distribution for the noise generation.
+        See generate_noise for supported strings.
+    @param device
+        The device used for calculations and needed for tensor definition.
+    @param loss_reduction
+        The method of loss reduction.
+        currently supported: mean
+    """
+
+    clipper = setup_clipper(model, max_norm, False, loss_reduction)
+
+    clipper.zero_grad()
+
+    clipper.clip_and_accumulate()
+    clip_values, batch_size = clipper.pre_step()
+
+    params = (p for p in model.parameters() if p.requires_grad)
+    for p, clip_value in zip(params, clip_values):
+        noise = generate_noise(clip_value, p, noise_multiplier, noise_type, device)
+        if loss_reduction == "mean":
+            noise /= batch_size
+        p.grad += noise
+
 # Client side Noise
 class PrivacyEngineXL(opacus.PrivacyEngine):
     """