Skip to main content
 首页 » 编程设计

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

2025年01月19日29Free-Thinker

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

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

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

我试过三种情况:

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

但所有情况都显示错误:

  1. ...必须来自与...相同的图表
  2. TypeError:无法 pickle _thread.RLock 对象
  3. TypeError:提要的值不能是 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 有什么想法吗??

请您参考如下方法:

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

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

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

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