关键步骤分为以下几步:
下载地址为:http://www.cs.toronto.edu/~guerzhoy/tf_alexnet/ ,所有的参数数据存放在bvlc_alexnet.npy文件中。
对网络模型定义的时候,关键在于变量的name、shape、dtype等属性的设置,要与bvlc_alexnet.npy文件中保存的一致。个人猜测是因为在加载恢复Alexnet网络模型的时候,代码中定义的变量要与文件中保存的变量一一对应。定义Alexnet代码如下 (可根据需要灵活修改,只要不改变所定义变量的name、shape、dtype等属性):
def alexnet(x, keep_prob, num_classes=1):
# conv1
with tf.variable_scope('conv1') as scope:
kernel = tf.Variable(tf.truncated_normal([11, 11, 3, 96], dtype=tf.float32,
stddev=1e-1), name='weights')
conv = tf.nn.conv2d(x, kernel, [1, 4, 4, 1], padding='SAME')
biases = tf.Variable(tf.constant(0.0, shape=[96], dtype=tf.float32),
trainable=True, name='biases')
bias = tf.nn.bias_add(conv, biases)
conv1 = tf.nn.relu(bias, name='conv1')
# lrn1
with tf.variable_scope('lrn1') as scope:
lrn1 = tf.nn.local_response_normalization(conv1,
alpha=1e-4,
beta=0.75,
depth_radius=2,
bias=2.0)
# pool1
with tf.variable_scope('pool1') as scope:
pool1 = tf.nn.max_pool(lrn1,
ksize=[1, 3, 3, 1],
strides=[1, 2, 2, 1],
padding='VALID')
# conv2
with tf.variable_scope('conv2') as scope:
pool1_groups = tf.split(axis=3, value=pool1, num_or_size_splits=2)
kernel = tf.Variable(tf.truncated_normal([5, 5, 48, 256], dtype=tf.float32,
stddev=1e-1), name='weights')
kernel_groups = tf.split(axis=3, value=kernel, num_or_size_splits=2)
conv_up = tf.nn.conv2d(pool1_groups[0], kernel_groups[0], [1, 1, 1, 1], padding='SAME')
conv_down = tf.nn.conv2d(pool1_groups[1], kernel_groups[1], [1, 1, 1, 1], padding='SAME')
biases = tf.Variable(tf.constant(0.0, shape=[256], dtype=tf.float32),
trainable=True, name='biases')
biases_groups = tf.split(axis=0, value=biases, num_or_size_splits=2)
bias_up = tf.nn.bias_add(conv_up, biases_groups[0])
bias_down = tf.nn.bias_add(conv_down, biases_groups[1])
bias = tf.concat(axis=3, values=[bias_up, bias_down])
conv2 = tf.nn.relu(bias, name='conv2')
# lrn2
with tf.variable_scope('lrn2') as scope:
lrn2 = tf.nn.local_response_normalization(conv2,
alpha=1e-4,
beta=0.75,
depth_radius=2,
bias=2.0)
# pool2
with tf.variable_scope('pool2') as scope:
pool2 = tf.nn.max_pool(lrn2,
ksize=[1, 3, 3, 1],
strides=[1, 2, 2, 1],
padding='VALID')
# conv3
with tf.variable_scope('conv3') as scope:
kernel = tf.Variable(tf.truncated_normal([3, 3, 256, 384],
dtype=tf.float32,
stddev=1e-1), name='weights')
conv = tf.nn.conv2d(pool2, kernel, [1, 1, 1, 1], padding='SAME')
biases = tf.Variable(tf.constant(0.0, shape=[384], dtype=tf.float32),
trainable=True, name='biases')
bias = tf.nn.bias_add(conv, biases)
conv3 = tf.nn.relu(bias, name='conv3')
# conv4
with tf.variable_scope('conv4') as scope:
conv3_groups = tf.split(axis=3, value=conv3, num_or_size_splits=2)
kernel = tf.Variable(tf.truncated_normal([3, 3, 192, 384],
dtype=tf.float32,
stddev=1e-1), name='weights')
kernel_groups = tf.split(axis=3, value=kernel, num_or_size_splits=2)
conv_up = tf.nn.conv2d(conv3_groups[0], kernel_groups[0], [1, 1, 1, 1], padding='SAME')
conv_down = tf.nn.conv2d(conv3_groups[1], kernel_groups[1], [1, 1, 1, 1], padding='SAME')
biases = tf.Variable(tf.constant(0.0, shape=[384], dtype=tf.float32),
trainable=True, name='biases')
biases_groups = tf.split(axis=0, value=biases, num_or_size_splits=2)
bias_up = tf.nn.bias_add(conv_up, biases_groups[0])
bias_down = tf.nn.bias_add(conv_down, biases_groups[1])
bias = tf.concat(axis=3, values=[bias_up, bias_down])
conv4 = tf.nn.relu(bias, name='conv4')
# conv5
with tf.variable_scope('conv5') as scope:
conv4_groups = tf.split(axis=3, value=conv4, num_or_size_splits=2)
kernel = tf.Variable(tf.truncated_normal([3, 3, 192, 256],
dtype=tf.float32,
stddev=1e-1), name='weights')
kernel_groups = tf.split(axis=3, value=kernel, num_or_size_splits=2)
conv_up = tf.nn.conv2d(conv4_groups[0], kernel_groups[0], [1, 1, 1, 1], padding='SAME')
conv_down = tf.nn.conv2d(conv4_groups[1], kernel_groups[1], [1, 1, 1, 1], padding='SAME')
biases = tf.Variable(tf.constant(0.0, shape=[256], dtype=tf.float32),
trainable=True, name='biases')
biases_groups = tf.split(axis=0, value=biases, num_or_size_splits=2)
bias_up = tf.nn.bias_add(conv_up, biases_groups[0])
bias_down = tf.nn.bias_add(conv_down, biases_groups[1])
bias = tf.concat(axis=3, values=[bias_up, bias_down])
conv5 = tf.nn.relu(bias, name='conv5')
# pool5
with tf.variable_scope('pool5') as scope:
pool5 = tf.nn.max_pool(conv5,
ksize=[1, 3, 3, 1],
strides=[1, 2, 2, 1],
padding='VALID', )
# flattened6
with tf.variable_scope('flattened6') as scope:
flattened = tf.reshape(pool5, shape=[-1, 6 * 6 * 256])
# fc6
with tf.variable_scope('fc6') as scope:
weights = tf.Variable(tf.truncated_normal([6 * 6 * 256, 100],
dtype=tf.float32,
stddev=1e-1), name='weights')
biases = tf.Variable(tf.constant(0.0, shape=[100], dtype=tf.float32),
trainable=True, name='biases')
bias = tf.nn.xw_plus_b(flattened, weights, biases)
fc6 = tf.nn.relu(bias)
# dropout6
with tf.variable_scope('dropout6') as scope:
dropout6 = tf.nn.dropout(fc6, keep_prob)
# fc7
with tf.variable_scope('fc7') as scope:
weights = tf.Variable(tf.truncated_normal([100, 1],
dtype=tf.float32,
stddev=1e-1), name='weights')
biases = tf.Variable(tf.constant(0.0, shape=[1], dtype=tf.float32),
trainable=True, name='biases')
bias = tf.nn.xw_plus_b(dropout6, weights, biases)
fc7 = tf.nn.relu(bias)
return fc7
加载模型参数的时候,可以只加载卷积层的参数,也可以加载全部的参数,根据需要自行决定。首先定义加载参数的函数:
def load_initial_weights(session):
"""
:param session: It's need to define tf.Session() first to run the graph of tensorflow.
:return:
"""
# Load the weights into memory.
# WEIGHTS_PATH: 存放 bvlc_alexnet.npy 文件的路径
weights_dict = np.load(os.path.join(WEIGHTS_PATH, 'bvlc_alexnet.npy'), encoding='bytes').item()
# Loop over all layer names stored in the weights dict
for op_name in weights_dict:
# Check if layer should be trained from scratch
# SKIP_LAYER: 指定对某些层的参数不进行恢复,比如只恢复卷积层的参数,就设为: SKIP_LAYER = ['fc6', 'fc7', 'fc8']
# 一定要注意 variable_scope 的相对范围。
if op_name not in SKIP_LAYER:
with tf.variable_scope(op_name, reuse=True):
# Assign weights/biases to their corresponding tf variable
for data in weights_dict[op_name]:
# Biases
if len(data.shape) == 1:
# var = tf.get_variable('biases') 会报错
var = slim.get_unique_variable(op_name + '/biases')
session.run(var.assign(data))
# Weights
else:
# var = tf.get_variable('weights')
var = slim.get_unique_variable(op_name + '/weights')
session.run(var.assign(data))
加载模型的示例代码如下:
sess.run(init)
# print(sess.run(slim.get_unique_variable('conv1/biases')))
load_initial_weights(sess)
# print(sess.run(slim.get_unique_variable('conv1/biases')))
check_init = tf.report_uninitialized_variables()
assert sess.run(check_init).size == 0
本文主要介绍了一种在tensorflow上加载其它格式(非ckpt格式)的Alexnet预训练模型的方法。在预训练模型的基础上,可以根据自己的任务需要修改loss函数、全连接层设置等进行fine tune。