5.4. 自定义层
Open the notebook in Colab
Open the notebook in Colab
Open the notebook in Colab

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

5.4.1. 不带参数的层

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

from mxnet import np, npx
from mxnet.gluon import nn

npx.set_np()

class CenteredLayer(nn.Block):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def forward(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()
import tensorflow as tf


class CenteredLayer(tf.keras.Model):
    def __init__(self):
        super().__init__()

    def call(self, inputs):
        return inputs - tf.reduce_mean(inputs)

让我们通过向其提供一些数据来验证该层是否按预期工作。

layer = CenteredLayer()
layer(np.array([1, 2, 3, 4, 5]))
array([-2., -1.,  0.,  1.,  2.])
layer = CenteredLayer()
layer(torch.FloatTensor([1, 2, 3, 4, 5]))
tensor([-2., -1.,  0.,  1.,  2.])
layer = CenteredLayer()
layer(tf.constant([1, 2, 3, 4, 5]))
<tf.Tensor: shape=(5,), dtype=int32, numpy=array([-2, -1,  0,  1,  2], dtype=int32)>

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

net = nn.Sequential()
net.add(nn.Dense(128), CenteredLayer())
net.initialize()
net = nn.Sequential(nn.Linear(8, 128), CenteredLayer())
net = tf.keras.Sequential([tf.keras.layers.Dense(128), CenteredLayer()])

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

Y = net(np.random.uniform(size=(4, 8)))
Y.mean()
array(3.783498e-10)
Y = net(torch.rand(4, 8))
Y.mean()
tensor(-3.7253e-09, grad_fn=<MeanBackward0>)
Y = net(tf.random.uniform((4, 8)))
tf.reduce_mean(Y)
<tf.Tensor: shape=(), dtype=float32, numpy=-4.656613e-10>

5.4.2. 带参数的层

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

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

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

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

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

dense = MyDense(units=3, in_units=5)
dense.params
mydense0_ (
  Parameter mydense0_weight (shape=(5, 3), dtype=<class 'numpy.float32'>)
  Parameter mydense0_bias (shape=(3,), dtype=<class 'numpy.float32'>)
)
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)

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

linear = MyLinear(5, 3)
linear.weight
Parameter containing:
tensor([[-0.8548,  0.3336,  0.4857],
        [ 0.3629,  1.3252,  0.2552],
        [ 0.4776,  0.1359,  0.3208],
        [-0.7780,  0.8034,  0.9160],
        [ 0.9042,  0.4283,  0.8978]], requires_grad=True)
class MyDense(tf.keras.Model):
    def __init__(self, units):
        super().__init__()
        self.units = units

    def build(self, X_shape):
        self.weight = self.add_weight(name='weight',
            shape=[X_shape[-1], self.units],
            initializer=tf.random_normal_initializer())
        self.bias = self.add_weight(
            name='bias', shape=[self.units],
            initializer=tf.zeros_initializer())

    def call(self, X):
        linear = tf.matmul(X, self.weight) + self.bias
        return tf.nn.relu(linear)

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

dense = MyDense(3)
dense(tf.random.uniform((2, 5)))
dense.get_weights()
[array([[-0.01717934, -0.0019761 , -0.04531218],
        [ 0.03347373, -0.03442077, -0.07020961],
        [ 0.09155419,  0.02316131,  0.00039486],
        [ 0.0046796 , -0.07209789, -0.04434246],
        [-0.01901274, -0.03947276, -0.02506148]], dtype=float32),
 array([0., 0., 0.], dtype=float32)]

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

dense.initialize()
dense(np.random.uniform(size=(2, 5)))
array([[0.        , 0.01633355, 0.        ],
       [0.        , 0.01581812, 0.        ]])
linear(torch.rand(2, 5))
tensor([[0.0000, 0.8970, 2.0717],
        [0.0000, 1.2271, 3.3805]])
dense(tf.random.uniform((2, 5)))
<tf.Tensor: shape=(2, 3), dtype=float32, numpy=
array([[0.02730487, 0.        , 0.        ],
       [0.09110244, 0.        , 0.        ]], dtype=float32)>

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

net = nn.Sequential()
net.add(MyDense(8, in_units=64),
        MyDense(1, in_units=8))
net.initialize()
net(np.random.uniform(size=(2, 64)))
array([[0.06508517],
       [0.0615553 ]])
net = nn.Sequential(MyLinear(64, 8), MyLinear(8, 1))
net(torch.rand(2, 64))
tensor([[0.4589],
        [0.0000]])
net = tf.keras.models.Sequential([MyDense(8), MyDense(1)])
net(tf.random.uniform((2, 64)))
<tf.Tensor: shape=(2, 1), dtype=float32, numpy=
array([[0.],
       [0.]], dtype=float32)>

5.4.3. 小结

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

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

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

5.4.4. 练习

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

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