将训练好的Tensorflow模型保存到另一台机器上进行推理[副本]

Question

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

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

如果有人能告诉我如何正确地保存训练过的模型，以及使用哪些变量来使用保存的模型进行推理，我会非常感激的。

Answer 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”目录下载到本地机器上，然后运行此代码并注释掉训练部分以实际使用模型。