Не удается сохранить метод в классе при загрузке модели в Tensorflow 2.0 - PullRequest
Новая
       107

Не удается сохранить метод в классе при загрузке модели в Tensorflow 2.0

1 голос
/ 14 июля 2020

Я сохраняюсь, пытаясь сохранить кастомную модель в Tensorflow 2.0. Я следовал руководству от Tensorflow.

Создание модели класса: 

class CustomModel(keras.Model):
    def __init__(self, hidden_units):
        super(CustomModel, self).__init__()
        self.dense_layers = [keras.layers.Dense(u) for u in hidden_units]

    def call(self, inputs):
        x = inputs
        for layer in self.dense_layers:
            x = layer(x)
        return x
    
    def get_layer_size(self):
        return len(self.dense_layers)

Я добавил метод get_layer_size для возврата длины слоев. Когда я сохраняю и загружаю модель. 

model = CustomModel([16, 16, 10])
# Build the model by calling it
input_arr = tf.random.uniform((1, 5))
outputs = model(input_arr)
# model.save("my_model1")
keras.models.save_model(model, filepath="model_1")

# Delete the custom-defined model class to ensure that the loader does not have
# access to it.
del CustomModel

loaded = keras.models.load_model("model_1")
np.testing.assert_allclose(loaded(input_arr), outputs)

print("Original model:", model)
print("Loaded model:", loaded)

Проблема, когда я вызываю метод get_layer_size.

Исходная модель: model.get_layer_size() return 3. Однако loaded.get_layer_size() возвращает ошибку

----------------------------------------- ---------------------------------- AttributeError Traceback (последний вызов последним) в () ----> 1 loaded.get_layer_size ()

AttributeError: объект 'CustomModel' не имеет атрибута 'get_layer_size'

Ссылка: https://colab.research.google.com/drive/14k4ehLdd_Yb-g8x9W3ZB9YLdJGcbIdMv?usp=sharing

Как мне решить эту проблему? Большое спасибо!

...