🐽 Add save_pretrained for TFAutoModel.

Author: dathudeptrai

tensorflow_tts/configs/base_config.py

@@ -22,11 +22,11 @@
 
 
 class BaseConfig(abc.ABC):
-    def set_pretrained_config(self, config):
-        self.config = config
+    def set_config_params(self, config_params):
+        self.config_params = config_params
 
     def save_pretrained(self, saved_path):
         """Save config to file"""
         os.makedirs(saved_path, exist_ok=True)
         with open(os.path.join(saved_path, CONFIG_FILE_NAME), "w") as file:
-            yaml.dump(self.config, file)
+            yaml.dump(self.config_params, file, Dumper=yaml.Dumper)

tensorflow_tts/inference/auto_config.py

@@ -72,13 +72,13 @@ def from_pretrained(cls, pretrained_path, **kwargs):
             )
 
         with open(pretrained_path) as f:
-            config = yaml.load(f, Loader=yaml.SafeLoader)
+            config = yaml.load(f, Loader=yaml.Loader)
 
         try:
             model_type = config["model_type"]
             config_class = CONFIG_MAPPING[model_type]
             config_class = config_class(**config[model_type + "_params"], **kwargs)
-            config_class.set_pretrained_config(config)
+            config_class.set_config_params(config)
             return config_class
         except Exception:
             raise ValueError(

tensorflow_tts/inference/auto_model.py

@@ -17,6 +17,7 @@
 import logging
 import warnings
 import os
+import copy
 
 from collections import OrderedDict
 
@@ -67,9 +68,7 @@ def __init__(self):
         raise EnvironmentError("Cannot be instantiated using `__init__()`")
 
     @classmethod
-    def from_pretrained(cls, config=None, pretrained_path=None, **kwargs):
-        is_build = kwargs.pop("is_build", True)
-
+    def from_pretrained(cls, pretrained_path=None, config=None, **kwargs):
         # load weights from hf hub
         if pretrained_path is not None:
             if not os.path.isfile(pretrained_path):
@@ -101,8 +100,8 @@ def from_pretrained(cls, config=None, pretrained_path=None, **kwargs):
                 config
             ):
                 model = model_class(config=config, **kwargs)
-                if is_build:
-                    model._build()
+                model.set_config(config)
+                model._build()
                 if pretrained_path is not None and ".h5" in pretrained_path:
                     try:
                         model.load_weights(pretrained_path)

tensorflow_tts/models/__init__.py

@@ -1,3 +1,4 @@
+from tensorflow_tts.models.base_model import BaseModel
 from tensorflow_tts.models.fastspeech import TFFastSpeech
 from tensorflow_tts.models.fastspeech2 import TFFastSpeech2
 from tensorflow_tts.models.melgan import (

tensorflow_tts/models/base_model.py

@@ -0,0 +1,33 @@
+# -*- coding: utf-8 -*-
+# Copyright 2020 TensorFlowTTS Team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Base Model for all model."""
+
+import tensorflow as tf
+import yaml
+import os
+import numpy as np
+
+from tensorflow_tts.utils.utils import MODEL_FILE_NAME, CONFIG_FILE_NAME
+
+
+class BaseModel(tf.keras.Model):
+    def set_config(self, config):
+        self.config = config
+
+    def save_pretrained(self, saved_path):
+        """Save config and weights to file"""
+        os.makedirs(saved_path, exist_ok=True)
+        self.config.save_pretrained(saved_path)
+        self.save_weights(os.path.join(saved_path, MODEL_FILE_NAME))

tensorflow_tts/models/fastspeech.py

@@ -17,6 +17,8 @@
 import numpy as np
 import tensorflow as tf
 
+from tensorflow_tts.models import BaseModel
+
 
 def get_initializer(initializer_range=0.02):
     """Creates a `tf.initializers.truncated_normal` with the given range.
@@ -746,7 +748,7 @@ def body(
         return outputs, encoder_masks
 
 
-class TFFastSpeech(tf.keras.Model):
+class TFFastSpeech(BaseModel):
     """TF Fastspeech module."""
 
     def __init__(self, config, **kwargs):

tensorflow_tts/models/hifigan.py

@@ -23,6 +23,7 @@
 from tensorflow_tts.utils import GroupConv1D
 from tensorflow_tts.utils import WeightNormalization
 
+from tensorflow_tts.models import BaseModel
 from tensorflow_tts.models import TFMelGANGenerator
 
 
@@ -133,7 +134,7 @@ def call(self, x, training=False):
         return xs / len(self.list_resblock)
 
 
-class TFHifiGANGenerator(tf.keras.Model):
+class TFHifiGANGenerator(BaseModel):
     def __init__(self, config, **kwargs):
         super().__init__(**kwargs)
         # check hyper parameter is valid or not
@@ -338,7 +339,7 @@ def _apply_weightnorm(self, list_layers):
                 pass
 
 
-class TFHifiGANMultiPeriodDiscriminator(tf.keras.Model):
+class TFHifiGANMultiPeriodDiscriminator(BaseModel):
     """Tensorflow Hifigan Multi Period discriminator module."""
 
     def __init__(self, config, **kwargs):

tensorflow_tts/models/mb_melgan.py

@@ -21,6 +21,7 @@
 import tensorflow as tf
 from scipy.signal import kaiser
 
+from tensorflow_tts.models import BaseModel
 from tensorflow_tts.models import TFMelGANGenerator
 
 

tensorflow_tts/models/melgan.py

@@ -17,6 +17,7 @@
 import numpy as np
 import tensorflow as tf
 
+from tensorflow_tts.models import BaseModel
 from tensorflow_tts.utils import GroupConv1D, WeightNormalization
 
 
@@ -186,7 +187,7 @@ def _apply_weightnorm(self, list_layers):
                 pass
 
 
-class TFMelGANGenerator(tf.keras.Model):
+class TFMelGANGenerator(BaseModel):
     """Tensorflow MelGAN generator module."""
 
     def __init__(self, config, **kwargs):
@@ -450,7 +451,7 @@ def _apply_weightnorm(self, list_layers):
                 pass
 
 
-class TFMelGANMultiScaleDiscriminator(tf.keras.Model):
+class TFMelGANMultiScaleDiscriminator(BaseModel):
     """MelGAN multi-scale discriminator module."""
 
     def __init__(self, config, **kwargs):

tensorflow_tts/models/parallel_wavegan.py

@@ -17,6 +17,8 @@
 
 import tensorflow as tf
 
+from tensorflow_tts.models import BaseModel
+
 
 def get_initializer(initializer_seed=42):
     """Creates a `tf.initializers.he_normal` with the given seed.
@@ -345,7 +347,7 @@ def call(self, c):
         return self.upsample(c_)
 
 
-class TFParallelWaveGANGenerator(tf.keras.Model):
+class TFParallelWaveGANGenerator(BaseModel):
     """Parallel WaveGAN Generator module."""
 
     def __init__(self, config, **kwargs):
@@ -491,7 +493,7 @@ def inference(self, mels):
         return x
 
 
-class TFParallelWaveGANDiscriminator(tf.keras.Model):
+class TFParallelWaveGANDiscriminator(BaseModel):
     """Parallel WaveGAN Discriminator module."""
 
     def __init__(self, config, **kwargs):

tensorflow_tts/models/tacotron2.py

@@ -27,6 +27,8 @@
 
 from tensorflow_tts.utils import dynamic_decode
 
+from tensorflow_tts.models import BaseModel
+
 
 def get_initializer(initializer_range=0.02):
     """Creates a `tf.initializers.truncated_normal` with the given range.
@@ -737,7 +739,7 @@ def step(self, time, inputs, state, training=False):
         return (outputs, next_state, next_inputs, finished)
 
 
-class TFTacotron2(tf.keras.Model):
+class TFTacotron2(BaseModel):
     """Tensorflow tacotron-2 model."""
 
     def __init__(self, config, **kwargs):
@@ -760,10 +762,10 @@ def __init__(self, config, **kwargs):
             units=config.n_mels, name="residual_projection"
         )
 
-        self.config = config
         self.use_window_mask = False
         self.maximum_iterations = 4000
         self.enable_tflite_convertible = enable_tflite_convertible
+        self.config = config
 
     def setup_window(self, win_front, win_back):
         """Call only for inference."""

test/test_auto.py

@@ -42,6 +42,8 @@
 )
 def test_auto_processor(mapper_path):
     processor = AutoProcessor.from_pretrained(pretrained_path=mapper_path)
+    processor.save_pretrained("./test_saved")
+    processor = AutoProcessor.from_pretrained("./test_saved/processor.json")
 
 
 @pytest.mark.parametrize(
@@ -65,7 +67,12 @@ def test_auto_processor(mapper_path):
 )
 def test_auto_model(config_path):
     config = AutoConfig.from_pretrained(pretrained_path=config_path)
-    model = TFAutoModel.from_pretrained(config=config, pretrained_path=None)
+    model = TFAutoModel.from_pretrained(pretrained_path=None, config=config)
 
     # test save_pretrained
-    config.save_pretrained("./")
+    config.save_pretrained("./test_saved")
+    model.save_pretrained("./test_saved")
+
+    # test from_pretrained
+    config = AutoConfig.from_pretrained("./test_saved/config.yml")
+    model = TFAutoModel.from_pretrained("./test_saved/model.h5", config=config)