5.4. 自定义层

深度学习成功背后的一个因素是神经网络的灵活性: 我们可以用创造性的方式组合不同的层,从而设计出适用于各种任务的架构。 例如,研究人员发明了专门用于处理图像、文本、序列数据和执行动态规划的层。 有时我们会遇到或要自己发明一个现在在深度学习框架中还不存在的层。 在这些情况下,必须构建自定义层。本节将展示如何构建自定义层。

5.4.1. 不带参数的层

首先,我们构造一个没有任何参数的自定义层。 回忆一下在 5.1节对块的介绍, 这应该看起来很眼熟。 下面的CenteredLayer类要从其输入中减去均值。 要构建它,我们只需继承基础层类并实现前向传播功能。

from mindspore import Parameter, nn, ops


class CenteredLayer(nn.Cell):
    def __init__(self):
        super().__init__()

    def construct(self, X):
        return X - X.mean()
import torch
import torch.nn.functional as F
from torch import nn


class CenteredLayer(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, X):
        return X - X.mean()

让我们向该层提供一些数据,验证它是否能按预期工作。

layer = CenteredLayer()
layer(ops.Tensor([1., 2, 3, 4, 5]))
Tensor(shape=[5], dtype=Float32, value= [-2.00000000e+00, -1.00000000e+00,  0.00000000e+00,  1.00000000e+00,  2.00000000e+00])
layer = CenteredLayer()
layer(torch.FloatTensor([1, 2, 3, 4, 5]))
tensor([-2., -1.,  0.,  1.,  2.])

现在,我们可以将层作为组件合并到更复杂的模型中。

net = nn.SequentialCell([nn.Dense(8, 128), CenteredLayer()])
net = nn.Sequential(nn.Linear(8, 128), CenteredLayer())

作为额外的健全性检查,我们可以在向该网络发送随机数据后,检查均值是否为0。 由于我们处理的是浮点数,因为存储精度的原因,我们仍然可能会看到一个非常小的非零数。

Y = net(ops.rand((4, 8)))
Y.mean()
Tensor(shape=[], dtype=Float32, value= -2.32831e-09)
Y = net(torch.rand(4, 8))
Y.mean()
tensor(-5.5879e-09, grad_fn=<MeanBackward0>)

5.4.2. 带参数的层

以上我们知道了如何定义简单的层,下面我们继续定义具有参数的层, 这些参数可以通过训练进行调整。 我们可以使用内置函数来创建参数,这些函数提供一些基本的管理功能。 比如管理访问、初始化、共享、保存和加载模型参数。 这样做的好处之一是:我们不需要为每个自定义层编写自定义的序列化程序。

现在,让我们实现自定义版本的全连接层。 回想一下,该层需要两个参数,一个用于表示权重,另一个用于表示偏置项。 在此实现中,我们使用修正线性单元作为激活函数。 该层需要输入参数:in_unitsunits,分别表示输入数和输出数。

class MyLinear(nn.Cell):
    def __init__(self, in_units, units):
        super().__init__()
        self.weight = Parameter(ops.randn((in_units, units)))
        self.bias = Parameter(ops.randn((units,)))

    def construct(self, X):
        linear = ops.matmul(X, self.weight) + self.bias
        return ops.relu(linear)
class MyLinear(nn.Module):
    def __init__(self, in_units, units):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(in_units, units))
        self.bias = nn.Parameter(torch.randn(units,))
    def forward(self, X):
        linear = torch.matmul(X, self.weight.data) + self.bias.data
        return F.relu(linear)

接下来,我们实例化MyDense类并访问其模型参数。

linear = MyLinear(5, 3)
linear.weight.value()
Tensor(shape=[5, 3], dtype=Float32, value=
[[-8.37819934e-01, -1.42728090e-01,  6.98715270e-01],
 [ 6.33553147e-01,  4.66425151e-01, -9.23575342e-01],
 [ 7.80555904e-02,  1.27004075e+00, -9.80480433e-01],
 [ 7.85815343e-02,  1.37643766e+00, -4.39103961e-01],
 [ 1.91481805e+00, -3.83728445e-01,  1.56912458e+00]])
linear = MyLinear(5, 3)
linear.weight
Parameter containing:
tensor([[ 0.1146,  0.2678, -2.7968],
        [-0.7714,  0.3853,  0.3882],
        [ 0.3864,  0.2533,  0.2811],
        [ 0.8532,  0.2337,  0.3863],
        [ 1.8150, -0.7465,  0.5738]], requires_grad=True)

我们可以使用自定义层直接执行前向传播计算。

linear(ops.rand((2, 5)))
Tensor(shape=[2, 3], dtype=Float32, value=
[[ 1.47905600e+00,  3.60304499e+00,  0.00000000e+00],
 [ 2.11579418e+00,  3.30277443e+00,  0.00000000e+00]])
linear(torch.rand(2, 5))
tensor([[2.4316, 0.0000, 0.0000],
        [2.4526, 0.0000, 0.0000]])

我们还可以使用自定义层构建模型,就像使用内置的全连接层一样使用自定义层。

net = nn.SequentialCell([MyLinear(64, 8), MyLinear(8, 1)])
net(ops.rand((2, 64)))
Tensor(shape=[2, 1], dtype=Float32, value=
[[ 0.00000000e+00],
 [ 0.00000000e+00]])
net = nn.Sequential(MyLinear(64, 8), MyLinear(8, 1))
net(torch.rand(2, 64))
tensor([[0.8125],
        [9.8983]])

5.4.3. 小结

  • 我们可以通过基本层类设计自定义层。这允许我们定义灵活的新层,其行为与深度学习框架中的任何现有层不同。

  • 在自定义层定义完成后,我们就可以在任意环境和网络架构中调用该自定义层。

  • 层可以有局部参数,这些参数可以通过内置函数创建。

5.4.4. 练习

  1. 设计一个接受输入并计算张量降维的层,它返回\(y_k = \sum_{i, j} W_{ijk} x_i x_j\)

  2. 设计一个返回输入数据的傅立叶系数前半部分的层。