Remove the remaining modality strings and combine hparams.{input_modalities,target_modality}.

PiperOrigin-RevId: 219617388

Author: Dustin Tran

tensor2tensor/data_generators/problem.py

@@ -28,7 +28,6 @@
 
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import text_encoder
-from tensor2tensor.layers import modalities
 from tensor2tensor.utils import data_reader
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import mlperf_log
@@ -1148,33 +1147,12 @@ def _create_modalities(problem_hparams, hparams):
   Returns:
     None
   """
-  input_modality_overrides = {}
-  if hasattr(hparams, "input_modalities"):
-    for override_str in hparams.input_modalities.split(";"):
-      if override_str != "default":
-        parts = override_str.split(":")
-        feature_name = parts[0]
-        modality_name = ":".join(parts[1:])
-        input_modality_overrides[feature_name] = modality_name
-
-  target_modality_name = None
-  if (hasattr(hparams, "target_modality") and
-      hparams.target_modality != "default"):
-    target_modality_name = hparams.target_modality
-
+  modality_overrides = getattr(hparams, "modality", {})
   modality = {}
   for feature_name, modality_cls in six.iteritems(problem_hparams.modality):
     vocab_size = problem_hparams.vocab_size[feature_name]
-    if feature_name in input_modality_overrides:
-      modality_obj = modalities.create_modality(
-          (input_modality_overrides[feature_name], vocab_size), hparams)
-    elif target_modality_name and feature_name == "targets":
-      # TODO(lukaszkaiser): allow overriding other target modalities.
-      modality_obj = modalities.create_modality(
-          (target_modality_name, vocab_size), hparams)
-    else:
-      modality_obj = modality_cls(hparams, vocab_size)
-    modality[feature_name] = modality_obj
+    modality_cls = modality_overrides.get(feature_name, modality_cls)
+    modality[feature_name] = modality_cls(hparams, vocab_size)
   problem_hparams.modality = modality
 
 
@@ -1200,17 +1178,10 @@ def _default_hparams():
       # token.
       stop_at_eos=False,
 
-      # Modalities used to map from input features to a space compatible with
-      # chosen model architecture.  One modality spec (which is a 2-tuple,
-      # (modality_full_name, vocab_size)) per feature key. modality_full_name
-      # is a string type:name, e.g. class_label:class_label_2d. Leaving off
-      # the name uses the default modality for that type (e.g. class_label ==
-      # class_label:default).
-      input_modality={},
-
-      # Modality used to map from hidden representation to the target space.
-      # Specified as a modality spec, a 2-tuple described above.
-      target_modality=None,
+      # Modalities used to map from features to a space compatible with
+      # chosen model architecture. It comprises key-value pairs of a feature
+      # name (str) and its modality class.
+      modality={},
 
       # Identifiers used to tell the model which input/target space will be
       # expected. For example, it can tell that we expect French as characters

tensor2tensor/layers/common_hparams.py

@@ -162,18 +162,14 @@ def basic_params1():
       # embeddings and the target embeddings.
       # You can also share the input embeddings with the target embeddings
       # by using a problem_hparams that uses the same modality object for
-      # the input_modality and target_modality.
+      # the input modality and target modality.
       shared_embedding=False,
       # In SymbolModality, skip the top layer, assume we're providing logits.
       symbol_modality_skip_top=False,
-      # For each feature for which you want to override the default input
-      # modality, add an entry to this semicolon-separated string. Entries are
-      # formatted "feature_name:modality_type:modality_name", e.g.
-      # "inputs:symbol:default;other_inputs:audio:identity".
-      input_modalities="default",  # We don't use empty string in params.
-      # To override the default target modality, specify
-      # "modality_type:modality_name", e.g. "symbol:ctc".
-      target_modality="default",
+      # Modalities used to map from features to a space compatible with
+      # chosen model architecture. It comprises key-value pairs of a feature
+      # name (str) and its modality class.
+      modality={},
       # The maximum length of "input" sequence.
       # Sequences longer than this value will be truncated. 0 or negative values
       # mean there is no maximum or truncation.

tensor2tensor/layers/common_image_attention.py

@@ -608,8 +608,14 @@ def prepare_image(inputs, hparams, name=None):
   channels = hparams.num_channels
 
   hidden_size = hparams.hidden_size
-  # Only do lookup if the modality is identity
-  if hparams.target_modality == "image:identity":
+  # TODO(trandustin): Check via modalities.IdentityModality and not its name.
+  # The current implementation is to avoid circular imports, modalities ->
+  # discretization -> common_image_attention -> modalities.
+  if "targets" in hparams.modality:
+    target_modality_name = hparams.modality["targets"].__name__
+  else:
+    target_modality_name = None
+  if target_modality_name == "IdentityModality":
     inputs = tf.to_int32(inputs)
     x = get_channel_embeddings(channels, inputs, hidden_size, name=name)
   else:

tensor2tensor/layers/modalities.py

@@ -25,7 +25,6 @@
 from tensor2tensor.layers import common_video
 from tensor2tensor.layers import discretization
 from tensor2tensor.utils import modality
-from tensor2tensor.utils import registry
 
 import tensorflow as tf
 
@@ -1070,96 +1069,3 @@ def top(self, body_output, _):
       x = body_output
       x = tf.expand_dims(x[:, -1], 1)  # Pick the last timestep
       return tf.layers.dense(x, self._vocab_size)
-
-
-def create_modality(modality_spec, model_hparams):
-  """Creates modality.
-
-  Args:
-    modality_spec: tuple ("modality_type:modality_name", vocab_size).
-    model_hparams: tf.contrib.training.HParams.
-
-  Returns:
-    Modality.
-
-  Raises:
-    LookupError: if modality_type is not recognized. See registry.Modalities for
-      accepted types.
-  """
-  modality_full_name, vocab_size = modality_spec
-  modality_type, modality_name = parse_modality_name(modality_full_name)
-
-  if modality_type == registry.Modalities.SYMBOL:
-    modality_collection = {
-        "default": SymbolModality,
-        "identity": IdentitySymbolModality,
-        "weights_all": SymbolModalityWeightsAll,
-        "one_hot": SymbolModalityOneHot,
-        "ctc": CTCSymbolModality,
-    }
-  elif modality_type == registry.Modalities.IMAGE:
-    modality_collection = {
-        "default": ImageModality,
-        "identity": IdentityModality,
-        "image_channel_compress": ImageChannelCompressModality,
-        "image_channel_bottom_identity": ImageChannelBottomIdentityModality,
-        "channel_embeddings_bottom": ImageChannelEmbeddingsBottom,
-    }
-  elif modality_type == registry.Modalities.AUDIO:
-    modality_collection = {
-        "default": SpeechRecognitionModality,
-        "identity": IdentityModality,
-        "spectral": AudioSpectralModality,
-        "speech": SpeechRecognitionModality,
-    }
-  elif modality_type == registry.Modalities.VIDEO:
-    modality_collection = {
-        "default": VideoModality,
-        "identity": IdentityModality,
-        "bitwise": VideoModalityBitwise,
-        "pixel_noise": VideoModalityPixelNoise,
-        "l1": VideoModalityL1,
-        "l2": VideoModalityL2,
-        "l2raw": VideoModalityL2Raw,
-        "l1raw": VideoModalityL1Raw,
-    }
-  elif modality_type == registry.Modalities.CLASS_LABEL:
-    modality_collection = {
-        "default": ClassLabelModality,
-        "identity": IdentityModality,
-        "multi_label": MultiLabelModality,
-        "onehot": OneHotClassLabelModality,
-        "sigmoid": SigmoidClassLabelModality,
-        "sigmoid_max_pooling": SigmoidMaxPoolingClassLabelModality,
-        "onehot_softmax_max_pooling": SoftmaxMaxPoolingClassLabelModality,
-        "onehot_softmax_average_pooling":
-            SoftmaxAveragePoolingClassLabelModality,
-        "onehot_softmax_last_timestep": SoftmaxLastTimestepClassLabelModality,
-    }
-  elif modality_type == registry.Modalities.GENERIC:
-    modality_collection = {
-        "default": IdentityModality,
-        "l2_loss": GenericL2LossModality,
-    }
-  elif modality_type == registry.Modalities.REAL:
-    modality_collection = {
-        "default": RealL2LossModality,
-        "identity": IdentityModality,
-        "l2_loss": RealL2LossModality,
-        "log_poisson_loss": RealLogPoissonLossModality,
-    }
-  else:
-    modality_types = ("symbol", "image", "audio", "video", "class_label",
-                      "generic", "real")
-    raise LookupError("Modality type %s not recognized. Options are: %s" %
-                      (modality_type, list(modality_types)))
-
-  return modality_collection[modality_name](model_hparams, vocab_size)
-
-
-def parse_modality_name(name):
-  name_parts = name.split(":")
-  if len(name_parts) < 2:
-    name_parts.append("default")
-  modality_type, modality_name = name_parts
-  return modality_type, modality_name

tensor2tensor/layers/modalities_test.py

@@ -23,7 +23,6 @@
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import modalities
 from tensor2tensor.utils import expert_utils
-from tensor2tensor.utils import registry
 
 import tensorflow as tf
 
@@ -113,22 +112,6 @@ def testSymbolModalityTargetsFactored(self):
     self.assertEqual(res1.shape, (batch_size, length, height, 1, vocab_size))
     self.assertEqual(res2.shape, ())
 
-  @tf.contrib.eager.run_test_in_graph_and_eager_modes()
-  def testCreateModality(self):
-    model_hparams = tf.contrib.training.HParams()
-
-    modality_spec = (registry.Modalities.SYMBOL, 2)
-    modality = modalities.create_modality(modality_spec, model_hparams)
-    self.assertIsInstance(modality, modalities.SymbolModality)
-
-    modality_spec = (registry.Modalities.CLASS_LABEL + ":onehot", None)
-    modality = modalities.create_modality(modality_spec, model_hparams)
-    self.assertIsInstance(modality, modalities.OneHotClassLabelModality)
-
-    modality_spec = (registry.Modalities.VIDEO + ":identity", None)
-    modality = modalities.create_modality(modality_spec, model_hparams)
-    self.assertIsInstance(modality, modalities.IdentityModality)
-
 
 if __name__ == "__main__":
   tf.test.main()

tensor2tensor/models/image_transformer.py

@@ -28,6 +28,7 @@
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_image_attention as cia
 from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
@@ -47,14 +48,16 @@ def body(self, features):
     hparams = copy.copy(self._hparams)
     targets = features["targets"]
     if (hparams.likelihood == cia.DistributionType.DMOL and
-        (hparams.target_modality != "image:image_channel_bottom_identity" or
+        (hparams.modality["targets"] !=
+         modalities.ImageChannelBottomIdentityModality or
          hparams.num_channels != 1)):
-      raise ValueError("When using DMOL for the likelihood, target_modality "
-                       "must be image:image_channel_bottom_identity and "
+      raise ValueError("When using DMOL for the likelihood,modality['targets'] "
+                       "must be ImageChannelBottomIdentityModality and "
                        "num_channels must be 1.")
     if (not tf.get_variable_scope().reuse and
         hparams.mode != tf.contrib.learn.ModeKeys.INFER and
-        hparams.target_modality != "image:image_channel_bottom_identity"):
+        hparams.modality["targets"] !=
+        modalities.ImageChannelBottomIdentityModality):
       tf.summary.image("targets", tf.to_float(targets), max_outputs=1)
 
     # Extra losses list if we want to use moe.
@@ -190,7 +193,7 @@ def image_transformer_base():
   hparams.optimizer_adam_beta1 = 0.9
   hparams.optimizer_adam_beta2 = 0.98
   hparams.label_smoothing = 0.0
-  hparams.target_modality = "image:identity"
+  hparams.modality["targets"] = modalities.IdentityModality
   hparams.norm_type = "layer"
   hparams.layer_prepostprocess_dropout = 0.0
   hparams.add_hparam("filter_size", 512)  # Add new ones like this.
@@ -277,7 +280,7 @@ def imagetransformer_cifar10_base_dmol():
   hparams = image_transformer_base()
   hparams.likelihood = cia.DistributionType.DMOL
   hparams.num_channels = 1
-  hparams.target_modality = "image:image_channel_bottom_identity"
+  hparams.modality["targets"] = modalities.ImageChannelBottomIdentityModality
   hparams.num_heads = 8
   hparams.batch_size = 8
   hparams.sampling_method = "random"
@@ -418,7 +421,7 @@ def imagetransformerpp_sep_channels_8l_8h():
   hparams = imagetransformer_base()
   hparams.likelihood = cia.DistributionType.DMOL
   hparams.num_channels = 1
-  hparams.target_modality = "image:image_channel_bottom_identity"
+  hparams.modality["targets"] = modalities.ImageChannelBottomIdentityModality
   hparams.num_heads = 8
   hparams.batch_size = 4
   hparams.attention_key_channels = hparams.attention_value_channels = 0
@@ -881,7 +884,7 @@ def imagetransformerpp_tiny():
   hparams = imagetransformer_tiny()
   hparams.likelihood = cia.DistributionType.DMOL
   hparams.num_channels = 1
-  hparams.target_modality = "image:image_channel_bottom_identity"
+  hparams.modality["targets"] = modalities.ImageChannelBottomIdentityModality
   return hparams
 
 

tensor2tensor/models/image_transformer_2d.py

@@ -29,6 +29,7 @@
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_image_attention as cia
 from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
@@ -381,7 +382,7 @@ def image_transformer2d_base():
   hparams.optimizer_adam_beta1 = 0.9
   hparams.optimizer_adam_beta2 = 0.98
   hparams.label_smoothing = 0.0
-  hparams.target_modality = "image:identity"
+  hparams.modality["targets"] = modalities.IdentityModality
   hparams.norm_type = "layer"
   hparams.layer_prepostprocess_dropout = 0.0
   hparams.add_hparam("filter_size", 512)  # Add new ones like this.

tensor2tensor/models/mtf_transformer.py

@@ -24,6 +24,7 @@
 
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
+from tensor2tensor.layers import modalities
 from tensor2tensor.models.research import moe
 from tensor2tensor.utils import mtf_model
 from tensor2tensor.utils import registry
@@ -772,8 +773,10 @@ def mtf_transformer_base():
   # These parameters make Transformer model compatible with MtfTransformer
   # Do not override these, as mtf_transformer does not support other options.
   hparams.clip_grad_norm = 0.  # i.e. no gradient clipping
-  hparams.target_modality = "symbol:identity"
-  hparams.input_modalities = "inputs:symbol:identity"
+  hparams.modality = {
+      "inputs": modalities.IdentitySymbolModality,
+      "targets": modalities.IdentitySymbolModality,
+  }
 
   # Parameters for computing the maximum decode length in beam search.
   # Maximum decode length is:

tensor2tensor/models/research/autoencoders.py

@@ -24,6 +24,7 @@
 from tensor2tensor.layers import common_layers
 from tensor2tensor.layers import discretization
 from tensor2tensor.layers import latent_layers
+from tensor2tensor.layers import modalities
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
@@ -1104,8 +1105,10 @@ def autoencoder_residual_text():
   hparams.hidden_size = 64
   hparams.max_hidden_size = 512
   hparams.bottleneck_noise = 0.0
-  hparams.target_modality = "symbol:identity"
-  hparams.input_modalities = "symbol:identity"
+  hparams.modality = {
+      "inputs": modalities.IdentitySymbolModality,
+      "targets": modalities.IdentitySymbolModality,
+  }
   hparams.autoregressive_mode = "none"
   hparams.sample_width = 1
   return hparams
@@ -1209,8 +1212,10 @@ def autoencoder_ordered_text():
   hparams.batch_size = 1024
   hparams.autoregressive_mode = "conv5"
   hparams.max_hidden_size = 1024
-  hparams.target_modality = "symbol:identity"
-  hparams.input_modalities = "symbol:identity"
+  hparams.modality = {
+      "inputs": modalities.IdentitySymbolModality,
+      "targets": modalities.IdentitySymbolModality,
+  }
   hparams.sample_height = 128
   hparams.sample_width = 1
   return hparams