Tensorflow:具有不同类别数

时间:2017-01-19 19:26:55

标签: tensorflow

如何在新数据集的tensorflow中微调预训练模型?在Caffe中,我可以简单地重命名最后一层并为随机初始化设置一些参数。张量流中有类似的东西吗?

假设我有一个检查点文件(deeplab_resnet.ckpt)和一些设置计算图的代码,我可以在其中修改最后一层,使得它具有与新数据集具有类相同的输出数。

然后我尝试像这样开始会话:

sess = tf.Session(config=config)
init = tf.initialize_all_variables()

sess.run(init)

trainable = tf.trainable_variables()
saver = tf.train.Saver(var_list=trainable, max_to_keep=40)
saver.restore(sess, 'ckpt_path/deeplab_resnet.ckpt')

然而,这在调用saver.restore函数时给出了一个错误,因为它期望与保存它的图形结构完全相同。 如何只加载'ckpt_path/deeplab_resnet.ckpt'文件中最后一层以外的所有权重? 我也尝试更改Classification图层名称,但没有运气......

我正在使用tensorflow-deeplab-resnet model

1 个答案:

答案 0 :(得分:3)

您可以指定要还原的变量的名称。

因此,您可以获取模型中所有变量的列表,并过滤掉最后一层的变量:

all_vars = tf.all_variables()
var_to_restore = [v for v in all_vars if not v.name.startswith('xxx')]

saver = tf.train.Saver(var_to_restore)

有关详细信息,请参阅documentation

或者,您可以尝试加载整个模型,创建一个新的"分支"超出最后一层之外的层,并在培训期间在成本函数中使用它。