MXNet的新接口Gluon

为什么要开发Gluon的接口

在MXNet中我们可以通过Sybmol模块来定义神经网络，并组通过Module模块提供的一些上层API来简化整个训练过程。那MXNet为什么还要重新开发一套Python的API呢，是否是重复造轮子呢？答案是否定的，Gluon主要是学习了Keras、Pytorch等框架的优点，支持动态图（Imperative）编程，更加灵活且方便调试。而原来MXNet基于Symbol来构建网络的方法是像TF、Caffe2一样静态图的编程方法。同时Gluon也继续了MXNet在静态图上的一些优化，比如节省显存，并行效率高等，运行起来比Pytorch更快。

更加简洁的接口

我们先来看一下用Gluon的接口，如果创建并组训练一个神经网络的，我们以mnist数据集为例：

import mxnet as mx
import mxnet.ndarray as nd
from mxnet import gluon
import mxnet.gluon.nn as nn

数据的读取

首先我们利用Gluon的data模块来读取mnist数据集

def transform(data, label):
    return data.astype('float32') / 255, label.astype('float32')

minist_train_dataset = gluon.data.vision.MNIST(train=True, transform=transform)
minist_test_dataset = gluon.data.vision.MNIST(train=False, transform=transform)

batch_size = 64
train_data = gluon.data.DataLoader(dataset=minist_train_dataset, shuffle=True, batch_size=batch_size)
test_data = gluon.data.DataLoader(dataset=minist_train_dataset, shuffle=False, batch_size=batch_size)

num_examples = len(train_data)
print(num_examples)

训练模型

这里我们使用Gluon来定义一个LeNet

# Step1 定义模型
lenet = nn.Sequential()
with lenet.name_scope():
    lenet.add(nn.Conv2D(channels=20, kernel_size=5, activation='relu'))
    lenet.add(nn.MaxPool2D(pool_size=2, strides=2))
    lenet.add(nn.Conv2D(channels=50, kernel_size=5, activation='relu'))
    lenet.add(nn.MaxPool2D(pool_size=2, strides=2))
    lenet.add(nn.Flatten())
    lenet.add(nn.Dense(128, activation='relu'))
    lenet.add(nn.Dense(10))
# Step2 初始化模型参数
lenet.initialize(ctx=mx.gpu())
# Step3 定义loss
softmax_loss = gluon.loss.SoftmaxCrossEntropyLoss()
# Step4 优化
trainer = gluon.Trainer(lenet.collect_params(), 'sgd', {'learning_rate': 0.5})

def accuracy(output, label):
     return nd.mean(output.argmax(axis=1)==label).asscalar()
def evaluate_accuracy(net, data_iter):
    acc = 0
    for data, label in data_iter:
        data = data.transpose((0,3,1,2))
        data = data.as_in_context(mx.gpu())
        label = label.as_in_context(mx.gpu())
        output = net(data)
        acc += accuracy(output, label)
    return acc / len(data_iter)

import mxnet.autograd as ag
epochs = 5
for e in range(epochs):
    total_loss = 0
    for data, label in train_data:
        data = data.transpose((0,3,1,2))
        data = data.as_in_context(mx.gpu())
        label = label.as_in_context(mx.gpu())
        with ag.record():
            output = lenet(data)
            loss = softmax_loss(output, label)
        loss.backward()
        trainer.step(batch_size)
        total_loss += nd.mean(loss).asscalar()
    print("Epoch %d, test accuracy: %f, average loss: %f" % (e, evaluate_accuracy(lenet, test_data), total_loss/num_examples))

背后的英雄 nn.Block

我们前面使用了nn.Sequential来定义一个模型，但是没有仔细介绍它，它其实是nn.Block的一个简单的形式。而nn.Block是一个一般化的部件。整个神经网络可以是一个nn.Block，单个层也是一个nn.Block。我们可以（近似）无限地嵌套nn.Block来构建新的nn.Block。nn.Block主要提供3个方向的功能：

1. 存储参数
2. 描述forward如何执行
3. 自动求导

所以nn.Sequential是一个nn.Block的容器，它通过add来添加nn.Block。它自动生成forward()函数。一个简单实现看起来如下：

class Sequential(nn.Block):
    def __init__(self, **kwargs):
        super(Sequential, self).__init__(**kwargs)
    def add(self, block):
        self._children.append(block)
    def forward(self, x):
        for block in self._children:
            x = block(x)
        return x

知道了nn.Block里的魔法后，我们就可以自定我们自己的nn.Block了，来实现不同的深度学习应用可能遇到的一些新的层。

在nn.Block中参数都是以一种Parameter的对象，通过这个对象的data()和grad()来访问对应的数据和梯度。

my_param = gluon.Parameter('my_params', shape=(3,3))
my_param.initialize()
(my_param.data(), my_param.grad())

每个nn.Block里都有一个类型为ParameterDict类型的成员变量params来保存所有这个层的参数。它其际上是一个名称到参数映射的字典。

pd = gluon.ParameterDict(prefix='custom_layer_name')
pd.get('custom_layer_param1', shape=(3,3))
pd

自义我们自己的全连接层

当我们要实现的功能在Gluon.nn模块中找不到对应的实现时，我们可以创建自己的层，它实际也就是一个nn.Block对象。要自定义一个nn.Block以，只需要继承nn.Block，如果该层需要参数，则在初始化函数中做好对应参数的初始化（实际只是分配的形状），然后再实现一个forward()函数来描述计算过程。

class MyDense(nn.Block):
    def __init__(self, units, in_units, **kwargs):
        super(MyDense, self).__init__(**kwargs)
        with self.name_scope():
            self.weight = self.params.get(
                'weight', shape=(in_units, units))
            self.bias = self.params.get('bias', shape=(units,))

    def forward(self, x):
        linear = nd.dot(x, self.weight.data()) + self.bias.data()
        return nd.relu(linear)

审视模型的参数

我们将从下面三个方面来详细讲解如何操作gluon定义的模型的参数。

1. 初始化
2. 读取参数
3. 参数的保存与加载

从上面我们们在mnist训练一个模型的步骤中可以看出，当我们定义好模型后，第一步就是需要调用initialize()对模型进行参数初始化。

def get_net():
    net = nn.Sequential()
    with net.name_scope():
        net.add(nn.Dense(4, activation='relu'))
        net.add(nn.Dense(2))
    return net
net = get_net()
net.initialize()

我们一直使用默认的initialize来初始化权重。实际上我们可以指定其他初始化的方法，mxnet.initializer模块中提供了大量的初始化权重的方法。比如非常流行的Xavier方法。

#net.initialize(init=mx.init.Xavier())
x = nd.random.normal(shape=(3,4))
net(x)

我们可以weight和bias来访问Dense的参数，它们是Parameter对象。

w = net[0].weight
b = net[0].bias
print('weight:', w.data())
print('weight gradient', w.grad())
print('bias:', b.data())
print('bias gradient', b.grad())

我们也可以通过collect_params来访问Block里面所有的参数（这个会包括所有的子Block）。它会返回一个名字到对应Parameter的dict。既可以用正常[]来访问参数，也可以用get()，它不需要填写名字的前缀。

params = net.collect_params()
print(params)
print(params['sequential18_dense0_weight'].data())
print(params.get('dense0_bias').data()) #不需要名字的前缀

延后的初始化

如果我们仔细分析过整个网络的初始化，我们会有发现，当我们没有给网络真正的输入数据时，网络中的很多参数是无法确认形状的。

net = get_net()
net.collect_params()

net.initialize()
net.collect_params()

我们注意到参数中的weight的形状的第二维都是0, 也就是说还没有确认。那我们可以肯定的是这些参数肯定是还没有分配内存的。

net(x)
net.collect_params()

当我们给这个网络一个输入数据后，网络中的数据参数的形状就固定下来了。而这个时候，如果我们给这个网络一个不同shape的输入数据，那运行中就会出现崩溃的问题。

模型参数的保存与加载

gluon.Sequential模块提供了save和load接口来方便我们对一个网络的参数进行保存与加载。

filename = "mynet.params"
net.save_params(filename)
net2 = get_net()
net2.load_params(filename, mx.cpu())

Hybridize

从上面我们使用gluon来训练mnist，可以看出，我们使用的是一种命令式的编程风格。大部分的深度学习框架只在命令式与符号式间二选一。那我们能不能拿到两种泛式全部的优点呢，事实上这一点可以做到。在MXNet的GluonAPI中，我们可以使用HybridBlock或者HybridSequential来构建网络。默认他们跟Block和Sequential一样是命令式的。但当我们调用.hybridize()后，系统会转撚成符号式来执行。

def get_net():
    net = nn.HybridSequential()
    with net.name_scope():
        net.add(
            nn.Dense(256, activation="relu"),
            nn.Dense(128, activation="relu"),
            nn.Dense(2)
        )
    net.initialize()
    return net

x = nd.random.normal(shape=(1, 512))
net = get_net()
net(x)

net.hybridize()
net(x)

注意到只有继承自HybridBlock的层才会被优化。HybridSequential和Gluon提供的层都是它的子类。如果一个层只是继承自Block，那么我们将跳过优化。我们可以将符号化的模型的定义保存下来，在其他语言API中加载。

x = mx.sym.var('data')
y = net(x)
print(y.tojson())

可以看出，对于HybridBlock的模块，既可以把NDArray作为输入，也可以把Symbol对象作为输入。当以Symbol作为输出时，它的结果就是一个Symbol对象。