tensorflow2.0 | 自定义Model

记录一下自定义keras模型的过程中遇到的坑。

1、自定义模型

自定义模型，要继承keras.Model，然后实现__init__和call两个方法，注意以下坑点：
__init__若没有定义layers，那么在后面应用梯度时会找不到trainable_variables；因此所有层的定义一定要在__init__里定义。

import tensroflow.keras as keras

class model(keras.Model):

    def __init__(self):
        super(VAE, self).__init__()
        self.dense1 = keras.layers.Dense(16, activation=tf.nn.relu)
        self.dense2 = keras.layers.Dense(8, activation=tf.nn.relu)
        self.dense3 = keras.layers.Dense(1)

    def call(self, inputs, training=None, mask=None):
        h = self.dense1(inputs)
        h = self.dense2(h)
        return self.dense3(h)

2、训练自定义模型

model = model()
model.build(input_shape=(None, 1))
model.summary()  # 必须在build之后
train_op = keras.optimizers.Adam(...)
model.compile(...)
model.fit(...)

3、保存自定义模型

注意保存模型时的坑：
subclass Model 是不能直接save的，但是能够save_weights，或者save_format=“tf”。按如下可以正确保存：

model.save_weights("model_weights.h5")

4、加载保存好的模型

加载模型后，要使用模型必需先调用build指定input_shape，不然模型无法确定输入数据的维度，从而无法创建对应的weight矩阵。

model = model()
model.build(input_shape=(None, 1))
model.summary()  # 必须在build之后
model.load_weights("model_weights.h5")  # 必须在build之后
# 加载模型权重后，可以继续训练，也可以做预测
y = model.predict(x)