Сохранение и загрузка пользовательских классов потерь Tensorflow - PullRequest
Новая
       14

Сохранение и загрузка пользовательских классов потерь Tensorflow

0 голосов
/ 21 апреля 2020

Я пытаюсь сохранить и загрузить модель с пользовательским классом потерь. Это простая модель nn с пользовательской функцией потерь (класс). Вот код (Код взят из книги Аурелиана Жерона «Руки на ML 2», глава 12): 

import tensorflow as tf
import tensorflow.keras as keras
from sklearn.datasets import fetch_california_housing
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

class HuberLoss(keras.losses.Loss):
    def __init__(self, threshold=1.0, **kwargs):
        self.threshold = threshold
        super().__init__(**kwargs)
    def call(self, y_true, y_pred):
        error = y_true - y_pred
        is_small_error = tf.abs(error) < self.threshold
        squared_loss = tf.square(error) / 2
        linear_loss  = self.threshold * tf.abs(error) - self.threshold**2 / 2
        return tf.where(is_small_error, squared_loss, linear_loss)
    def get_config(self):
        base_config = super().get_config()
        return {**base_config, "threshold": self.threshold}

housing = fetch_california_housing()
X_train_full, X_test, y_train_full, y_test = train_test_split(
    housing.data, housing.target.reshape(-1, 1), random_state=42)
X_train, X_valid, y_train, y_valid = train_test_split(
    X_train_full, y_train_full, random_state=42)


scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_valid_scaled = scaler.transform(X_valid)
X_test_scaled = scaler.transform(X_test)


input_shape = X_train.shape[1:]

model = keras.models.Sequential([
    keras.layers.Dense(30, activation="selu", kernel_initializer="lecun_normal",
                       input_shape=input_shape),
    keras.layers.Dense(1),
])

model.compile(loss=HuberLoss(2.), optimizer="nadam", metrics=["mae"])

model.fit(X_train_scaled, y_train, epochs=2,
          validation_data=(X_valid_scaled, y_valid))

Модель прекрасно компилируется без ошибок. Вывод: 

Train on 11610 samples, validate on 3870 samples
Epoch 1/2
11610/11610 [==============================] - 5s 422us/sample - loss: 0.7539 - mae: 0.9434 - val_loss: 0.5205 - val_mae: 0.6497
Epoch 2/2
11610/11610 [==============================] - 2s 214us/sample - loss: 0.2592 - mae: 0.5231 - val_loss: 0.4114 - val_mae: 0.5850
<tensorflow.python.keras.callbacks.History at 0x1be93b8fc48>

Затем сохранил модель так: 

model.save("my_model_1.h5")

Но когда я пытаюсь загрузить модель, используя словарь keras.models.load_model и custom_objects, выдает ошибку : 

model = keras.models.load_model(
    "my_model_1.h5", 
    custom_objects={
        'HuberLoss': HuberLoss
    })

Вывод: 

    ---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-3-a00475905a88> in <module>
      2     "my_model_1.h5",
      3     custom_objects={
----> 4         'HuberLoss': HuberLoss
      5     })

~\.conda\envs\HandsOnML2\lib\site-packages\tensorflow_core\python\keras\saving\save.py in load_model(filepath, custom_objects, compile)
    144   if (h5py is not None and (
    145       isinstance(filepath, h5py.File) or h5py.is_hdf5(filepath))):
--> 146     return hdf5_format.load_model_from_hdf5(filepath, custom_objects, compile)
    147 
    148   if isinstance(filepath, six.string_types):

~\.conda\envs\HandsOnML2\lib\site-packages\tensorflow_core\python\keras\saving\hdf5_format.py in load_model_from_hdf5(filepath, custom_objects, compile)
    182       # Compile model.
    183       model.compile(**saving_utils.compile_args_from_training_config(
--> 184           training_config, custom_objects))
    185 
    186       # Set optimizer weights.

~\.conda\envs\HandsOnML2\lib\site-packages\tensorflow_core\python\keras\saving\saving_utils.py in compile_args_from_training_config(training_config, custom_objects)
    232   loss_config = training_config['loss']  # Deserialize loss class.
    233   if isinstance(loss_config, dict) and 'class_name' in loss_config:
--> 234     loss_config = losses.get(loss_config)
    235   loss = nest.map_structure(
    236       lambda obj: custom_objects.get(obj, obj), loss_config)

~\.conda\envs\HandsOnML2\lib\site-packages\tensorflow_core\python\keras\losses.py in get(identifier)
   1184     return deserialize(identifier)
   1185   if isinstance(identifier, dict):
-> 1186     return deserialize(identifier)
   1187   elif callable(identifier):
   1188     return identifier

~\.conda\envs\HandsOnML2\lib\site-packages\tensorflow_core\python\keras\losses.py in deserialize(name, custom_objects)
   1173       module_objects=globals(),
   1174       custom_objects=custom_objects,
-> 1175       printable_module_name='loss function')
   1176 
   1177 

~\.conda\envs\HandsOnML2\lib\site-packages\tensorflow_core\python\keras\utils\generic_utils.py in deserialize_keras_object(identifier, module_objects, custom_objects, printable_module_name)
    290     config = identifier
    291     (cls, cls_config) = class_and_config_for_serialized_keras_object(
--> 292         config, module_objects, custom_objects, printable_module_name)
    293 
    294     if hasattr(cls, 'from_config'):

~\.conda\envs\HandsOnML2\lib\site-packages\tensorflow_core\python\keras\utils\generic_utils.py in class_and_config_for_serialized_keras_object(config, module_objects, custom_objects, printable_module_name)
    248     cls = module_objects.get(class_name)
    249     if cls is None:
--> 250       raise ValueError('Unknown ' + printable_module_name + ': ' + class_name)
    251 
    252   cls_config = config['config']

ValueError: Unknown loss function: HuberLoss

Как это исправить?

...