如何在 Tensorflow 2.0 中保存在 GPU 上训练以加载到 CPU 上的模型答案

【问题标题】：How do I save a model in Tensorflow 2.0 that was trained on a GPU to be loaded on a CPU如何在 Tensorflow 2.0 中保存在 GPU 上训练以加载到 CPU 上的模型
【发布时间】：2025-05-18 12:15:02
【问题描述】：

我有一袋模型，我在一台带有 GPU 的机器上训练并保存了这些模型。我使用了以下经过训练并在 GPU 上运行的模型

    model = Sequential()
    model.add(CuDNNLSTM(units=30,input_shape=(None, 11), return_sequences=True, name='LAYER1'))
    model.add(Dropout(.9, name='LAYER2'))
    model.add(Dense(units=10, activation="relu",name='LAYER3'))
    model.add(Dropout(.1, name='LAYER4'))
    model.add(CuDNNLSTM(units=20,return_sequences=False,name='LAYER5'))
    model.add(Dropout(.1, name='LAYER6'))
    model.add(Dense(units=3,activation="linear",name='LEVEL7'))
    rmsprop_opt = RMSprop(lr=learning_rate)
    model.compile(loss="mse", optimizer=rmsprop_opt)

我使用以下方法保存模型图：

    model_json_dict = json.loads(model.to_json())
    json.dump(model_json_dict, open("my_model_graph.json", "w"))

然后我使用检查点方法保存了权重：

    callback_checkpoint = ModelCheckpoint(filepath="model_checkpoint.h5",
                                          monitor='val_loss',
                                          verbose=1,
                                          save_weights_only=True,
                                          save_best_only=True)

    callbacks = [callback_checkpoint]

我使用以下方法拟合模型：

    history = model.fit(feature_train,
                        label_train,
                        batch_size=batch_size,
                        epochs=epochs,
                        validation_split=validation_split,
                        callbacks=callbacks)

我想将模型读回到只有 CPU 的机器中进行预测。我已经在第二台机器上加载了模型并按如下方式加载了权重，TF 抱怨 CPU/GPU 问题。

    model = model_from_json(json.dumps(json.load(open("my_model_graph.json","r"))))
    model.load_weights("model_checkpoint.h5")

所以问题是我如何将这些保存的模型及其权重转换为可以重新加载到只有 CPU 的第二台机器的形式？

它对执行此操作的正确方法感到困惑。有一个 SO 显示使用 Saver() 类。 Tensorflow: how to save/restore a model?。另一个帖子说它不能完成，另一个帖子说它是透明的。转换这些现有模型的推荐方法是什么？（再培训他们不是一种选择！）

【问题讨论】：

你得到了什么确切的错误？很可能是你使用CuDNNLSTM有问题，只能在GPU上运行。
看来我能做到这一点的唯一方法就是重建图表。
数据图的json加载不起作用。

标签： python tensorflow keras cross-platform

【解决方案1】：

我就是这样解决的。我的模型看起来像这样：

def build_model(layers, machine, learning_rate, dropout_rate):
    model = Sequential() #   tf.keras.models.Sequential()
    if machine == "GPU":
        model.add(
            CuDNNLSTM(
                units=layers[1],
                input_shape=(None, layers[0]),
                return_sequences=True,
                name='FIRST_LAYER')
        )
    else:
        model.add(
            LSTM(
                units=layers[1],
                input_shape=(None, layers[0]),
                return_sequences=True,
                name='FIRST_LAYER')
        )
...

我这样创建初始模型：

my_model = build_model(layers=layer_sizes, machine='GPU', learning_rate=0.003,dropout_rate=0.05)

...然后训练模型并使用以下方法保存权重：

my_model.save_weights("my_model_weights.h5")

现在我切换到CPU实例如下：

my_model = build_model(layers=layer_sizes, machine='CPU', learning_rate=0.003,dropout_rate=0.05)

...然后我可以按如下方式加载保存的权重：

my_model.load_weights("my_model_weights.h5")

很遗憾 model_from_json api 不能跨机器工作！

【讨论】：