mobilenet_v2 tensorlayer 实现

mobilenet_v2 使用 tensorlayer 实现
训练数据使用 cifar10

完整代码 github 地址
https://github.com/One-sixth/mobilenet_v2-tensorlayer
启动训练
python3 cifar10_use_mobilenet_v2.py

mobilenet_v2 定义
like_mobilenet_v2.py

import tensorflow as tf
import tensorlayer as tl
import numpy as np

def get_model(x, k=None, is_train=True, reuse=False):
    act = tf.nn.leaky_relu
    with tf.variable_scope('mobilenet_v2', reuse=reuse) as vs:
        def bottleneck(x, n_filter, stride=1, t=6, name='bt'):
            # tl.layers.SeparableConv2d(x, n_filter, (3, 3), stride, act, 'same')
            net = tl.layers.Conv2d(x, n_filter * t, (1, 1), b_init=None, name='%s_expand' % name)
            net = tl.layers.BatchNormLayer(net, act=act, is_train=is_train, name='%s_bn1' % name)
            net = tl.layers.DepthwiseConv2d(net, (3, 3), (stride, stride), b_init=None, name='%s_dwise1' % name)
            net = tl.layers.BatchNormLayer(net, act=act, is_train=is_train, name='%s_bn2' % name)
            net = tl.layers.Conv2d(net, n_filter, (1, 1), b_init=None, name='%s_project' % name)
            net = tl.layers.BatchNormLayer(net, act=tf.identity, is_train=is_train, name='%s_bn3' % name)
            if stride == 1:
                last_n_filter = x.outputs.get_shape()[-1]
                if n_filter > last_n_filter:
                    shortcut = tl.layers.PadLayer(x, [[0,0],[0,0],[0,0],[0, n_filter-last_n_filter]])
                elif n_filter < last_n_filter:
                    shortcut = tl.layers.Conv2d(x, n_filter, 1, 1, name='%s_shortcut' % name)
                else:
                    shortcut = x
                net = tl.layers.ElementwiseLayer([net, shortcut], tf.add)
            return net

        # 224 x 224 x 3
        net = tl.layers.InputLayer(x)
        net = tl.layers.Conv2d(net, 32, (3, 3), (2, 2), b_init=None, name='conv1')
        net = tl.layers.BatchNormLayer(net, act=act, is_train=is_train, name='bn1')
        for i in range(1):
            stride = 1
            net = bottleneck(net, 16, stride, 1, 'bt0_%d' % i)
        for i in range(2):
            if i == 0:
                stride = 2
            else:
                stride = 1
            net = bottleneck(net, 24, stride, 6, 'bt1_%d' % i)
        for i in range(3):
            if i == 0:
                stride = 2
            else:
                stride = 1
            net = bottleneck(net, 32, stride, 6, 'bt2_%d' % i)
        for i in range(4):
            if i == 0:
                stride = 2
            else:
                stride = 1
            net = bottleneck(net, 64, stride, 6, 'bt3_%d' % i)
        for i in range(3):
            stride = 1
            net = bottleneck(net, 96, stride, 6, 'bt4_%d' % i)
        for i in range(3):
            if i == 0:
                stride = 2
            else:
                stride = 1
            net = bottleneck(net, 160, stride, 6, 'bt5_%d' % i)
        for i in range(1):
            stride = 1
            net = bottleneck(net, 320, stride, 6, 'bt6_%d' % i)

        if k != None:
            net = tl.layers.Conv2d(net, 1280, 1, 1, b_init=None, name='conv2')
            net = tl.layers.BatchNormLayer(net, act=act, is_train=is_train, name='bn2')
            net = tl.layers.MeanPool2d(net, net.outputs.get_shape()[1:3], 1, 'VALID', 'meanpool1')
            net = tl.layers.Conv2d(net, k, 1, 1, name='output')
            net = tl.layers.FlattenLayer(net)

        ws = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, vs.name + '.*kernel.*')
        l2_loss = tf.add_n([tf.nn.l2_loss(w) for w in ws])

    return net, l2_loss

if __name__ == '__main__':
    x = tf.placeholder(tf.float32, (None, 224, 224, 3))
    net, l2_loss = get_model(x, 1280, True, False)
    writer = tf.summary.FileWriter('log/', graph=tf.get_default_graph())
    writer.close()
    # net = tl.layers.SeparableConv2d(tl.layers.InputLayer(x), 100, padding='SAME')

使用 cifar10 训练 mobilenet_v2

cifar10_use_mobilenet_v2.py

import tensorflow as tf
import tensorlayer as tl
import numpy as np
from progressbar import progressbar
import like_mobilenet_v2

x_train, y_train, x_test, y_test = tl.files.load_cifar10_dataset(shape=(-1, 32, 32, 3))


img = tf.placeholder(tf.float32, (None, 32, 32, 3))
label = tf.placeholder(tf.int64, (None,))
# img_scale = tf.image.resize_images(img, [224, 224])

# net, l2_loss = like_mobilenet_v2.get_model(img_scale, 10, True, False)
net, l2_loss = like_mobilenet_v2.get_model(img, 10, True, False)

cost = tf.losses.sparse_softmax_cross_entropy(label, net.outputs)

correct_prediction = tf.equal(tf.argmax(net.outputs, 1), label)
acc_op = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

train_op = tf.train.AdamOptimizer(0.001).minimize(cost, var_list=net.all_params)

sess = tf.Session()
tl.layers.initialize_global_variables(sess)

epoch = 200
batch_size = 500
train_batch_count = int(np.ceil(len(x_train) / batch_size))
test_batch_count = int(np.ceil(len(x_test) / batch_size))

for e in range(epoch):
    loss = 0
    for b in progressbar(range(train_batch_count)):
        feed_dict = {img:x_train[b*batch_size : (b+1)*batch_size], label:y_train[b*batch_size : (b+1)*batch_size]}
        los, _ = sess.run([cost, train_op], feed_dict)
        loss += los
    print('train loss', loss / train_batch_count)
    if e % 5 == 1:
        loss = 0
        acc = 0
        for b in progressbar(range(test_batch_count)):
            feed_dict = {img: x_test[b * batch_size: (b + 1) * batch_size],
                         label: y_test[b * batch_size: (b + 1) * batch_size]}
            los, ac = sess.run([cost, acc_op], feed_dict)
            loss += los
            acc += ac
        print('test loss', loss / test_batch_count, 'test acc', acc / test_batch_count)