tensorflow 2.0 技巧 | 自定义tf.keras.Model的坑

自定义tf.keras.Model需要注意的点

model.save()

• subclass Model 是不能直接save的，save成.h5，但是能够save_weights，或者save_format="tf"

NotImplementedError: Saving the model to HDF5 format requires the model to be a Functional model or a Sequential model. It does not work for subclassed models, because such models are defined via the body of a Python method, which isn't safely serializable. Consider saving to the Tensorflow SavedModel format (by setting save_format="tf") or using `save_weights`.

model.trainable_variables

• __init__若没有注册该layers，那么在后面应用梯度时会找不到model.trainable_variables。
  
  像下面这样是不行的：

class Map_model(tf.keras.Model):
    def __init__(self, is_train=False):
        super(Map_model, self).__init__()
    def call(self, x):
        x = tf.keras.layers.Dense(10, activation='relu')
        return x

model.summary()

• 需要先指定input_shape，或者你直接fit一遍它也能自动确定

    model.build(input_shape=(None, 448, 448, 3))
    print(model.summary())

class Map_model(tf.keras.Model):
    def __init__(self, is_train=False):
        super(Map_model, self).__init__()
        self.map_f1 = tf.keras.layers.Dense(10, activation='relu', trainable=is_train)
        # self.map_f2 = tf.keras.layers.Dense(6, activation='relu')
        self.map_f3 = tf.keras.layers.Dense(3, activation='softmax', trainable=is_train)

    def call(self, x):
        x = self.map_f1(x)
        # x = self.map_f2(x)
        return self.map_f3(x)

@tf.function
def train_step(mmodel, label, L_label, loss_object, train_loss, train_accuracy, optimizer):
    with tf.GradientTape() as tape:
        L_label_pred = mmodel(label)
        loss = loss_object(L_label, L_label_pred)
    gradient_l = tape.gradient(loss, mmodel.trainable_variables)
    train_loss(loss)
    train_accuracy(L_label, L_label_pred)
    optimizer.apply_gradients(zip(gradient_l, mmodel.trainable_variables))

def train():
    mmodel = Map_model(is_train=True)
    optimizer = tf.keras.optimizers.Adam(0.01)
    loss_object = tf.keras.losses.CategoricalCrossentropy()
    train_loss = tf.keras.metrics.Mean(name='train_loss')
    train_accuracy = tf.keras.metrics.CategoricalAccuracy(name='train_accuracy')

    EPOCHS = 0
    labels = range(1, 30)  # labels = truth_label -1
    L_labels = [int(prpcs.map2Lclass(l)) for l in labels]
    labels = [l - 1 for l in labels]
    labels_onehot = tf.one_hot(labels, depth=29)
    L_labels_onehot = tf.one_hot(L_labels, depth=3)
    EPS = 1e-6
    loss_e = 0x7f7f7f
    while loss_e > EPS:
        EPOCHS += 1
        train_loss.reset_states()
        train_accuracy.reset_states()
        train_step(mmodel, labels_onehot, L_labels_onehot, loss_object, train_loss, train_accuracy, optimizer)

        template = 'Epoch {}, Loss: {}, Accuracy: {}'
        print(template.format(EPOCHS,
                              train_loss.result(),
                              train_accuracy.result() * 100))
        loss_e = train_loss.result()
    print("labels_onehot shape:", labels_onehot.shape)
    model_path = r'./models/'
    if not os.path.exists(model_path):
        os.makedirs(model_path)
    mmodel.save(os.path.join(model_path, 'map_model_{}'.format(EPS)))
    mmodel.save_weights(os.path.join(model_path, 'map_model_weights_{}'.format(EPS)))
    print("Save model!")