Pytorch自定义网络层

Pytorch、Tensoflow等许多深度学习框架集成了大量常见的网络层，为我们搭建神经网络提供了诸多便利。但在实际工作中，因为项目要求、研究需要或者发论文需要等等，大家一般都会需要自己发明一个现在在深度学习框架中还不存在的层。 在这些情况下，就必须构建自定义层。

博主在学习了沐神的动手学深度学习这本书之后，学到了许多东西。这里记录一下书中基于Pytorch实现简单自定义网络层的方法，仅供参考。

一、不带参数的层

首先，我们构造一个没有任何参数的自定义层，要构建它，只需继承基础层类并实现前向传播功能。

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()
print(layer(torch.FloatTensor([1, 2, 3, 4, 5])))

输出结果如下：

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

运行正常，表明网络没有问题。

现在将我们自建的网络层作为组件合并到更复杂的模型中，并输入数据进行验证：

net = nn.Sequential(nn.Linear(8, 128), CenteredLayer())
Y = net(torch.rand(4, 8))
print(Y.mean())  # 因为模型参数较多，输出也较多，所以这里输出Y的均值，验证模型可运行即可

结果如下：

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

二、带参数的层

这里使用内置函数来创建参数，这些函数可以提供一些基本的管理功能，使用更加方便。

这里实现了一个简单的自定义的全连接层，大家可根据需要自行修改即可。

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)

接下来实例化类并访问其模型参数：

linear = MyLinear(5, 3)
print(linear.weight)

结果如下：

Parameter containing:
tensor([[-0.3708,  1.2196,  1.3658],
        [ 0.4914, -0.2487, -0.9602],
        [ 1.8458,  0.3016, -0.3956],
        [ 0.0616, -0.3942,  1.6172],
        [ 0.7839,  0.6693, -0.8890]], requires_grad=True)

而后输入一些数据，查看模型输出结果：

print(linear(torch.rand(2, 5)))
# 结果如下
tensor([[1.2394, 0.0000, 0.0000],
        [1.3514, 0.0968, 0.6667]])

我们还可以使用自定义层构建模型，使用方法与使用内置的全连接层相同。

net = nn.Sequential(MyLinear(64, 8), MyLinear(8, 1))
print(net(torch.rand(2, 64)))
# 结果如下
tensor([[4.1416],
        [0.2567]])

三、总结

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

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

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

四、参考

《动手学深度学习》 — 动手学深度学习 2.0.0-beta0 documentationhttps://zh-v2.d2l.ai/