PyTorch教程-6.5. 自定义图层

深度学习成功背后的一个因素是广泛的层的可用性，这些层可以以创造性的方式组合以设计适合各种任务的架构。例如，研究人员发明了专门用于处理图像、文本、循环顺序数据和执行动态规划的层。迟早，您会遇到或发明深度学习框架中尚不存在的层。在这些情况下，您必须构建自定义层。在本节中，我们将向您展示如何操作。

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类只是从其输入中减去平均值。要构建它，我们只需要继承基础层类并实现前向传播功能。

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)

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

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]))

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]))

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

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。因为我们处理的是浮点数，由于量化，我们可能仍然会看到非常小的非零数。

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

tensor(0., grad_fn=)

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

array(3.783498e-10)

Here we utilize the init_with_output method which returns both the output of the network as well as the parameters. In this case we only focus on the output.

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

Array(5.5879354e-09, dtype=float32)

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

6.5.2. 带参数的图层

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

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

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([[-1.2894e+00, 6.5869e-01, -1.3933e+00],
    [ 7.2590e-01, 7.1593e-01, 1.8115e-03],
    [-1.5900e+00, 4.1654e-01, -1.3358e+00],
    [ 2.2732e-02, -2.1329e+00, 1.8811e+00],
    [-1.0993e+00, 2.9763e-01, -1.4413e+00]], 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)

Next, we instantiate the MyDense class and access its model parameters.

dense = MyDense(units=3, in_units=5)
dense.params

mydense0_ (
 Parameter mydense0_weight (shape=(5, 3), dtype=)
 Parameter mydense0_bias (shape=(3,), dtype=)
)

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)

Next, we instantiate the MyDense class and access its model parameters.

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

FrozenDict({
  params: {
    weight: Array([[-0.02040312, 1.0439496 , -2.3386796 ],
        [ 1.1002127 , -1.780812 , -0.32284564],
        [-0.6944499 , -1.8438653 , -0.5338283 ],
        [ 1.3954164 , 1.5816483 , 0.0469989 ],
        [-0.12351853, 1.2818031 , 0.7964193 ]], 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)

Next, we instantiate the MyDense class and access its model parameters.

dense = MyDense(3)
dense(tf.random.uniform((2, 5)))
dense.get_weights()

[array([[-0.08860754, -0.04000078, -0.03810905],
    [-0.0543257 , -0.01143957, -0.06748273],
    [-0.05273567, -0.01696461, -0.00552523],
    [ 0.00193098, 0.0662979 , -0.05486313],
    [-0.08595717, 0.08563109, 0.04592342]], dtype=float32),
 array([0., 0., 0.], dtype=float32)]

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

linear(torch.rand(2, 5))

tensor([[0.0000, 1.7772, 0.0000],
    [0.0000, 1.0303, 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.05256309, 0.    , 0.    ],
    [0.3500959 , 0.    , 0.30999148]], dtype=float32)

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

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

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

tensor([[0.],
    [0.]])

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([[8.348445 ],
    [2.0591817]], dtype=float32)

net = tf.keras.models.Sequential([MyDense(8), MyDense(1)])
net(tf.random.uniform((2, 64)))

6.5.3. 概括

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

6.5.4. 练习

设计一个接受输入并计算张量缩减的层，即它返回yk=∑i,jWijkxixj.

设计一个返回数据傅里叶系数前半部分的层。