如何保存模型在训练期间创建的嵌入?

时间:2021-03-22 21:47:36

标签: python tensorflow keras

我正在处理图像相似性问题,并希望保存模型在训练期间创建的图像嵌入。有没有办法在嵌入传递到损失函数之前捕获嵌入?

这是我的模型

model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(filters=128, kernel_size=2, padding='same', activation='relu', input_shape=(32,32,3)),
    tf.keras.layers.MaxPooling2D(pool_size=2),
    tf.keras.layers.Dropout(0.3),
    tf.keras.layers.Conv2D(filters=64, kernel_size=2, padding='same', activation='relu'),
    tf.keras.layers.MaxPooling2D(pool_size=2),
    tf.keras.layers.Dropout(0.3),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(32, activation=None), # No activation on final dense layer
    tf.keras.layers.Lambda(lambda x: tf.math.l2_normalize(x, axis=1)) # L2 normalize embeddings

])

model.compile(
    optimizer=tf.keras.optimizers.Adam(0.001),
    loss=tfa.losses.TripletSemiHardLoss(),
    metrics = ["accuracy"])

history = model.fit(train_dataset, epochs=3, validation_data=test_dataset)

明确地说,我不希望仅显示来自最后一层的输出 here. 我想保存从我的模型输出的最终结果向量。

1 个答案:

答案 0 :(得分:1)

您可以创建一个自定义回调,并在每个批次的训练步骤结束时,使用输入批次调用您的模型并获取输出并保存。

class SaveEmbeddingCallback(tf.keras.callbacks.Callback):
    def on_train_batch_end(self, batch, logs=None):
        embedding = self.model.predict(batch)
        # IN THIS STAGE YOU HAVE THE OUTPUT OF THE MODEL
        # YOU CAN SAVE IT OR WHATEVER YOU WANT
        ...


model.compile(
    optimizer=tf.keras.optimizers.Adam(0.001),
    loss=tfa.losses.TripletSemiHardLoss(),
    metrics = ["accuracy"])

history = model.fit(train_dataset,
    epochs=3,
    validation_data=test_dataset,
    callbacks=[SaveEmbeddingCallback()])

有关 keras 自定义回调的更多信息,请阅读此 TensorFlow tutorial

相关问题