TensorFlow 使用预训练模型 ResNet-50
模型定义、训练及预训练参数导入
首先，我们来梳理一下要使用预训练模型来微调神经网络需要做的事情：1.定义神经网络结构；2.导入预训练模型参数；3.读取数据进行训练；4.使用 Tensorboard 可视化训练过程（此处略，留到以后单独讲）。清楚了以上步骤之后，我们来看如下全部代码：
# -*- coding: utf-8 -*-
"""
Created on Tue May 8 13:58:54 2018
@author: shirhe-lyh
"""
import numpy as np
import os
import tensorflow as tf
from tensorflow.contrib.slim import nets
slim = tf.contrib.slim
def get_next_batch(batch_size=64, ...):
"""Get a batch set of training data.
Args:
batch_size: An integer representing the batch size.
...: Additional arguments.
Returns:
images: A 4-D numpy array with shape [batch_size, height, width,
num_channels] representing a batch of images.
labels: A 1-D numpy array with shape [batch_size] representing
the groundtruth labels of the corresponding images.
"""
... # Get images and the corresponding groundtruth labels.
return images, labels
if __name__ == '__main__':
# Specify which gpu to be used
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
batch_size = 64
num_classes = 5
num_steps = 10000
resnet_model_path = '···/resnet_v1_50.ckpt' # Path to the pretrained model
model_save_path = '···/model' # Path to the model.ckpt-(num_steps) will be saved
... # Any other constant variables
inputs = tf.placeholder(tf.float32, shape=[None, 224, 224, 3], name='inputs')
labels = tf.placeholder(tf.int32, shape=[None], name='labels')
is_training = tf.placeholder(tf.bool, name='is_training')
with slim.arg_scope(nets.resnet_v1.resnet_arg_scope()):
net, endpoints = nets.resnet_v1.resnet_v1_50(inputs, num_classes=None,
is_training=is_training)
with tf.variable_scope('Logits'):
net = tf.squeeze(net, axis=[1, 2])
net = slim.dropout(net, keep_prob=0.5, scope='scope')
logits = slim.fully_connected(net, num_outputs=num_classes,
activation_fn=None, scope='fc')
checkpoint_exclude_scopes = 'Logits'
exclusions = None
if checkpoint_exclude_scopes:
exclusions = [
scope.strip() for scope in checkpoint_exclude_scopes.split(',')]
variables_to_restore = []
for var in slim.get_model_variables():
excluded = False
for exclusion in exclusions:
if var.op.name.startswith(exclusion):
excluded = True
if not excluded:
variables_to_restore.append(var)
losses = tf.nn.sparse_softmax_cross_entropy_with_logits(
labels=labels, logits=logits)
loss = tf.reduce_mean(losses)
logits = tf.nn.softmax(logits)
classes = tf.argmax(logits, axis=1, name='classes')
accuracy = tf.reduce_mean(tf.cast(
tf.equal(tf.cast(classes, dtype=tf.int32), labels), dtype=tf.float32))
optimizer = tf.train.AdamOptimizer(learning_rate=0.0001)
train_step = optimizer.minimize(loss)
init = tf.global_variables_initializer()
saver_restore = tf.train.Saver(var_list=variables_to_restore)
saver = tf.train.Saver(tf.global_variables())
config = tf.ConfigProto(allow_soft_placement = True)
config.gpu_options.per_process_gpu_memory_fraction = 0.95
with tf.Session(config=config) as sess:
sess.run(init)
# Load the pretrained checkpoint file xxx.ckpt
saver_restore.restore(sess, resnet_model_path)
for i in range(num_steps):
images, groundtruth_lists = get_next_batch(batch_size, ...)
train_dict = {inputs: images,
labels: groundtruth_lists,
is_training: True}
sess.run(train_step, feed_dict=train_dict)
loss_, acc_ = sess.run([loss, accuracy], feed_dict=train_dict)
train_text = 'Step: {}, Loss: {:.4f}, Accuracy: {:.4f}'.format(
i+1, loss_, acc_)
print(train_text)
if (i+1) % 1000 == 0:
saver.save(sess, model_save_path, global_step=i+1)
print('save mode to {}'.format(model_save_path)
main 函数的第一行 os.environ["CUDA_VISIBLE_DEVICES"] = "0" 指定系统（服务器或电脑等）中哪些 GPU 是对 TensorFlow 可见的，这里说 0 号 GPU 是可见的。如果你的系统只有一个 GPU，那么这行是多余的（请注释掉）；如果你的系统有多个 GPU，那么你可以通过 , 分隔的方式指定同时使用多个 GPU，如 os.environ["CUDA_VISIBLE_DEVICES"] = "0,1" 表示使用 0 号和 1 号 GPU。接下来，定义了几个变量，顾名思义，分别指定了批量大小、要分类的类别数目、迭代的总次数、预训练模型 ResNet-50 的存储路径和训练后模型要保存到的路径（这里最后的 model 不是指文件夹，而是说最后保存的模型要命名为 model-xxx.ckpt，xxx 表示模型保存时对应的训练次数。如果你想保存为其它格式的名字，请任取）。
再接着往下看，定义了 3 个占位符，分别是：一个批量的图像数据入口、对应的类标号入口和是否是将模型用于训练还是推断的界定布尔值。再接着便到了模型定义阶段。第一个 with 语句，直接使用 slim 模块封装好的 ResNet-50 模型，注意要将最后的输出层禁用，即要将参数 num_classes 设置为 None。第二个 with 语句，是自定义的输出层，指定模型最后输出 num_classes 个值。这里，为了与前一个 with 语句相区别，以便以后导入预训练模型参数，一定要指定变量的命名空间 tf.variable_scope()，其中的名字可自取，只需要保证不与 ResNet-50 定义中的变量命名空间重名即可。这个 with 语句有三句话，第一句是去掉多余的维度：假设输入到 ResNet-50 的数据形状为 [None, 224, 224, 3]，则上一个 with 语句的输出形状为 [None, 1, 1, 2048]，这句话的作用就是将数据的形状变为 [None, 2048]，去掉其中形状为 1 的第 1，2 个索引维度；第二句话 net = slim.dropout() 使用了 CNN 常用的正则化手段 dropout；第三句话使用一个全连接层指定输出大小。
到此，模型定义就完成了。接下来，我们导入预训练参数。我们前面定义的神经网络包含两个阶段：1.使用 ResNet-50 来提取图像特征阶段；2.使用自定义的输出层来得到分类结果，它们分别对应前面的两个 with 语句。我们要导入的 ResNet-50 预训练参数只对应第一个阶段的模型参数，因此导入时就需要把第二阶段的参数排除掉。所以，接下来的一串代码
checkpoint_exclude_scopes = 'Logits'
exclusions = None
if checkpoint_exclude_scopes:
exclusions = [
scope.strip() for scope in checkpoint_exclude_scopes.split(',')]
variables_to_restore = []
for var in slim.get_model_variables():
excluded = False
for exclusion in exclusions:
if var.op.name.startswith(exclusion):
excluded = True
if not excluded:
variables_to_restore.append(var)
的作用就是将所有需要恢复的变量取出来，所有预训练模型中没有的变量排除掉。代码首先指定要排除的变量的命名空间为 Logits，对应上面的第二个 with 语句下的所有变量。然后通过 slim.get_model_variables 函数得到所有的模型变量，对这些变量进行遍历，如果变量不在 Logits 这个命名空间中就记录到要从预训练模型恢复的变量列表 variables_to_restore 中，否则就排除掉。
当我们记录下所有要恢复的变量之后，就可以做其它任何事情了，如接下来的代码定义了损失函数 loss、准确率 accuracy、优化器 optimizer 等。然后，定义了两个 tf.train.Saver 类的对象：
saver_restore = tf.train.Saver(var_list=variables_to_restore)
saver = tf.train.Saver(tf.global_variables())
其中，第一个对象 saver_restore 用于恢复预训练模型中的参数，var_list 指定要从预训练模型中恢复的变量列表；第二个对象 saver 用来保存我们微调过后的训练模型，因此要保留所有的模型变量，从而 var_list=tf.global_variables()。接下来的 config 语句是说只使用每个 GPU 的 95% 的显存（config.gpu_options.per_process_gpu_memory_fraction = 0.95
），还允许 TensorFlow 将计算放到 CPU 上（allow_soft_placement=True），如果你不想这么设置直接将有关 config 的所有语句删除即可。
再然后，通过 with tf.Session() as sess 生成一个会话，然后将所有变量作为结点 sess.run(init) 写入 TensorFlow graph。之后，又到了一个重要阶段，即导入预训练模型参数阶段。导入预训练模型参数只需要一句话就可以完成：
saver_restore.restore(sess, resnet_model_path)
当然，我们前面还是做了一些必要的准确： saver_restore = tf.train.Saver(var_list=variables_to_restore)。最后，就是对模型使用自己的数据进行微调以及模型保存等，略过。关于模型保存，请访问 TensorFlow 模型保存与恢复。
作者：公输睚信
链接：https://www.jianshu.com/p/a5ec1dff490d
來源：简书
简书著作权归作者所有，任何形式的转载都请联系作者获得授权并注明出处。