[英]Saving a custom tf.estimator trained model for tensorflow serving
[英]Tensorflow, use a tf.estimator trained model within another tf.estimator model_fn
有没有办法在另一个模型B中使用tf.estimator训练的模型A?
这是一种情况,假设我有一个受过训练的“ model A”,带有model_a_fn()。 “模型A”获取图像作为输入,并输出一些类似于MNIST分类器的矢量浮点值。 在model_b_fn()中定义了另一个“模型B”。 它也获取图像作为输入,并在训练“模型B”时需要“模型A”的矢量输出。
因此,基本上我想训练需要输入作为“模型A”的图像和预测输出的“模型B”。 (不再需要训练“模型A”,只需要在训练“模型B”时获得预测输出即可)
我尝试了三种情况:
但是所有情况都显示错误:
这是我使用的代码...仅附加重要部分
def model_a_fn(features, labels, mode, params):
# ...
# ...
# ...
return
def main():
# model checkpoint location
model_a_dir = './model_a'
# create estimator for Model A
model_a = tf.estimator.Estimator(model_fn=model_a_fn, model_dir=model_a_dir)
# train Model A
model_a.train(input_fn=lambda : input_fn_a)
# ...
# ...
# ...
# export model a
model_a.export_savedmodel(model_a_dir, serving_input_receiver_fn=serving_input_receiver_fn)
# exported to ./model_a/123456789
return
if __name__ == '__main__':
main()
# follows model_a's input format
def bypass_input_fn(x):
features = {
'x': x,
}
return features
def model_b_fn(features, labels, mode, params):
# parse input
inputs = tf.reshape(features['x'], shape=[-1, 28, 28, 1])
# get Model A's response
model_a = params['model_a']
predictions = model_a.predict(
input_fn=lambda: bypass_input_fn(inputs)
)
for results in predictions:
# Error occurs!!!
model_a_output = results['class_id']
# build Model B
layer1 = tf.layers.conv2d(inputs, 32, 5, same, activation=tf.nn.relu)
layer1 = tf.layers.max_pooling2d(layer1, pool_size=[2, 2], strides=2)
# ...
# some layers added...
# ...
flatten = tf.layers.flatten(prev_layer)
layern = tf.layers.dense(10)
# let say layern's output shape and model_a_output's output shape is same
add_layer = tf.add(flatten, model_a_output)
# ...
# do more... stuff
# ...
return
def main():
# load pretrained model A
model_a_dir = './model_a'
model_a = tf.estimator.Estimator(model_fn=model_a_fn, model_dir=model_a_dir)
# model checkpoint location
model_b_dir = './model_b/'
# create estimator for Model A
model_b = tf.estimator.Estimator(
model_fn=model_b_fn,
model_dir=model_b_dir,
params={
'model_a': model_a,
}
)
# train Model B
model_b.train(input_fn=lambda : input_fn_b)
return
if __name__ == '__main__':
main()
def model_b_fn(features, labels, mode, params):
# parse input
inputs = tf.reshape(features['x'], shape=[-1, 28, 28, 1])
# get Model A's response
model_a_predict_fn = params['model_a_predict_fn']
model_a_prediction = model_a_predict_fn(
{
'x': inputs
}
)
model_a_output = model_a_prediction['output']
# build Model B
layer1 = tf.layers.conv2d(inputs, 32, 5, same, activation=tf.nn.relu)
layer1 = tf.layers.max_pooling2d(layer1, pool_size=[2, 2], strides=2)
# ...
# some layers added...
# ...
flatten = tf.layers.flatten(prev_layer)
layern = tf.layers.dense(10)
# let say layern's output shape and model_a_output's output shape is same
add_layer = tf.add(flatten, model_a_output)
# ...
# do more... stuff
# ...
return
def main():
# load pretrained model A
model_a_dir = './model_a/123456789'
model_a_predict_fn = tf.contrib.predictor.from_saved_model(export_dir=model_a_dir)
# model checkpoint location
model_b_dir = './model_b/'
# create estimator for Model A
# Error occurs!!!
model_b = tf.estimator.Estimator(
model_fn=model_b_fn,
model_dir=model_b_dir,
params={
'model_a_predict_fn': model_a_predict_fn,
}
)
# train Model B
model_b.train(input_fn=lambda : input_fn_b)
return
if __name__ == '__main__':
main()
def model_b_fn(features, labels, mode, params):
# parse input
inputs = tf.reshape(features['x'], shape=[-1, 28, 28, 1])
# get Model A's response
model_a_predict_fn = tf.contrib.predictor.from_saved_model(export_dir=params['model_a_dir'])
# Error occurs!!!
model_a_prediction = model_a_predict_fn(
{
'x': inputs
}
)
model_a_output = model_a_prediction['output']
# build Model B
layer1 = tf.layers.conv2d(inputs, 32, 5, same, activation=tf.nn.relu)
layer1 = tf.layers.max_pooling2d(layer1, pool_size=[2, 2], strides=2)
# ...
# some layers added...
# ...
flatten = tf.layers.flatten(prev_layer)
layern = tf.layers.dense(10)
# let say layern's output shape and model_a_output's output shape is same
add_layer = tf.add(flatten, model_a_output)
# ...
# do more... stuff
# ...
return
def main():
# load pretrained model A
model_a_dir = './model_a/123456789'
# model checkpoint location
model_b_dir = './model_b/'
# create estimator for Model A
# Error occurs!!!
model_b = tf.estimator.Estimator(
model_fn=model_b_fn,
model_dir=model_b_dir,
params={
'model_a_dir': model_a_dir,
}
)
# train Model B
model_b.train(input_fn=lambda : input_fn_b)
return
if __name__ == '__main__':
main()
因此,关于在另一个tf.estimator中使用训练有素的自定义tf.estimator的任何想法请吗?
我已经找到解决这个问题的一种方法。
如果遇到同样的问题,可以使用这种方法。
例如,请参阅https://github.com/moono/tf-cnn-mnist/blob/master/4_3_estimator_within_estimator.py了解简单用例。
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.