add parameter documentation

privacy_engine_xl.py

@@ -6,6 +6,24 @@ import os
 def generate_noise(max_norm, parameter, noise_multiplier, noise_type, device):
     """
     A noise generation function that can utilize different distributions for noise generation.
+
+    @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 parameter
+        The parameter, based on which the dimension of the noise tensor
+        will be determined
+    @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.
+
+    @return
+        a tensor of noise in the same shape as ``parameter``.
     """
     if noise_multiplier > 0:
         mean = 0
@@ -32,6 +50,22 @@ def generate_noise(max_norm, parameter, noise_multiplier, noise_type, device):
 def apply_noise(weights, batch_size, noise_multiplier, noise_type, device, loss_reduction="mean"):
     """
     A function for applying noise to weights on the (intermediate) server side that utilizes the generate_noise function above.
+
+    @param weights
+        The weights to which to apply the noise.
+    @param batch_size
+        Batch size used for averaging.
+    @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
     """
     for p in weights.values():
         noise = generate_noise(0, p, noise_multiplier, noise_type, device)
@@ -44,6 +78,42 @@ class PrivacyEngineXL(opacus.PrivacyEngine):
     """
     A privacy engine that can utilize different distributions for noise generation, based on opacus' privacy engine.
     It gets attached to the optimizer just like the privacy engine from opacus.
+
+    @param module:
+        The Pytorch module to which we are attaching the privacy engine
+    @param batch_size
+        Training batch size. Used in the privacy accountant.
+    @param sample_size
+        The size of the sample (dataset). Used in the privacy accountant.
+    @param alphas
+        A list of RDP orders
+    @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 max_grad_norm
+        The maximum norm of the per-sample gradients. Any gradient with norm
+        higher than this will be clipped to this value.
+    @param secure_rng
+        If on, it will use ``torchcsprng`` for secure random number generation. Comes with
+        a significant performance cost, therefore it's recommended that you turn it off when
+        just experimenting.
+    @param grad_norm_type
+        The order of the norm. For instance, 2 represents L-2 norm, while
+        1 represents L-1 norm.
+    @param batch_first
+        Flag to indicate if the input tensor to the corresponding module
+        has the first dimension representing the batch. If set to True,
+        dimensions on input tensor will be ``[batch_size, ..., ...]``.
+    @param target_delta
+        The target delta
+    @param loss_reduction
+        Indicates if the loss reduction (for aggregating the gradients)
+        is a sum or a mean operation. Can take values "sum" or "mean"
+    @param noise_type
+        Sets the distribution for the noise generation.
+        See generate_noise for supported strings.
+    @param **misc_settings
+        Other arguments to the init
     """
 
     def __init__(
@@ -87,4 +157,17 @@ class PrivacyEngineXL(opacus.PrivacyEngine):
         self.noise_type = noise_type
 
     def _generate_noise(self, max_norm, parameter):
+        """
+        Generates a tensor of noise in the same shape as ``parameter``.
+
+        @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 parameter
+            The parameter, based on which the dimension of the noise tensor
+            will be determined
+
+        @return
+            a tensor of noise in the same shape as ``parameter``.
+        """
         return generate_noise(max_norm, parameter, self.noise_multiplier, self.noise_type, self.device)