从save_model.pb加载预训练的模型以使用Custom Estimator进行转移学习

Question

我的任务是从“原始”张量流中移植现有模型。

即）

sess = tf.Session()

sess.run(tf.global_variables_initializer())
train_op = ...
accuracy = ...
for i in range(100):
        print('EPOCH', i)
        _, acc = sess.run([train_op, accuracy], feed_dict={x: inputs, y: labels})

自定义估算器（基于我开发的自定义估算器的通用模板）。

此现有模型利用 saved_model 中的预训练模型作为其初始层进行特征提取，并增加了几层+ softmax进行分类。

总体思路：

def create_estimator(config, hyper_params):
    def _model_fn(features, labels, mode):

         <feed features into pretrained model>
         features_layer = <get final layer from pretrained model>

         logits = tf.train.Dense.apply(NUM_CLASSES,
                              activation=None,use_bias=True).
                              apply(feature_layer)
         predictions = tf.nn.softmax(predictions)
         if(mode == Modes.TRAIN or mode == Modes.EVAL):
             loss = ...
             if(Modes == Modes.TRAIN):
                 train_op = 
                    tf.train.AdamOptimizer(
                    learning_rate=learning_rate).minimize(loss)
         if(mode == Modes.PREDICT):
             ...
         return tf.estimator.EstimatorSpec(
             mode=mode,loss=cost,
             train_op=optimizer,
             training_hooks=[...])

我最初的想法是编写一个从tf.train.SessionRunHook派生的类，因为我之前曾使用它来从检查点进行部分变量还原，并且我在下面包括了一些起始代码...但是在这一点我停留在如何实现自己的目标上。

class RestoreFromSavedModel(tf.train.SessionRunHook):
        def begin():
            graph = tf.get_default_graph()
            with tf.gfile.GFile(os.path.join(
                    MODEL_DIR, IMAGE_CLASSIFIER_FILE), 'rb') as f:
                graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
                _ = tf.import_graph_def(graph_def, name='')
        output_tensor = graph.get_tensor_by_name('prelogits')

任何人都可以就我如何将保存的模型读入图形，将输入数据输入模型并提取特征张量以输入到小模型中的方式提供一些反馈。