Tensorflow:在具有不同类别数量的新数据集上微调预训练模型

Question

如何在新数据集上微调 tensorflow 中的预训练模型？在 Caffe 中，我可以简单地重命名最后一层并设置一些参数以进行随机初始化。 tensorflow 中是否可能出现类似的情况？

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

然后我尝试像这样启动 session :

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

Answer 1

您可以指定要恢复的变量的名称。

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

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了解详情。

或者，您可以尝试加载整个模型，并在最后一层之前的层中创建一个新的“分支”，并在训练期间在成本函数中使用它。