internal merge of PR tensorflow#1389

PiperOrigin-RevId: 230774856

Author: googlehjx

tensor2tensor/data_generators/problem.py

@@ -877,7 +877,7 @@ def serving_input_fn(self, hparams):
         dtype=tf.string, shape=[None], name="serialized_example")
     dataset = tf.data.Dataset.from_tensor_slices(serialized_example)
     dataset = dataset.map(self.decode_example)
-    dataset = dataset.map(lambda ex: self.preprocess_example(ex, mode, hparams))   
+    dataset = dataset.map(lambda ex: self.preprocess_example(ex, mode, hparams))
     dataset = dataset.map(data_reader.cast_ints_to_int32)
     dataset = dataset.padded_batch(
         tf.shape(serialized_example, out_type=tf.int64)[0],

tensor2tensor/models/research/universal_transformer.py

@@ -438,21 +438,14 @@ def update_hparams_for_universal_transformer(hparams):
 
 @registry.register_hparams
 def universal_transformer_base():
-  hparams = transformer.transformer_base()
-  # To have a similar capacity to the transformer_base with 6 layers,
-  # we need to increase the size of the UT's layer
-  # since, in fact, UT has a single layer repeating multiple times.
-  hparams.hidden_size = 1024
-  hparams.filter_size = 4096
-  hparams.num_heads = 16
-  hparams.layer_prepostprocess_dropout = 0.3
+  hparams = transformer.transformer_big()
   hparams = update_hparams_for_universal_transformer(hparams)
   return hparams
 
 
 @registry.register_hparams
 def universal_transformer_base_tpu():
-  hparams = universal_transformer_base()
+  hparams = transformer.transformer_big()
   hparams = update_hparams_for_universal_transformer(hparams)
   transformer.update_hparams_for_tpu(hparams)
   hparams.add_step_timing_signal = False
@@ -461,7 +454,7 @@ def universal_transformer_base_tpu():
 
 @registry.register_hparams
 def universal_transformer_big():
-  hparams = universal_transformer_base()
+  hparams = transformer.transformer_big()
   hparams = update_hparams_for_universal_transformer(hparams)
   hparams.hidden_size = 2048
   hparams.filter_size = 8192