pylinting

Author: sayakpaul

.gitmodules

@@ -13,14 +13,14 @@ Minimal implementation of [Denoised Smoothing: A Provable Defense for Pretrained
 
 Randomized Smoothing is a well-tested method to provably defend against _l2_ adversarial attacks under a specific radii. But it assumes that a classifier performs well under Gaussian noisy perturbations which may not always be the case.  
 
-**Note**: I utilized many scripts from the [official repository](https://github.com/microsoft/denoised-smoothing) of Denoised Smoothing to develop this repository. My aim with this repository is to provide a template for researchers to conduct certification tests with Keras/TensorFlow models. I encourage the readers to check out the original repository, it's really well-developed. 
+**Note**: Many scripts have been utilized from the [official repository](https://github.com/microsoft/denoised-smoothing) of Denoised Smoothing to develop this. 
 
 ## Further notes
 
 * The Denoised Smoothing process is demonstrated on the CIFAR-10 dataset. 
-* You can train a classifier quickly with the [`Train_Classifier.ipynb`](https://colab.research.google.com/github/sayakpaul/Denoised-Smoothing-TF/blob/main/Train_Classifier.ipynb) notebook.
-* Training the denoiser is demonstrated in the [`Train_Denoiser.ipynb`](https://colab.research.google.com/github/sayakpaul/Denoised-Smoothing-TF/blob/main/Train_Denoiser.ipynb) notebook. 
-* Certification tests are in [`Certification_Test.ipynb`](https://colab.research.google.com/github/sayakpaul/Denoised-Smoothing-TF/blob/main/Certification_Test.ipynb) notebook. 
+* You can train a classifier quickly with the [`Train_Classifier.ipynb`](https://colab.research.google.com/github/sayakpaul/neural-structured-learning/blob/master/research/denoised_smoothing/notebooks/Train_Classifier.ipynb) notebook.
+* Training of the denoiser is demonstrated in the [`Train_Denoiser.ipynb`](https://colab.research.google.com/github/sayakpaul/neural-structured-learning/blob/master/research/denoised_smoothing/notebooks/Train_Denoiser.ipynb) notebook. 
+* Certification tests are in [`Certification_Test.ipynb`](https://colab.research.google.com/github/sayakpaul/neural-structured-learning/blob/master/research/denoised_smoothing/notebooks/Certification_Test.ipynb) notebook. 
 
 All the notebooks can be executed on Colab! You also have the option to train using the free TPUs.
 
@@ -34,7 +34,7 @@ As we can see prepending a pre-trained denoiser is extremely helpful for our pur
 
 ## Models
 
-The models are available inside [`models.tar.gz`](https://github.com/sayakpaul/Denoised-Smoothing-TF/blob/main/models.tar.gz) in the `SavedModel` format. In the interest of reproducibility, the initial model weights are also provided. 
+The models are available inside [`models.tar.gz`](https://github.com/sayakpaul/Denoised-Smoothing-TF/blob/main/models.tar.gz) in `SavedModel` format. In the interest of reproducibility, the initial model weights are also provided. 
 
 ## Acknowledgements
 

research/denoised_smoothing/README.md

@@ -1,43 +1,73 @@
-# Referenced from: https://github.com/microsoft/denoised-smoothing/blob/master/code/archs/dncnn.py
-# Original: https://github.com/cszn/DnCNN/blob/master/TrainingCodes/dncnn_keras/main_train.py
-
+"""
+# Referenced from:
+https://github.com/microsoft/denoised-smoothing/blob/master/code/archs/dncnn.py
+# Original:
+https://github.com/cszn/DnCNN/blob/master/TrainingCodes/dncnn_keras/main_train.py
+"""
 from tensorflow.keras import layers
 import tensorflow as tf
 
 WEIGHT_DECAY = 1e-4
 
-def conv_block(x, channels=64, ksize=3, padding="same", bn=False):
-    x = layers.Conv2D(
-        channels,
-        kernel_size=ksize,
-        padding=padding,
-        use_bias=True if bn else False,
-        kernel_initializer="orthogonal"
-    )(x)
-    if bn:
-        x = layers.BatchNormalization(momentum=0.0, epsilon=0.0001)(x)
-    x = layers.Activation("relu")(x)
-    return x
-
-
-def run_dncnn(x, image_chnls=3, depth=17, n_channels=64):
-    x = conv_block(x, channels=n_channels)
-    for _ in range(depth - 2):
-        x = conv_block(x, channels=n_channels, bn=True)
-
-    outputs = layers.Conv2D(
-        image_chnls,
-        kernel_size=3,
-        padding="same",
-        use_bias=False,
-        kernel_initializer="orthogonal",
-    )(x)
-    return outputs
-
-
-def get_dncnn(image_size=32, image_chnls=3,
-              depth=17, n_channels=64):
-    inputs = layers.Input((image_size, image_size, image_chnls))
-    outputs = run_dncnn(inputs, depth=depth, n_channels=n_channels)
-    outputs = layers.Subtract()([inputs, outputs])
-    return tf.keras.Model(inputs, outputs)
+def conv_block(x: tf.Tensor, channels: int = 64, ksize: int = 3,
+  padding: str = "same", bn: bool = False) -> tf.Tensor:
+  """Constructs a convolution block.
+
+  :param x: inputs [batch_size, h, w, nb_channels]
+  :param channels: number of channels (int)
+  :param ksize: kernel size (int or a tuple of ints)
+  :param padding: padding type (str)
+  :param bn: whether to use Batch Normalization
+  :return: output of the convolution block
+  """
+  x = layers.Conv2D(
+    channels,
+    kernel_size=ksize,
+    padding=padding,
+    use_bias=True if bn else False,
+    kernel_initializer="orthogonal"
+  )(x)
+  if bn:
+    x = layers.BatchNormalization(momentum=0.0, epsilon=0.0001)(x)
+  x = layers.Activation("relu")(x)
+  return x
+
+
+def run_dncnn(x: tf.Tensor, image_chnls: int = 3, depth: int = 17,
+  n_channels: int = 64) -> tf.Tensor:
+  """Runs a DNCNN block.
+
+  :param x: inputs [batch_size, h, w, nb_channels]
+  :param image_chnls: number of channels in the output images (int)
+  :param depth: depth of the network (int)
+  :param n_channels: number of channels in the convolutional layers (int)
+  :return: batch of images
+  """
+  x = conv_block(x, channels=n_channels)
+  for _ in range(depth - 2):
+    x = conv_block(x, channels=n_channels, bn=True)
+
+  outputs = layers.Conv2D(
+    image_chnls,
+    kernel_size=3,
+    padding="same",
+    use_bias=False,
+    kernel_initializer="orthogonal"
+  )(x)
+  return outputs
+
+
+def get_dncnn(image_size: int = 32, image_chnls: int = 3,
+  depth: int = 17, n_channels: int = 64) -> tf.keras.Model:
+  """Constructs a DNCNN model.
+
+  :param image_size: size of the images (int)
+  :param image_chnls: number of channels in the inputs images (int)
+  :param depth: depth of the network (int)
+  :param n_channels: number of channels in the convolutional layers (int)
+  :return: DNCNN model
+  """
+  inputs = layers.Input((image_size, image_size, image_chnls))
+  outputs = run_dncnn(inputs, depth=depth, n_channels=n_channels)
+  outputs = layers.Subtract()([inputs, outputs])
+  return tf.keras.Model(inputs, outputs)