Tensorflow:如何在新图中使用预训练权重? [英] Tensorflow: how to use pretrained weights in new graph?

问题描述

我正在尝试使用 tensorflow 和 python 框架构建一个带有 CNN 的对象检测器.我想训练我的模型首先只进行对象识别(分类)，然后使用预训练模型的几个卷积层训练它来预测边界框.我需要替换完全连接的层，可能还有一些最后的卷积层.因此，出于这个原因，我想知道是否可以从用于训练对象分类器的张量流图中仅导入权重到我将训练进行对象检测的新定义图.所以基本上我想做这样的事情:

I'm trying to build an object detector with CNN using tensorflow with python framework. I would like to train my model to do just object recognition (classification) at first and then using several convolutional layers of the pretarined model train it to predict bounding boxes. I will need to replace fully connected layers and probably some last convolutional layers. So, for this reason, I would like to know if it is possible to import only weights from tensorflow graph that was used to train object classifier to a newly defined graph that I will train to do object detection. So basically I would like to do something like this:

# here I initialize the new graph
conv_1=tf.nn.conv2d(in, weights_from_old_graph)
conv_2=tf.nn.conv2d(conv_1, weights_from_old_graph)
...
conv_n=tf.nn.nnconv2d(conv_n-1,randomly_initialized_weights)
fc_1=tf.matmul(conv_n, randomly_initalized_weights)

推荐答案

使用不带参数的 saver 来保存整个模型.

Use saver with no arguments to save the entire model.

tf.reset_default_graph()
v1 = tf.get_variable("v1", [3], initializer = tf.initializers.random_normal)
v2 = tf.get_variable("v2", [5], initializer = tf.initializers.random_normal)
saver = tf.train.Saver()

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    saver.save(sess, save_path='./test-case.ckpt')

    print(v1.eval())
    print(v2.eval())
saver = None

v1 = [ 2.1882825   1.159807   -0.26564872]
v2 = [0.11437789 0.5742971 ]

然后在要恢复到某些值的模型中，将要恢复的变量名称列表或 {"variable name": variable} 字典传递给 Saver>.

Then in the model you want to restore to certain values, pass a list of variable names you want to restore or a dictionary of {"variable name": variable} to the Saver.

tf.reset_default_graph()
b1 = tf.get_variable("b1", [3], initializer= tf.initializers.random_normal)
b2 = tf.get_variable("b2", [3], initializer= tf.initializers.random_normal)
saver = tf.train.Saver(var_list={'v1': b1})

with tf.Session() as sess:
  saver.restore(sess, "./test-case.ckpt")
  print(b1.eval())
  print(b2.eval())

INFO:tensorflow:Restoring parameters from ./test-case.ckpt
b1 = [ 2.1882825   1.159807   -0.26564872]
b2 = FailedPreconditionError: Attempting to use uninitialized value b2

这篇关于Tensorflow:如何在新图中使用预训练权重?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！