下面我有一些代码用于部署训练有素的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
答案 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
调用获得)输入到原始图形。