remove aux dropout option from conformer and from init_model_fn signature for fastmri, vit and criteo

Author: priyakasimbeg

algoperf/workloads/criteo1tb/criteo1tb_jax/workload.py

@@ -73,11 +73,9 @@ def init_model_fn(
       self,
       rng: spec.RandomState,
       dropout_rate: Optional[float] = None,
-      aux_dropout_rate: Optional[float] = None,
       tabulate: Optional[bool] = False,
   ) -> spec.ModelInitState:
     """Only dropout is used."""
-    del aux_dropout_rate
     if self.use_resnet:
       model_class = models.DLRMResNet
     else:

algoperf/workloads/fastmri/fastmri_jax/workload.py

@@ -22,9 +22,8 @@ def init_model_fn(
       self,
       rng: spec.RandomState,
       dropout_rate: Optional[float] = None,
-      aux_dropout_rate: Optional[float] = None) -> spec.ModelInitState:
+      ) -> spec.ModelInitState:
     """aux_dropout_rate is unused."""
-    del aux_dropout_rate
     fake_batch = jnp.zeros((13, 320, 320))
     if dropout_rate is None:
       self._model = UNet(

algoperf/workloads/imagenet_vit/imagenet_jax/workload.py

@@ -33,9 +33,7 @@ def initialized(self, key: spec.RandomState,
   def init_model_fn(
       self,
       rng: spec.RandomState,
-      dropout_rate: Optional[float] = None,
-      aux_dropout_rate: Optional[float] = None) -> spec.ModelInitState:
-    del aux_dropout_rate
+      dropout_rate: Optional[float] = None) -> spec.ModelInitState:
     if dropout_rate is None:
       self._model = models.ViT(
           num_classes=self._num_classes,

algoperf/workloads/librispeech_conformer/librispeech_jax/workload.py

@@ -61,24 +61,32 @@ def init_model_fn(
       self,
       rng: spec.RandomState,
       dropout_rate: Optional[float] = None,
-      aux_dropout_rate: Optional[float] = None) -> spec.ModelInitState:
+      ) -> spec.ModelInitState:
     """Conformer model init function.
 
-    Here we use dropout_rate as *_residual_dropout_rate, and aux_dropout_rate as
+    Here we use dropout_rate as *_residual_dropout_rate, and for
     input_dropout_rate.
     """
     if self.use_gelu:
       activation_function_name = 'gelu'
     else:
       activation_function_name = 'swish'
-    model_config = models.ConformerConfig(
-        attention_residual_dropout_rate=dropout_rate,
-        feed_forward_residual_dropout_rate=dropout_rate,
-        input_dropout_rate=aux_dropout_rate,
-        use_specaug=self.use_specaug,
-        attention_temperature=self.attention_temperature,
-        use_post_layer_norm=self.use_post_layer_norm,
-        activation_function_name=activation_function_name)
+    if dropout_rate is None:
+      model_config = models.ConformerConfig(
+          attention_residual_dropout_rate=dropout_rate,
+          feed_forward_residual_dropout_rate=dropout_rate,
+          input_dropout_rate=dropout_rate,
+          use_specaug=self.use_specaug,
+          attention_temperature=self.attention_temperature,
+          use_post_layer_norm=self.use_post_layer_norm,
+          activation_function_name=activation_function_name)
+    else:
+      model_config = models.ConformerConfig(
+          use_specaug=self.use_specaug,
+          attention_temperature=self.attention_temperature,
+          use_post_layer_norm=self.use_post_layer_norm,
+          activation_function_name=activation_function_name)
+
     self._model = models.Conformer(model_config)
     input_shape = [(320000,), (320000,)]
     fake_input_batch = [np.zeros((2, *x), jnp.float32) for x in input_shape]