add dropout piping for conformer and deepspeech

Author: priyakasimbeg

algoperf/workloads/fastmri/fastmri_jax/workload.py

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

algoperf/workloads/librispeech_conformer/librispeech_jax/models.py

@@ -38,13 +38,10 @@ class ConformerConfig:
   num_attention_heads: int = 8
   num_encoder_layers: int = 4
   attention_dropout_rate: float = 0.0
-  # If None, defaults to 0.1.
-  attention_residual_dropout_rate: Optional[float] = 0.1
-  # If None, defaults to 0.0.
+  attention_residual_dropout_rate: Optional[float] = 0.0
   conv_residual_dropout_rate: Optional[float] = 0.0
   feed_forward_dropout_rate: float = 0.0
-  # If None, defaults to 0.1.
-  feed_forward_residual_dropout_rate: Optional[float] = 0.1
+  feed_forward_residual_dropout_rate: Optional[float] = 0.0
   convolution_kernel_size: int = 5
   feed_forward_expansion_factor: int = 4
   freq_mask_count: int = 2

algoperf/workloads/librispeech_conformer/librispeech_jax/workload.py

@@ -61,7 +61,7 @@ def init_model_fn(
       self,
       rng: spec.RandomState,
       dropout_rate: Optional[float] = None,
-      ) -> spec.ModelInitState:
+  ) -> spec.ModelInitState:
     """Conformer model init function.
 
     Here we use dropout_rate as *_residual_dropout_rate, and for

algoperf/workloads/librispeech_deepspeech/librispeech_jax/models.py

@@ -75,6 +75,9 @@ class Subsample(nn.Module):
   @nn.compact
   def __call__(self, inputs, output_paddings, train, dropout_rate=None):
     config = self.config
+    if dropout_rate is None:
+      dropout_rate = config.dropout_rate
+
     outputs = jnp.expand_dims(inputs, axis=-1)
 
     outputs, output_paddings = Conv2dSubsampling(
@@ -111,7 +114,9 @@ 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, rate=dropout_rate)(
+            outputs, rate=dropout_rate)
 
     return outputs, output_paddings
 
@@ -187,7 +192,13 @@ class FeedForwardModule(nn.Module):
   config: DeepspeechConfig
 
   @nn.compact
-  def __call__(self, inputs, input_paddings=None, train=False):
+  def __call__(self,
+               inputs,
+               input_paddings=None,
+               train=False,
+               dropout_rate=None):
+    if dropout_rate is None:
+      dropout_rate = self.config.feed_forward_dropout_rate
     padding_mask = jnp.expand_dims(1 - input_paddings, -1)
     config = self.config
 
@@ -211,12 +222,8 @@ def __call__(self, inputs, input_paddings=None, train=False):
       inputs = nn.relu(inputs)
     inputs *= padding_mask
 
-    if config.feed_forward_dropout_rate is None:
-      feed_forward_dropout_rate = 0.1
-    else:
-      feed_forward_dropout_rate = config.feed_forward_dropout_rate
-    inputs = Dropout(rate=feed_forward_dropout_rate)(
-        inputs, deterministic=not train)
+    inputs = Dropout(rate=dropout_rate)(
+        inputs, deterministic=not train, rate=dropout_rate)
 
     return inputs
 
@@ -472,8 +479,10 @@ def setup(self):
     )
 
   @nn.compact
-  def __call__(self, inputs, input_paddings, train):
+  def __call__(self, inputs, input_paddings, train, dropout_rate=None):
     config = self.config
+    if dropout_rate is None:
+      dropout_rate = config.dropout_rate
 
     outputs = inputs
     output_paddings = input_paddings
@@ -493,7 +502,7 @@ def __call__(self, inputs, input_paddings, train):
 
     # Subsample input by a factor of 4 by performing strided convolutions.
     outputs, output_paddings = Subsample(
-        config=config)(outputs, output_paddings, train)
+        config=config)(outputs, output_paddings, train, dropout_rate=dropout_rate)
 
     # Run the lstm layers.
     for _ in range(config.num_lstm_layers):
@@ -507,9 +516,8 @@ def __call__(self, inputs, input_paddings, train):
         outputs = outputs + FeedForwardModule(config=self.config)(
             outputs, output_paddings, train)
       else:
-        outputs = FeedForwardModule(config=self.config)(outputs,
-                                                        output_paddings,
-                                                        train)
+        outputs = FeedForwardModule(config=self.config)(
+            outputs, output_paddings, train, dropout_rate=dropout_rate)
 
     # Run the decoder which in this case is a trivial projection layer.
     if config.enable_decoder_layer_norm:

algoperf/workloads/librispeech_deepspeech/librispeech_jax/workload.py

@@ -18,24 +18,31 @@ class LibriSpeechDeepSpeechWorkload(LibriSpeechConformerWorkload):
   def init_model_fn(
       self,
       rng: spec.RandomState,
-      dropout_rate: Optional[float] = None,
-      aux_dropout_rate: Optional[float] = None) -> spec.ModelInitState:
+      dropout_rate: Optional[float] = None) -> spec.ModelInitState:
     """Deepspeech model init function.
-
-    Here we use dropout_rate as feed_forward_dropout_rate, and aux_dropout_rate
-    as input_dropout_rate.
     """
-    model_config = models.DeepspeechConfig(
-        feed_forward_dropout_rate=dropout_rate,
-        use_specaug=self.use_specaug,
-        input_dropout_rate=aux_dropout_rate,
-        use_tanh=self.use_tanh,
-        enable_residual_connections=self.enable_residual_connections,
-        enable_decoder_layer_norm=self.enable_decoder_layer_norm,
-        layernorm_everywhere=self.layernorm_everywhere,
-        freq_mask_count=self.freq_mask_count,
-        time_mask_count=self.time_mask_count,
-    )
+    if dropout_rate is None:
+      model_config = models.DeepspeechConfig(
+          use_specaug=self.use_specaug,
+          use_tanh=self.use_tanh,
+          enable_residual_connections=self.enable_residual_connections,
+          enable_decoder_layer_norm=self.enable_decoder_layer_norm,
+          layernorm_everywhere=self.layernorm_everywhere,
+          freq_mask_count=self.freq_mask_count,
+          time_mask_count=self.time_mask_count,
+      )
+    else:
+      model_config = models.DeepspeechConfig(
+          feed_forward_dropout_rate=dropout_rate,
+          use_specaug=self.use_specaug,
+          input_dropout_rate=dropout_rate,
+          use_tanh=self.use_tanh,
+          enable_residual_connections=self.enable_residual_connections,
+          enable_decoder_layer_norm=self.enable_decoder_layer_norm,
+          layernorm_everywhere=self.layernorm_everywhere,
+          freq_mask_count=self.freq_mask_count,
+          time_mask_count=self.time_mask_count,
+      )
     self._model = models.Deepspeech(model_config)
     input_shape = [(320000,), (320000,)]
     fake_input_batch = [np.zeros((2, *x), jnp.float32) for x in input_shape]
@@ -64,6 +71,7 @@ def model_fn(
       rng: spec.RandomState,
       update_batch_norm: bool,
       use_running_average_bn: Optional[bool] = None
+      dropout_rate: Optional[bool] = None
   ) -> Tuple[spec.Tensor, spec.ModelAuxiliaryState]:
     variables = {'params': params, **model_state}
     inputs, input_paddings = augmented_and_preprocessed_input_batch['inputs']
@@ -75,7 +83,8 @@ def model_fn(
           input_paddings,
           train=True,
           rngs={'dropout' : rng},
-          mutable=['batch_stats'])
+          mutable=['batch_stats'],
+          dropout_rate=dropout_rate)
       return (logits, logit_paddings), new_model_state
     else:
       logits, logit_paddings = self._model.apply(

algoperf/workloads/ogbg/ogbg_jax/models.py

@@ -48,9 +48,9 @@ class GNN(nn.Module):
 
   @nn.compact
   def __call__(self, graph, train, dropout_rate=None):
-    if not dropout_rate:
+    if dropout_rate is not None:
       dropout_rate = self.dropout_rate
-    dropout = Dropout(deterministic=not train, rate=dropout_rate)
+    dropout = Dropout(dropout_rate, deterministic=not train)(dropout_rate)
 
     graph = graph._replace(
         globals=jnp.zeros([graph.n_node.shape[0], self.num_outputs]))

algoperf/workloads/ogbg/ogbg_jax/workload.py

@@ -20,18 +20,24 @@ class OgbgWorkload(BaseOgbgWorkload):
   def init_model_fn(
       self,
       rng: spec.RandomState,
-      dropout_rate: Optional[float] = None,
-      aux_dropout_rate: Optional[float] = None) -> spec.ModelInitState:
+      dropout_rate: Optional[float] = None) -> spec.ModelInitState:
     """aux_dropout_rate is unused."""
-    del aux_dropout_rate
     rng, params_rng, dropout_rng = jax.random.split(rng, 3)
-    self._model = models.GNN(
-        self._num_outputs,
-        dropout_rate=dropout_rate,
-        activation_fn_name=self.activation_fn_name,
-        hidden_dims=self.hidden_dims,
-        latent_dim=self.latent_dim,
-        num_message_passing_steps=self.num_message_passing_steps)
+    if dropout_rate is None:
+      self._model = models.GNN(
+          self._num_outputs,
+          activation_fn_name=self.activation_fn_name,
+          hidden_dims=self.hidden_dims,
+          latent_dim=self.latent_dim,
+          num_message_passing_steps=self.num_message_passing_steps)
+    else:
+      self._model = models.GNN(
+          self._num_outputs,
+          dropout_rate=dropout_rate,
+          activation_fn_name=self.activation_fn_name,
+          hidden_dims=self.hidden_dims,
+          latent_dim=self.latent_dim,
+          num_message_passing_steps=self.num_message_passing_steps)
     init_fn = jax.jit(functools.partial(self._model.init, train=False))
     fake_batch = jraph.GraphsTuple(
         n_node=jnp.asarray([1]),

algoperf/workloads/wmt/wmt_jax/workload.py

@@ -209,10 +209,7 @@ def translate_and_calculate_bleu(self,
   def init_model_fn(
       self,
       rng: spec.RandomState,
-      dropout_rate: Optional[float] = 0.0,
-      aux_dropout_rate: Optional[float] = 0.0) -> spec.ModelInitState:
-    """aux_dropout_rate is used as attention_dropout_rate."""
-
+      dropout_rate: Optional[float] = 0.0) -> spec.ModelInitState:
     init_fake_batch_size = 2
     input_shape = (init_fake_batch_size, 256)
     target_shape = (init_fake_batch_size, 256)
@@ -224,13 +221,20 @@ def init_model_fn(
     else:
       raise ValueError(f'Unknown activation function {self.activation}.')
 
+    if dropout_rate is None:
+      model_config = models.TransformerConfig(
+          pre_ln=self.pre_ln,
+          attention_temp=self.attention_temp,
+          activation=activation,
+          glu=self.glu)
+    else:
     model_config = models.TransformerConfig(
-        dropout_rate=dropout_rate,
-        attention_dropout_rate=aux_dropout_rate,
-        pre_ln=self.pre_ln,
-        attention_temp=self.attention_temp,
-        activation=activation,
-        glu=self.glu)
+          dropout_rate=dropout_rate,
+          attention_dropout_rate=dropout_rate,
+          pre_ln=self.pre_ln,
+          attention_temp=self.attention_temp,
+          activation=activation,
+          glu=self.glu)
     self._train_model = models.Transformer(model_config)
     eval_config = replace(model_config, deterministic=True)
     self._eval_model = models.Transformer(eval_config)