swap out lstm layer

Author: priyakasimbeg

algoperf/workloads/librispeech_deepspeech/librispeech_jax/models.py

@@ -439,6 +439,8 @@ class BatchRNN(nn.Module):
   @nn.compact
   def __call__(self, inputs, input_paddings, train):
     config = self.config
+    lengths = jnp.sum(1 - input_paddings, axis=-1, dtype=jnp.int32)
+
 
     if config.layernorm_everywhere:
       inputs = LayerNorm(config.encoder_dim)(inputs)
@@ -452,7 +454,7 @@ def __call__(self, inputs, input_paddings, train):
     output, _ = lstm.LSTM(
         hidden_size=config.encoder_dim // 2,
         bidirectional=config.bidirectional,
-        num_layers=1)(inputs, input_paddings)
+        num_layers=1)(inputs, lengths)
 
     return output
 

algoperf/workloads/librispeech_deepspeech/librispeech_jax/workload.py

@@ -80,8 +80,6 @@ def model_fn(
           train=True,
           rngs={'dropout' : rng},
           mutable=['batch_stats'])
-      if 'batch_stats' in new_model_state and new_model_state['batch_stats']:
-        new_model_state = jax.lax.pmean(new_model_state, 'batch')
       return (logits, logit_paddings), new_model_state
     else:
       logits, logit_paddings = self._model.apply(

scoring/plot_utils/plot_curves.py

@@ -10,17 +10,17 @@
 flags.DEFINE_string(
     'experiment_dir',
     # '/home/kasimbeg/algoperf-runs-internal/experiments/jit_switch_debug_conformer_old_step_hint',
-    '/home/kasimbeg/submissions_algorithms/logs/external_tuning/baseline',
+    '/home/kasimbeg/algoperf-runs-internal/experiments/jit_switch_debug_deepspeech_nadamw_jit_branch',
     'Path to experiment dir.')
 flags.DEFINE_string(
     'workloads',
-    'librispeech_conformer_jax',
+    'librispeech_deepspeech_jax',
     'Filter only for workload e.g. fastmri_jax. If None include all workloads in experiment.'
 )
 flags.DEFINE_string('project_name',
-                    'visualize-training-curves-legacy-stephint-conformer',
+                    'visualize-training-curves-legacy-stephint-deepspeech',
                     'Wandb project name.')
-flags.DEFINE_string('run_postfix', 'pmap', 'Postfix for wandb runs.')
+flags.DEFINE_string('run_postfix', 'jit_legacy_lstm', 'Postfix for wandb runs.')
 
 FLAGS = flags.FLAGS