AlexNet_加载ImageNet上的预训练模型_tensorflow版本

文章目录

  • 1. 下载 Alexnet的预训练模型参数
  • 2. 定义Alexnet网络模型
  • 3. 加载模型参数

在 tensorflow的GitHub仓库中没有直接给出Alexnet在ImageNet上的预训练模型供tensorflow调用。本文主要通过参考以下资料,实现了使用tensorflow加载在ImageNet上预训练好的其它格式保存的Alexnet模型:

  • https://github.com/stephen-v/tensorflow_alexnet_classify
  • http://www.cnblogs.com/vipyoumay/p/7686230.html

关键步骤分为以下几步:

  1. 下载 Alexnet的预训练模型参数
  2. 定义Alexnet网络模型
  3. 加载模型参数

1. 下载 Alexnet的预训练模型参数

下载地址为:http://www.cs.toronto.edu/~guerzhoy/tf_alexnet/ ,所有的参数数据存放在bvlc_alexnet.npy文件中。

2. 定义Alexnet网络模型

对网络模型定义的时候,关键在于变量的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

3. 加载模型参数

加载模型参数的时候,可以只加载卷积层的参数,也可以加载全部的参数,根据需要自行决定。首先定义加载参数的函数:

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。

你可能感兴趣的:(TensorFlow)