Hparam specification for solving Librispeech with Transformer

PiperOrigin-RevId: 185972342

Author: T2T Team

tensor2tensor/data_generators/librispeech.py

@@ -185,3 +185,10 @@ def add_librispeech_hparams(hparams):
   hparams.train_steps = 5000000
   hparams.num_hidden_layers = 4
   return hparams
+
+
+def set_librispeech_length_hparams(hparams):
+  hparams.max_length = 1650 * 80  # this limits inputs[1] * inputs[2]
+  hparams.max_input_seq_length = 1650
+  hparams.max_target_seq_length = 350
+  return hparams

tensor2tensor/models/transformer.py

@@ -31,6 +31,7 @@
 
 from six.moves import xrange  # pylint: disable=redefined-builtin
 
+from tensor2tensor.data_generators import librispeech
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
@@ -1423,3 +1424,32 @@ def transformer_lm_tpu_1():
   hparams.hidden_size = 2048
   hparams.filter_size = 8192
   return hparams
+
+
+@registry.register_hparams
+def transformer_librispeech():
+  """Hparams for training ASR model on Librispeech."""
+  hparams = transformer_base()
+
+  hparams.num_heads = 4
+  hparams.filter_size = 1024
+  hparams.hidden_size = 256
+  hparams.num_encoder_layers = 5
+  hparams.num_decoder_layers = 3
+  hparams.learning_rate = 0.15
+  hparams.batch_size = 6000000
+
+  librispeech.set_librispeech_length_hparams(hparams)
+  return hparams
+
+
+@registry.register_hparams
+def transformer_librispeech_tpu():
+  """Hparams for training ASR model on Librispeech on TPU."""
+  hparams = transformer_librispeech()
+  update_hparams_for_tpu(hparams)
+
+  hparams.batch_size = 32
+  librispeech.set_librispeech_length_hparams(hparams)
+  return hparams
+