reformatting and dropout fixes to fastmri and vit

Author: priyakasimbeg

algoperf/workloads/criteo1tb/criteo1tb_jax/models.py

@@ -7,6 +7,7 @@
 
 from algoperf.jax_utils import Dropout
 
+
 class DLRMResNet(nn.Module):
   """Define a DLRMResNet model.
 
@@ -30,7 +31,7 @@ class DLRMResNet(nn.Module):
   @nn.compact
   def __call__(self, x, train, dropout_rate=None):
     if not dropout_rate:
-      dropout_rate=self.dropout_rate
+      dropout_rate = self.dropout_rate
 
     bot_mlp_input, cat_features = jnp.split(x, [self.num_dense_features], 1)
     cat_features = jnp.asarray(cat_features, dtype=jnp.int32)
@@ -157,7 +158,7 @@ class DlrmSmall(nn.Module):
   def __call__(self, x, train, dropout_rate=None):
     if not dropout_rate:
       dropout_rate = self.dropout_rate
-      
+
     bot_mlp_input, cat_features = jnp.split(x, [self.num_dense_features], 1)
     cat_features = jnp.asarray(cat_features, dtype=jnp.int32)
 
@@ -219,6 +220,6 @@ def scaled_init(key, shape, dtype=jnp.float_):
       if (dropout_rate is not None and dropout_rate > 0.0 and
           layer_idx == num_layers_top - 2):
         top_mlp_input = Dropout(deterministic=not train)(
-                top_mlp_input, rate=dropout_rate)
+            top_mlp_input, rate=dropout_rate)
     logits = top_mlp_input
     return logits

algoperf/workloads/fastmri/fastmri_jax/models.py

@@ -21,6 +21,7 @@
 
 from algoperf.jax_utils import Dropout
 
+
 def _instance_norm2d(x, axes, epsilon=1e-5):
   # promote x to at least float32, this avoids half precision computation
   # but preserves double or complex floating points
@@ -57,13 +58,13 @@ class UNet(nn.Module):
   num_channels: int = 32
   num_pool_layers: int = 4
   out_channels = 1
-  dropout_rate: Optional[float] = 0.0  # If None, defaults to 0.0.
+  dropout_rate: float = 0.0
   use_tanh: bool = False
   use_layer_norm: bool = False
 
   @nn.compact
   def __call__(self, x, train=True, dropout_rate=None):
-    if not dropout_rate:
+    if dropout_rate is None:
       dropout_rate = self.dropout_rate
 
     # pylint: disable=invalid-name
@@ -138,7 +139,7 @@ class ConvBlock(nn.Module):
   dropout_rate: Dropout probability.
   """
   out_channels: int
-  dropout_rate: float
+  dropout_rate: float = 0.0
   use_tanh: bool
   use_layer_norm: bool
 
@@ -152,8 +153,8 @@ def __call__(self, x, train=True, dropout_rate=None):
     Returns:
         jnp.array: Output tensor of shape `(N, H, W, out_channels)`.
     """
-    if not dropout_rate:
-      dropout_rate=self.dropout_rate
+    if dropout_rate is None:
+      dropout_rate = self.dropout_rate
     x = nn.Conv(
         features=self.out_channels,
         kernel_size=(3, 3),
@@ -174,8 +175,9 @@ def __call__(self, x, train=True, dropout_rate=None):
     x = activation_fn(x)
     # Ref code uses dropout2d which applies the same mask for the entire channel
     # Replicated by using broadcast dims to have the same filter on HW
-    x = Dropout(broadcast_dims=(1, 2), deterministic=not train)(
-            x, rate=dropout_rate )
+    x = Dropout(
+        dropout_rate, broadcast_dims=(1, 2), deterministic=not train)(
+            x, rate=dropout_rate)
     x = nn.Conv(
         features=self.out_channels,
         kernel_size=(3, 3),
@@ -188,7 +190,7 @@ def __call__(self, x, train=True, dropout_rate=None):
       x = _instance_norm2d(x, (1, 2))
     x = activation_fn(x)
     x = Dropout(
-        broadcast_dims=(1, 2), deterministic=not train)(
+        dropout_rate, broadcast_dims=(1, 2), deterministic=not train)(
             x, rate=dropout_rate)
     return x
 

algoperf/workloads/fastmri/fastmri_jax/workload.py

@@ -26,12 +26,21 @@ def init_model_fn(
     """aux_dropout_rate is unused."""
     del aux_dropout_rate
     fake_batch = jnp.zeros((13, 320, 320))
-    self._model = UNet(
-        num_pool_layers=self.num_pool_layers,
-        num_channels=self.num_channels,
-        use_tanh=self.use_tanh,
-        use_layer_norm=self.use_layer_norm,
-        dropout_rate=dropout_rate)
+    if dropout_rate is None:
+      self._model = UNet(
+          num_pool_layers=self.num_pool_layers,
+          num_channels=self.num_channels,
+          use_tanh=self.use_tanh,
+          use_layer_norm=self.use_layer_norm,
+      )
+    else:
+      self._model = UNet(
+          num_pool_layers=self.num_pool_layers,
+          num_channels=self.num_channels,
+          use_tanh=self.use_tanh,
+          use_layer_norm=self.use_layer_norm,
+          dropout_rate=dropout_rate)
+
     params_rng, dropout_rng = jax.random.split(rng)
     variables = jax.jit(
         self._model.init)({'params': params_rng, 'dropout': dropout_rng},

algoperf/workloads/imagenet_vit/imagenet_jax/models.py

@@ -39,9 +39,12 @@ class MlpBlock(nn.Module):
   dropout_rate: float = 0.0
 
   @nn.compact
-  def __call__(self, x: spec.Tensor, train: bool = True, dropout_rate=None) -> spec.Tensor:
+  def __call__(self,
+               x: spec.Tensor,
+               train: bool = True,
+               dropout_rate=None) -> spec.Tensor:
     """Applies Transformer MlpBlock module."""
-    if not dropout_rate:
+    if dropout_rate is None:
       dropout_rate = self.dropout_rate
 
     inits = {
@@ -57,7 +60,7 @@ def __call__(self, x: spec.Tensor, train: bool = True, dropout_rate=None) -> spe
       y = nn.Dense(self.mlp_dim, **inits)(x)
       x = x * y
 
-    x = Dropout()(x, train, rate=dropout_rate)
+    x = Dropout(dropout_rate)(x, train, rate=dropout_rate)
     x = nn.Dense(d, **inits)(x)
     return x
 
@@ -71,9 +74,12 @@ class Encoder1DBlock(nn.Module):
   dropout_rate: float = 0.0
 
   @nn.compact
-  def __call__(self, x: spec.Tensor, train: bool = True, dropout_rate=dropout_rate) -> spec.Tensor:
-    if not dropout_rate:
-      dropout_rate=self.dropout_rate
+  def __call__(self,
+               x: spec.Tensor,
+               train: bool = True,
+               dropout_rate=dropout_rate) -> spec.Tensor:
+    if dropout_rate is None:
+      dropout_rate = self.dropout_rate
 
     if not self.use_post_layer_norm:
       y = nn.LayerNorm(name='LayerNorm_0')(x)
@@ -83,15 +89,14 @@ def __call__(self, x: spec.Tensor, train: bool = True, dropout_rate=dropout_rate
           deterministic=train,
           name='MultiHeadDotProductAttention_1')(
               y)
-      y = Dropout()(y, train, dropout_rate=dropout_rate)
+      y = Dropout(dropout_rate)(y, train, dropout_rate=dropout_rate)
       x = x + y
 
       y = nn.LayerNorm(name='LayerNorm_2')(x)
       y = MlpBlock(
-          mlp_dim=self.mlp_dim,
-          use_glu=self.use_glu,
-          name='MlpBlock_3')(y, train, dropout_rate=dropout_rate)
-      y = Dropout()(y, train, rate=dropout_rate)
+          mlp_dim=self.mlp_dim, use_glu=self.use_glu, name='MlpBlock_3')(
+              y, train, dropout_rate=dropout_rate)
+      y = Dropout(dropout_rate)(y, train, rate=dropout_rate)
       x = x + y
     else:
       y = x
@@ -101,16 +106,18 @@ def __call__(self, x: spec.Tensor, train: bool = True, dropout_rate=dropout_rate
           deterministic=train,
           name='MultiHeadDotProductAttention_1')(
               y)
-      y = Dropout()(y, train, rate=dropout_rate)
+      y = Dropout(dropout_rate)(y, train, rate=dropout_rate)
       x = x + y
       x = nn.LayerNorm(name='LayerNorm_0')(x)
 
       y = x
       y = MlpBlock(
           mlp_dim=self.mlp_dim,
           use_glu=self.use_glu,
-          name='MlpBlock_3')(y, train, dropout_rate=dropout_rate)
-      y = Dropout()(y, train)(rate=dropout_rate)
+          name='MlpBlock_3',
+          dropout_rate=dropout_rate)(
+              y, train, dropout_rate=dropout_rate)
+      y = Dropout(dropout_rate)(y, train)(rate=dropout_rate)
       x = x + y
       x = nn.LayerNorm(name='LayerNorm_2')(x)
 
@@ -127,9 +134,12 @@ class Encoder(nn.Module):
   use_post_layer_norm: bool = False
 
   @nn.compact
-  def __call__(self, x: spec.Tensor, train: bool = True, dropout_rate=None) -> spec.Tensor:
-    if not dropout_rate:
-      dropout_rate=self.dropout_rate
+  def __call__(self,
+               x: spec.Tensor,
+               train: bool = True,
+               dropout_rate=None) -> spec.Tensor:
+    if dropout_rate is None:
+      dropout_rate = self.dropout_rate
 
     # Input Encoder
     for lyr in range(self.depth):
@@ -139,7 +149,8 @@ def __call__(self, x: spec.Tensor, train: bool = True, dropout_rate=None) -> spe
           num_heads=self.num_heads,
           use_glu=self.use_glu,
           use_post_layer_norm=self.use_post_layer_norm,
-          )(dropout_rate=dropout_rate)
+          dropout_rate=dropout_rate)(
+              dropout_rate=dropout_rate)
       x = block(x, train)
     if not self.use_post_layer_norm:
       return nn.LayerNorm(name='encoder_layernorm')(x)
@@ -151,9 +162,12 @@ class MAPHead(nn.Module):
   """Multihead Attention Pooling."""
   mlp_dim: Optional[int] = None  # Defaults to 4x input dim
   num_heads: int = 12
+  dropout_rate: 0.0
 
   @nn.compact
-  def __call__(self, x):
+  def __call__(self, x, dropout_rate=None):
+    if dropout_rate is None:
+      dropout_rate = self.dropout_rate
     n, _, d = x.shape
     probe = self.param('probe',
                        nn.initializers.xavier_uniform(), (1, 1, d),
@@ -166,7 +180,7 @@ def __call__(self, x):
         kernel_init=nn.initializers.xavier_uniform())(probe, x)
 
     y = nn.LayerNorm()(x)
-    x = x + MlpBlock(mlp_dim=self.mlp_dim)(y)
+    x = x + MlpBlock(mlp_dim=self.mlp_dim, dropout_rate=dropout_rate)(y)
     return x[:, 0]
 
 
@@ -180,7 +194,7 @@ class ViT(nn.Module):
   mlp_dim: Optional[int] = None  # Defaults to 4x input dim.
   num_heads: int = 12
   rep_size: Union[int, bool] = True
-  dropout_rate: Optional[float] = 0.0  # If None, defaults to 0.0.
+  dropout_rate: Optional[float] = 0.0
   reinit: Optional[Sequence[str]] = None
   head_zeroinit: bool = True
   use_glu: bool = False
@@ -194,8 +208,12 @@ def get_posemb(self,
     return posemb_sincos_2d(*seqshape, width, dtype=dtype)
 
   @nn.compact
-  def __call__(self, x: spec.Tensor, *, train: bool = False, dropout_rate=None) -> spec.Tensor:
-    if not dropout_rate:
+  def __call__(self,
+               x: spec.Tensor,
+               *,
+               train: bool = False,
+               dropout_rate=None) -> spec.Tensor:
+    if dropout_rate is None:
       dropout_rate = self.dropout_rate
     # Patch extraction
     x = nn.Conv(
@@ -212,19 +230,24 @@ def __call__(self, x: spec.Tensor, *, train: bool = False, dropout_rate=None) ->
     # Add posemb before adding extra token.
     x = x + self.get_posemb((h, w), c, x.dtype)
 
-    x = Dropout()(x, not train, rate=dropout_rate)
+    x = Dropout(dropout_rate)(x, not train, rate=dropout_rate)
 
     x = Encoder(
         depth=self.depth,
         mlp_dim=self.mlp_dim,
         num_heads=self.num_heads,
         use_glu=self.use_glu,
         use_post_layer_norm=self.use_post_layer_norm,
-        name='Transformer')(
+        name='Transformer',
+        dropout_rate=dropout_rate)(
             x, train=not train, dropout_rate=dropout_rate)
 
     if self.use_map:
-      x = MAPHead(num_heads=self.num_heads, mlp_dim=self.mlp_dim)(x)
+      x = MAPHead(
+          num_heads=self.num_heads,
+          mlp_dim=self.mlp_dim,
+          dropout_rate=dropout_rate)(
+              x, dropout_rate=dropout_rate)
     else:
       x = jnp.mean(x, axis=1)
 

algoperf/workloads/imagenet_vit/imagenet_jax/workload.py

@@ -36,13 +36,21 @@ def init_model_fn(
       dropout_rate: Optional[float] = None,
       aux_dropout_rate: Optional[float] = None) -> spec.ModelInitState:
     del aux_dropout_rate
-    self._model = models.ViT(
-        dropout_rate=dropout_rate,
-        num_classes=self._num_classes,
-        use_glu=self.use_glu,
-        use_post_layer_norm=self.use_post_layer_norm,
-        use_map=self.use_map,
-        **decode_variant('S/16'))
+    if dropout_rate is None:
+      self._model = models.ViT(
+          num_classes=self._num_classes,
+          use_glu=self.use_glu,
+          use_post_layer_norm=self.use_post_layer_norm,
+          use_map=self.use_map,
+          **decode_variant('S/16'))
+    else:
+      self._model = models.ViT(
+          dropout_rate=dropout_rate,
+          num_classes=self._num_classes,
+          use_glu=self.use_glu,
+          use_post_layer_norm=self.use_post_layer_norm,
+          use_map=self.use_map,
+          **decode_variant('S/16'))
     params, model_state = self.initialized(rng, self._model)
     self._param_shapes = param_utils.jax_param_shapes(params)
     self._param_types = param_utils.jax_param_types(self._param_shapes)

algoperf/workloads/librispeech_deepspeech/librispeech_jax/models.py

@@ -111,9 +111,7 @@ def __call__(self, inputs, output_paddings, train, dropout_rate=None):
       input_dropout_rate = 0.1
     else:
       input_dropout_rate = config.input_dropout_rate
-    outputs = Dropout(
-        rate=input_dropout_rate, deterministic=not train)(
-            outputs)
+    outputs = Dropout(rate=input_dropout_rate, deterministic=not train)(outputs)
 
     return outputs, output_paddings