Don't export functions when saving hparams to json.

PiperOrigin-RevId: 237176017

Author: Dustin Tran

tensor2tensor/utils/hparam.py

@@ -542,15 +542,15 @@ def to_json(self, indent=None, separators=None, sort_keys=False):
       A JSON string.
     """
     def remove_callables(x):
-      if callable(x):
-        return x.__name__
+      """Omit callable elements from input with arbitrary nesting."""
       if isinstance(x, dict):
-        return {k: remove_callables(v) for k, v in six.iteritems(x)}
-      if isinstance(x, list):
-        return [remove_callables(i) for i in x]
+        return {k: remove_callables(v) for k, v in six.iteritems(x)
+                if not callable(v)}
+      elif isinstance(x, list):
+        return [remove_callables(i) for i in x if not callable(i)]
       return x
     return json.dumps(
-        {k: remove_callables(v) for k, v in six.iteritems(self.values())},
+        remove_callables(self.values()),
         indent=indent,
         separators=separators,
         sort_keys=sort_keys)

tensor2tensor/utils/hparam_test.py

@@ -136,6 +136,15 @@ def testSetFromMap(self):
     self.assertDictEqual({'d': [0.1, 0.2, 0.3], 'x': 1, 'b': 2.0},
                          hparams.values())
 
+  def testFunction(self):
+    def f(x):
+      return x
+    hparams = hparam.HParams(function=f)
+    self.assertEqual(hparams.function, f)
+
+    json_str = hparams.to_json()
+    self.assertEqual(json_str, '{}')
+
   def testBoolParsing(self):
     for value in 'true', 'false', 'True', 'False', '1', '0':
       for initial in False, True:

tensor2tensor/utils/hparams_lib.py

@@ -64,6 +64,14 @@ def create_hparams_from_json(json_path, hparams=None):
   tf.logging.info("Loading hparams from existing json %s" % json_path)
   with tf.gfile.Open(json_path, "r") as f:
     hparams_values = json.load(f)
+    # Prevent certain keys from overwriting the passed-in hparams.
+    # TODO(trandustin): Remove this hack after registries are available to avoid
+    # saving them as functions.
+    hparams_values.pop("bottom", None)
+    hparams_values.pop("loss", None)
+    hparams_values.pop("name", None)
+    hparams_values.pop("top", None)
+    hparams_values.pop("weights_fn", None)
     new_hparams = HParams(**hparams_values)
     # Some keys are in new_hparams but not hparams, so we need to be more
     #   careful than simply using parse_json() from HParams