6.5. 自定义层
在 Colab 中打开 Notebook
在 Colab 中打开 Notebook
在 Colab 中打开 Notebook
在 Colab 中打开 Notebook
在 SageMaker Studio Lab 中打开 Notebook

深度学习成功背后的一个因素是其拥有范围广大的层，我们可以通过创造性的方式组合这些层，从而设计出适用于各种任务的架构。例如，研究人员发明了专门用于处理图像、文本、循环遍历时序数据和执行动态规划的层。迟早有一天，你会需要一个在深度学习框架中尚不存在的层。在这些情况下，你必须构建一个自定义层。在本节中，我们将向你展示如何构建。

pytorchmxnetjaxtensorflow
import torch
from torch import nn
from torch.nn import functional as F
from d2l import torch as d2l

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

npx.set_np()

import jax
from flax import linen as nn
from jax import numpy as jnp
from d2l import jax as d2l

No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

import tensorflow as tf
from d2l import tensorflow as d2l

6.5.1. 不带参数的层

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

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

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

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

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

class CenteredLayer(nn.Module):
    def __call__(self, X):
        return X - X.mean()

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

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

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

pytorchmxnetjaxtensorflow
layer = CenteredLayer()
layer(torch.tensor([1.0, 2, 3, 4, 5]))

tensor([-2., -1.,  0.,  1.,  2.])

layer = CenteredLayer()
layer(np.array([1.0, 2, 3, 4, 5]))

[21:49:18] ../src/storage/storage.cc:196: Using Pooled (Naive) StorageManager for CPU

array([-2., -1.,  0.,  1.,  2.])

layer = CenteredLayer()
layer(jnp.array([1.0, 2, 3, 4, 5]))

Array([-2., -1.,  0.,  1.,  2.], dtype=float32)

layer = CenteredLayer()
layer(tf.constant([1.0, 2, 3, 4, 5]))

<tf.Tensor: shape=(5,), dtype=float32, numpy=array([-2., -1.,  0.,  1.,  2.], dtype=float32)>

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

pytorchmxnetjaxtensorflow
net = nn.Sequential(nn.LazyLinear(128), CenteredLayer())

net = nn.Sequential()
net.add(nn.Dense(128), CenteredLayer())
net.initialize()

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

net = tf.keras.Sequential([tf.keras.layers.Dense(128), CenteredLayer()])

作为额外的健全性检查，我们可以通过网络发送随机数据并检查均值是否确实为0。因为我们处理的是浮点数，由于量化，我们可能仍然会看到一个非常小的非零数。

pytorchmxnetjaxtensorflow
Y = net(torch.rand(4, 8))
Y.mean()

tensor(-6.5193e-09, grad_fn=<MeanBackward0>)

Y = net(np.random.rand(4, 8))
Y.mean()

array(3.783498e-10)

这里我们利用 init_with_output 方法，它同时返回网络的输出和参数。在这种情况下，我们只关注输出。

Y, _ = net.init_with_output(d2l.get_key(), jax.random.uniform(d2l.get_key(),
                                                              (4, 8)))
Y.mean()

Array(-3.7252903e-09, dtype=float32)

Y = net(tf.random.uniform((4, 8)))
tf.reduce_mean(Y)

<tf.Tensor: shape=(), dtype=float32, numpy=-1.8626451e-09>

6.5.2. 带参数的层

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

现在让我们实现我们自己的全连接层版本。回想一下，该层需要两个参数，一个表示权重，另一个表示偏置。在此实现中，我们将ReLU激活作为默认值。该层需要两个输入参数：in_units 和 units，分别表示输入和输出的数量。

pytorchmxnetjaxtensorflow
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.4783,  0.4284, -0.0899],
        [-0.6347,  0.2913, -0.0822],
        [-0.4325, -0.1645, -0.3274],
        [ 1.1898,  0.6482, -1.2384],
        [-0.1479,  0.0264, -0.9597]], requires_grad=True)

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 MyDense(nn.Module):
    in_units: int
    units: int

    def setup(self):
        self.weight = self.param('weight', nn.initializers.normal(stddev=1),
                                 (self.in_units, self.units))
        self.bias = self.param('bias', nn.initializers.zeros, self.units)

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

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

dense = MyDense(5, 3)
params = dense.init(d2l.get_key(), jnp.zeros((3, 5)))
params

FrozenDict({
    params: {
        weight: Array([[-0.23823419, -0.70915407,  0.72494346],
               [ 0.2568525 , -0.20872341, -0.8993567 ],
               [ 0.80883664,  0.16673394,  0.75610644],
               [-0.35652584,  0.13841456, -1.0971175 ],
               [ 0.3117082 ,  1.2280334 , -1.0946037 ]], dtype=float32),
        bias: Array([0., 0., 0.], dtype=float32),
    },
})

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.01007051, -0.05935554,  0.03142897],
        [ 0.02453684, -0.01833588, -0.03096254],
        [-0.09680572, -0.01736571, -0.00858052],
        [-0.02245625,  0.02958351, -0.05780673],
        [ 0.03997313,  0.01949595, -0.00150928]], dtype=float32),
 array([0., 0., 0.], dtype=float32)]

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

pytorchmxnetjaxtensorflow
linear(torch.rand(2, 5))

tensor([[0.0000, 0.9316, 0.0000],
        [0.1808, 1.4208, 0.0000]])

dense.initialize()
dense(np.random.uniform(size=(2, 5)))

array([[0.        , 0.01633355, 0.        ],
       [0.        , 0.01581812, 0.        ]])

dense.apply(params, jax.random.uniform(d2l.get_key(),
                                       (2, 5)))

Array([[0.3850514 , 0.        , 0.49188882],
       [0.46509624, 0.26056105, 0.        ]], dtype=float32)

dense(tf.random.uniform((2, 5)))

<tf.Tensor: shape=(2, 3), dtype=float32, numpy=
array([[0., 0., 0.],
       [0., 0., 0.]], dtype=float32)>

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

pytorchmxnetjaxtensorflow
net = nn.Sequential(MyLinear(64, 8), MyLinear(8, 1))
net(torch.rand(2, 64))

tensor([[ 0.0000],
        [13.0800]])

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([MyDense(64, 8), MyDense(8, 1)])
Y, _ = net.init_with_output(d2l.get_key(), jax.random.uniform(d2l.get_key(),
                                                              (2, 64)))
Y

Array([[0.],
       [0.]], dtype=float32)

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.00309976]], dtype=float32)>

6.5.3. 小结

我们可以通过基础层类设计自定义层。这使我们能够定义灵活的新层，其行为与库中任何现有层都不同。一旦定义，自定义层可以在任意上下文和架构中调用。层可以有局部参数，这些参数可以通过内置函数创建。

6.5.4. 练习

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

  2. 设计一个层，返回数据的前一半傅里叶系数。

pytorchmxnetjaxtensorflow

讨论

讨论

讨论

讨论