如何在 TensorFlow 2 中保存/加载模型的一部分?

Posted

技术标签:

【中文标题】如何在 TensorFlow 2 中保存/加载模型的一部分?【英文标题】:How to save/load part of a model in TensorFlow 2? 【发布时间】:2020-06-02 18:43:36 【问题描述】:

我在 TensorFlow 2 中制作了一个自动编码器。首先,我需要在数据集上训练模型。训练自动编码器后,我想将自动编码器的编码器部分用于不同的数据集,同时保留训练后的权重。有没有办法这样做?我找不到一个。任何指导方针都会有所帮助。这是我正在尝试使用的粗略自动编码器:

from tensorflow.keras.layers import Input, Dense
from tensorflow.keras import Model

def Autoencoder():
    # Encoder
    inp = Input(shape=(100,))
    x = Dense(1000)(inp)
    x = Dense(500)(x)
    code = Dense(10)(x)

    # Decoder    
    x = Dense(500)(code)
    x = Dense(1000)(x)
    out = Dense(100)(x)

    model = Model(inp,out)
    return model

autoencoder = Autoencoder()

【问题讨论】:

【参考方案1】:

如果您只想使用编码器部分,您可以编写一个 CustomCheckpoint 类并保存编码器,如下所示。

class CustomCheckpoint(Callback):

def __init__(self, filepath, encoder):
    self.monitor = 'val_loss'
    self.monitor_op = np.less
    self.best = np.Inf

    self.filepath = filepath
    self.encoder = encoder

def on_epoch_end(self, epoch, logs=None):
    current = logs.get(self.monitor)
    if self.monitor_op(current, self.best):
        self.best = current
        # self.encoder.save_weights(self.filepath, overwrite=True)
        self.encoder.save(self.filepath, overwrite=True) # Whichever you prefer  

要从保存的权重中加载编码器部分,您可以执行以下代码。

from keras.models import load_model
autoencoder = load_model("path_to_file")
encoder = Model(autoencoder.layers[0].input, autoencoder.layers[1].output)

【讨论】:

以上是关于如何在 TensorFlow 2 中保存/加载模型的一部分?的主要内容,如果未能解决你的问题,请参考以下文章

转 tensorflow模型保存 与 加载

如何使用 c++ 在 tensorflow 中保存模型

Tensorflow加载预训练模型和保存模型

如何保存经过训练的模型(估计器)并将其加载回来以使用 Tensorflow 中的数据对其进行测试?

在 TensorFlow Functional API 中保存和加载具有相同图形的多个模型

TensorFlow 保存和加载不一致