【MXNet】（二十一）：实现VGGNet

论文原文《Very Deep Convolutional Networks for Large-Scale Image Recognition》。

from mxnet import gluon, init, nd
from mxnet.gluon import nn

def vgg_block(num_convs, num_channels):
    blk = nn.Sequential()
    for _ in range(num_convs):
        blk.add(nn.Conv2D(num_channels, kernel_size=3, padding=1, activation='relu'))
    blk.add(nn.MaxPool2D(pool_size=2, strides=2))
    return blk
           

这里实现的是VGG-11，它有5个卷积块，前2块使用单卷积层，而后3块使用双卷积层。第一块的输出通道是64，之后每次对输出通道数翻倍，直到变为512。

conv_arch = ((1, 64), (1, 128), (2, 256), (2, 512), (2, 512))
           

def vgg(conv_arch):
    net = nn.Sequential()
    for (num_convs, num_channels) in conv_arch:
        net.add(vgg_block(num_convs, num_channels))
    net.add(nn.Dense(4096, activation='relu'), nn.Dropout(0.5),
           nn.Dense(4096, activation='relu'), nn.Dropout(0.5),
           nn.Dense(10))
    return net

net = vgg(conv_arch)
           

打印各层的尺寸看一下。

net.initialize()
X = nd.random.uniform(shape=(1, 1, 224, 224))
for blk in net:
    X = blk(X)
    print(blk.name, 'output shape:\t', X.shape)
           

为简单起见数据集还是使用Fashion-MNIST。

lr, num_epochs, batch_size, ctx = 0.05, 5, 128, try_gpu()
net.initialize(force_reinit=True, ctx=ctx, init=init.Xavier())
trainer = gluon.Trainer(net.collect_params(), 'sgd', {'learning_rate': lr})
train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size, resize=224)
train(net, train_iter, test_iter, batch_size, trainer, ctx, num_epochs)
           

其中try_gpu（）函数和train（）函数的实现请参考链接：【MXNet】（二十）：实现AlexNet。

VGG网络比较深，参数量很大，因此需要的GPU显存也是比较多的。我的GPU是NVIDIA Geforce GTX 1050Ti，4G显存，把batch_size改为16才能训练。

慢的感人！！！