Tensorflow，在另一个 tf.estimator model_fn 中使用一个 tf.estimator 训练模型

Question

有没有办法在另一个模型 B 中使用 tf.estimator 训练的模型 A？

这里的情况， 假设我有一个带有 model_a_fn() 的训练有素的“模型 A”。 'Model A' 获取图像作为输入，并输出一些类似于 MNIST 分类器的向量浮点值。 在 model_b_fn() 中定义了另一个“模型 B”。 它还获取图像作为输入，并且在训练“模型 B”时需要“模型 A”的矢量输出。

所以基本上我想训练需要输入作为图像的“模型 B”和“模型 A”的预测输出。 （不再需要训练'Model A'，只需要在训练'Model B'时得到预测输出）

我试过三种情况：

1. 在 model_b_fn() 中使用估算器对象（'Model A'）
2. 使用 tf.estimator.export_savedmodel() 导出“模型 A”，并创建预测函数。使用 params dict 将其传递给 model_b_fn()。
3. 与 2 相同，但在 model_b_fn() 中恢复“模型 A”

但所有情况都显示错误：

1. ...必须与...来自同一图表
2. TypeError：无法腌制 _thread.RLock 对象
3. TypeError：Feed 的值不能是 tf.Tensor 对象。

这是我使用的代码...仅附加重要部分

train_model_a.py

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()

train_model_b_case_1.py

# 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()

train_model_b_case_2.py

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()

train_model_b_case_3.py

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 吗？

Answer 1

我已经找到了解决这个问题的方法。

如果遇到同样的问题，可以使用此方法。

1. 创建一个运行 tensorflow.contrib.predictor.from_saved_model() 的函数 -> 将其称为“pretrained_predictor()”
2. 在 Model B 的 model_fn() 中，调用上面预定义的 'pretrained_predictor()' 并用 tensorflow.py_func() 包装它

示例案例，简单用例见https://github.com/moono/tf-cnn-mnist/blob/master/4_3_estimator_within_estimator.py。