将训练有素的 Tensorflow 模型保存到另一台机器上进行推理 [重复]

Posted 2023-02-23

【中文标题】将训练有素的 Tensorflow 模型保存到另一台机器上进行推理 [重复]

【英文标题】：Saving trained Tensorflow model to inference on another machine [duplicate]

【发布时间】：2019-01-11 16:48:34

【问题描述】：

我对机器学习和 Tensorflow 框架比较陌生。我试图使用 MNIST 手写数字数据集，对我创建的测试示例进行推理，从而使我的训练模型受到here 提供的代码的严重影响。但是，我正在使用 GPU 的远程机器上进行训练，并尝试将数据保存到一个目录中，以便我可以在本地机器上传输数据和推理

似乎我可以使用tf.saved_model.simple_save 保存一些模型，但是，我不确定如何使用保存的数据进行推理以及使用这些数据在给定新图像的情况下进行预测。似乎有多种方法可以保存模型，但我不确定使用 Tensorflow 框架的约定或“正确方法”是什么。

到目前为止，这是我认为我需要的行，但不确定它是否正确。

            tf.saved_model.simple_save(sess, 'mnist_model',                                                                                 
                inputs='x': self.x,                                                                                                   
                outputs='y_': self.y_, 'y_conv':self.y_conv)                                                                          

如果有人能指出如何正确保存经过训练的模型以及使用哪些变量来使用保存的模型进行推断的方向，我将不胜感激。

【参考方案1】：

您可以这样做的一种方法是在您的图形定义中创建一个tf.train.Saver() 对象，然后使用它将网络保存到指定目录。然后可以将该目录中的权重从远程计算机下载到本地计算机并在本地恢复。这是一个小示例网络：

import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data

mnist = input_data.read_data_sets('MNIST_data', one_hot=True)


# >>>> Config. Vars <<<<

TRAIN_STEPS = 1000

SAVE_EVERY  = 100


# >>>> Network <<<<

inputs = tf.placeholder(tf.float32, shape=[None, 784])

labels = tf.placeholder(tf.float32, shape=[None, 10])

h1     = tf.layers.dense(inputs, 256, activation=tf.nn.relu, use_bias=True)

logits = tf.layers.dense(h1, 10, use_bias=True)

predictions = tf.nn.softmax(logits)

prediction_ids = tf.argmax(predictions, axis=1)

# >>>> Loss & Optimisation <<<<

loss = tf.nn.softmax_cross_entropy_with_logits_v2(labels=labels, logits=logits)

opt  = tf.train.AdamOptimizer().minimize(loss)

# >>>> Utilities <<<<

init  = tf.global_variables_initializer()

saver = tf.train.Saver()


with tf.Session() as sess:

    sess.run(init)

    # >>>> Training - run on remote, comment out locally <<<<

    for i in range(TRAIN_STEPS):

        print("Train step ".format(i), end="\r")

        batch_data, batch_labels = mnist.train.next_batch(batch_size=128)

        feed_dict = 
            inputs: batch_data,
            labels: batch_labels
        

        l, _ = sess.run([loss, opt], feed_dict=feed_dict)

        if i % SAVE_EVERY == 0:

            saver.save(sess, "saved_model/network_weights.ckpt")


    # >>>> Using the network - run locally to use the network <<<

    saver.restore(sess, "saved_model/network_weights.ckpt")

    test_data, test_labels = mnist.test.images, mnist.test.labels

    feed_dict = 
        inputs: test_data,
        labels: test_labels
    

    preds = sess.run(prediction_ids, feed_dict=feed_dict)

    print(preds)

因此，一旦您在网络中定义了保护程序，您就可以使用它将权重保存到指定目录 - 在本例中是在目录“saved_models”中，您需要在运行此特定代码之前创建该目录.

恢复模型非常简单，只需调用saver.restore() 然后将会话和权重存储路径传递给它。因此，您可以在远程机器上运行此代码，将“saved_models”目录下载到本地机器，然后运行此代码并注释掉训练部分以实际使用模型。

【讨论】：

哦，我明白了，我想我没有意识到在恢复图形时需要重建模型并重新定义变量。谢谢！