AlexNet_加载ImageNet上的预训练模型_tensorflow版本_imagenet预训练的alexnet.h5下载

在 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
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139

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))
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

加载模型的示例代码如下：

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
1
2
3
4
5
6
7
8

本文主要介绍了一种在tensorflow上加载其它格式（非ckpt格式）的Alexnet预训练模型的方法。在预训练模型的基础上，可以根据自己的任务需要修改loss函数、全连接层设置等进行fine tune。