MXNet自定义层

1.不含模型参数的自定义层

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

class CenteredLayer(nn.Block):
    # 自定义层
    def __init__(self, **kwargs):
        super(CenteredLayer, self).__init__(**kwargs)
        
    def forward(self, x):
        return x-x.mean()

# 这里构造的是没有自定义参数的层，所以都不需要调用初始化函数initialze（）
# 其实就是forward函数中没有用到参数，所以不会报错
layer = CenteredLayer()
layer(nd.array([1,2,3,4,5]))

# 自定义层也可以用于构造更复杂的模型
net = nn.Sequential()
net.add(nn.Dense(128),
       CenteredLayer())
net.initialize()
y = net(nd.random.uniform(shape=(4, 8)))
y.mean().asscalar

2.含模型参数的自定义层

这里使用Block模型里面自带的ParameterDict类型的成员变量params，params中的get方法可以创建Paramter实例。

class MyDense(nn.Block):
    # units为该层的输出个数，in_units为该层的输入个数
    def __init__(self, units, in_units, **kwargs):
        super(MyDense, self).__init__(**kwargs)
        '''用到成员变量params，通过get方法定义模型参数'''
        self.weight = self.params.get('weight', shape=(in_units, units))
        self.bias = self.params.get('bias', shape=(units,))
        
    def forward(self, x):
        '''这里forward函数使用了自定义的模型参数'''
        linear = nd.dot(x, self.weight.data()) + self.bias.data()
        return nd.relu(linear)

# 自立由于使用到了参数，所以需要初始化
dense = MyDense(units=3, in_units=5)
dense.initialize()
dense(nd.random.uniform(shape=(2,5)))

# 也可以使用自定义层构造模型
net = nn.Sequential()
net.add(MyDense(8, in_units=64),
       MyDense(1, in_units=8))
net.initialize()
net(nd.random.uniform(shape=(2, 64)))