TensorFlow:在图形定义中存储扭曲

时间:2016-07-13 15:00:42

标签: python machine-learning tensorflow

下面我有一些代码用于部署训练有素的TensorFlow模型。它基本上只是从.pb文件加载模型,获取模型的第一层和最后一层并评估图像。这很好用,但我想部署许多具有不同图像尺寸和扭曲的模型,如模糊,美白和旋转。

我的问题是:图像失真序列可以存储在.pb文件中吗?如果是这样,怎么样?

目标是最小化部署脚本中的代码量。

import base64
import math
import os
import tensorflow as tf

def get_graph():
    if os.path.isfile('./graph.pb'):
        graph_def = tf.GraphDef()
        with open('./graph.pb', 'rb') as graph_file:
            graph_def.ParseFromString(graph_file.read())
    else:
        raise Exception('Graph file \'./graph.pb\' does not exist')
    return graph_def

def init(event):
    graph_def = get_graph()

    with tf.Session() as session:
        session.graph.as_default()
        tf.import_graph_def(graph_def, name = '')

        stringified = base64.b64decode(event['image'].split(',')[1])
        decoded = tf.image.decode_jpeg(stringified, channels = 3)
        decoded.set_shape([event['height'], event['width'], 3])
        image = tf.cast(decoded, tf.float32)

        evaluation = image.eval(sessions = sess)

        input_tensor = sess.graph.get_tensor_by_name('input_placeholder:0')
        output_tensor = sess.graph.get_tensor_by_name('softmax_linear/softmax_linear:0')

        feed_dict = { input_tensor: evaluation }
        result = sess.run([output_tensor], feed_dict = feed_dict)
        return result

1 个答案:

答案 0 :(得分:0)

答案取决于原始graph.pb的生成方式。

如果您修改了生成原始graph.pb的脚本,则只需将重新整形,扭曲等操作添加到该脚本并重新生成graph.pb。请注意,您必须删除旧的input_placeholder操作,并将预处理操作的输出连接到input_placeholder正在馈送的操作的输入。然后,您的新占位符将仅以stringified作为输入。

如果无法修改生成原始graph.pb的脚本,则可以通过将预处理子图保存到自己的.pb来减少部署脚本中的代码量。类似的东西:

raw_input = tf.placeholder(tf.string)
decoded = tf.image.decode_jpeg(raw_input, channels = 3)
decoded.set_shape([event['height'], event['width'], 3])
image = tf.cast(decoded, tf.float32)
with open('preprocess.pb', 'w') as f:
  f.write(tf.get_default_graph().as_graph_def())

然后,您可以将原始输入提供给预处理子图,并将输出(仍然通过eval调用获得)输入到原始图形。