PyTorch教程-6.5. 自定義圖層

深度學(xué)習(xí)成功背后的一個(gè)因素是廣泛的層的可用性，這些層可以以創(chuàng)造性的方式組合以設(shè)計(jì)適合各種任務(wù)的架構(gòu)。例如，研究人員發(fā)明了專(zhuān)門(mén)用于處理圖像、文本、循環(huán)順序數(shù)據(jù)和執(zhí)行動(dòng)態(tài)規(guī)劃的層。遲早，您會(huì)遇到或發(fā)明深度學(xué)習(xí)框架中尚不存在的層。在這些情況下，您必須構(gòu)建自定義層。在本節(jié)中，我們將向您展示如何操作。

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. 沒(méi)有參數(shù)的圖層

首先，我們構(gòu)建一個(gè)自定義層，它自己沒(méi)有任何參數(shù)。如果您還記得我們?cè)诘?6.1 節(jié)中對(duì)模塊的介紹，這應(yīng)該看起來(lái)很熟悉。以下 CenteredLayer類(lèi)只是從其輸入中減去平均值。要構(gòu)建它，我們只需要繼承基礎(chǔ)層類(lèi)并實(shí)現(xiàn)前向傳播功能。

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)

讓我們通過(guò)提供一些數(shù)據(jù)來(lái)驗(yàn)證我們的層是否按預(yù)期工作。

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

我們現(xiàn)在可以將我們的層合并為構(gòu)建更復(fù)雜模型的組件。

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

作為額外的健全性檢查，我們可以通過(guò)網(wǎng)絡(luò)發(fā)送隨機(jī)數(shù)據(jù)并檢查均值實(shí)際上是否為 0。因?yàn)槲覀兲幚淼氖歉↑c(diǎn)數(shù)，由于量化，我們可能仍然會(huì)看到非常小的非零數(shù)。

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. 帶參數(shù)的圖層

現(xiàn)在我們知道如何定義簡(jiǎn)單的層，讓我們繼續(xù)定義具有可通過(guò)訓(xùn)練調(diào)整的參數(shù)的層。我們可以使用內(nèi)置函數(shù)來(lái)創(chuàng)建參數(shù)，這些參數(shù)提供了一些基本的內(nèi)務(wù)處理功能。特別是，它們管理訪問(wèn)、初始化、共享、保存和加載模型參數(shù)。這樣，除了其他好處之外，我們將不需要為每個(gè)自定義層編寫(xiě)自定義序列化例程。

現(xiàn)在讓我們實(shí)現(xiàn)我們自己的全連接層版本。回想一下，該層需要兩個(gè)參數(shù)，一個(gè)代表權(quán)重，另一個(gè)代表偏差。在此實(shí)現(xiàn)中，我們將 ReLU 激活作為默認(rèn)值進(jìn)行烘焙。該層需要兩個(gè)輸入?yún)?shù)： in_units和units，分別表示輸入和輸出的數(shù)量。

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)

接下來(lái)，我們實(shí)例化該類(lèi)MyLinear并訪問(wèn)其模型參數(shù)。

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

我們可以直接使用自定義層進(jìn)行前向傳播計(jì)算。

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

我們還可以使用自定義層構(gòu)建模型。一旦我們有了它，我們就可以像使用內(nèi)置的全連接層一樣使用它。

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. 概括

我們可以通過(guò)基本層類(lèi)來(lái)設(shè)計(jì)自定義層。這使我們能夠定義靈活的新層，這些層的行為不同于庫(kù)中的任何現(xiàn)有層。一旦定義好，自定義層就可以在任意上下文和架構(gòu)中被調(diào)用。層可以有局部參數(shù)，可以通過(guò)內(nèi)置函數(shù)創(chuàng)建。

6.5.4. 練習(xí)

設(shè)計(jì)一個(gè)接受輸入并計(jì)算張量縮減的層，即它返回yk=∑i,jWijkxixj.

設(shè)計(jì)一個(gè)返回?cái)?shù)據(jù)傅里葉系數(shù)前半部分的層。