Moving gin_register_keras_layers() to nn.py.

Ethan Manilow · Magenta Team · commit f0756291c2c8 · 2022-04-13T13:31:23.000-07:00
PiperOrigin-RevId: 441565076
diff --git a/ddsp/training/ddsp_run.py b/ddsp/training/ddsp_run.py
@@ -183,7 +183,6 @@ def main(unused_argv):
   gfile.makedirs(restore_dir)  # Only makes dirs if they don't exist.
   parse_gin(restore_dir)
   logging.info('Operative Gin Config:\n%s', gin.config.config_str())
-  train_util.gin_register_keras_layers()
 
   if FLAGS.allow_memory_growth:
     allow_memory_growth()
diff --git a/ddsp/training/nn.py b/ddsp/training/nn.py
@@ -32,6 +32,21 @@
 # pylint: disable=redundant-keyword-arg
 
 
+def gin_register_keras_layers():
+  """Registers all keras layers and Sequential to be referenceable in gin."""
+  # Register sequential model.
+  gin.external_configurable(tf.keras.Sequential, 'tf.keras.Sequential')
+
+  # Register all the layers.
+  for k, v in inspect.getmembers(tf.keras.layers):
+    # Duck typing for tf.keras.layers.Layer since keras uses metaclasses.
+    if hasattr(v, 'variables'):
+      gin.external_configurable(v, f'tf.keras.layers.{k}')
+
+
+gin_register_keras_layers()
+
+
 class DictLayer(tfkl.Layer):
   """Wrap a Keras Layer to take dictionary inputs and outputs.
 
diff --git a/ddsp/training/train_util.py b/ddsp/training/train_util.py
@@ -15,7 +15,6 @@
 # Lint as: python3
 """Library of training functions."""
 
-import inspect
 import json
 import os
 import time
@@ -209,18 +208,6 @@ def format_for_tensorboard(line):
     summary_writer.flush()
 
 
-def gin_register_keras_layers():
-  """Registers all keras layers and Sequential to be referenceable in gin."""
-  # Register sequential model.
-  gin.external_configurable(tf.keras.Sequential, 'tf.keras.Sequential')
-
-  # Register all the layers.
-  for k, v in inspect.getmembers(tf.keras.layers):
-    # Duck typing for tf.keras.layers.Layer since keras uses metaclasses.
-    if hasattr(v, 'variables'):
-      gin.external_configurable(v, f'tf.keras.layers.{k}')
-
-
 # ------------------------ Training Loop ---------------------------------------
 @gin.configurable
 def train(data_provider,
