pytorch教程之nn.Module类详解—使用Module类来自定义网络层

前言：前面介绍了如何自定义一个模型——通过继承nn.Module类来实现，在__init__构造函数中申明各个层的定义，在forward中实现层之间的连接关系，实际上就是前向传播的过程。

事实上，在pytorch里面自定义层也是通过继承自nn.Module类来实现的，我前面说过，pytorch里面一般是没有层的概念，层也是当成一个模型来处理的，这里和keras是不一样的。前面介绍过，我们当然也可以直接通过继承torch.autograd.Function类来自定义一个层，但是这很不推荐，不提倡，至于为什么后面会介绍。记住一句话，keras更加注重的是层Layer、pytorch更加注重的是模型Module。

所以本文就专门来介绍如何通过nn.Module类来实现自定义层。

一、从系统预定义的层说起
1.1 Linear层的代码

import math
import torch
from torch.nn.parameter import Parameter
from .. import functional as F
from .. import init
from .module import Module
from …_jit_internal import weak_module, weak_script_method

class Linear(Module):
__constants__ = [‘bias’]

def __init__(self, in_features, out_features, bias=True):
super(Linear, self).__init__()
self.in_features = in_features
self.out_features = out_features
self.weight = Parameter(torch.Tensor(out_features, in_features))
if bias:
self.bias = Parameter(torch.Tensor(out_features))
else:
self.register_parameter(‘bias’, None)
self.reset_parameters()

def reset_parameters(self):
init.kaiming_uniform_(self.weight, a=math.sqrt(5))
if self.bias is not None:
fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
bound = 1 / math.sqrt(fan_in)
init.uniform_(self.bias, -bound, bound)

@weak_script_method
def forward(self, input):
return F.linear(input, self.weight, self.bias)

def extra_repr(self):
return ‘in_features={}, out_features={}, bias={}’.format(
self.in_features, self.out_features, self.bias is not None
)
1.2 Conv2d类的实现

class Conv2d(_ConvNd):
def __init__(self, in_channels, out_channels, kernel_size, stride=1,
padding=0, dilation=1, groups=1,
bias=True, padding_mode=’zeros’):
kernel_size = _pair(kernel_size)
stride = _pair(stride)
padding = _pair(padding)
dilation = _pair(dilation)
super(Conv2d, self).__init__(
in_channels, out_channels, kernel_size, stride, padding, dilation,
False, _pair(0), groups, bias, padding_mode)

@weak_script_method
def forward(self, input):
if self.padding_mode == ‘circular’:
expanded_padding = ((self.padding[1] + 1) // 2, self.padding[1] // 2,
(self.padding[0] + 1) // 2, self.padding[0] // 2)
return F.conv2d(F.pad(input, expanded_padding, mode=’circular’),
self.weight, self.bias, self.stride,
_pair(0), self.dilation, self.groups)
return F.conv2d(input, self.weight, self.bias, self.stride,
self.padding, self.dilation, self.groups)
1.3 初步总结

我在前面的文章里面说过，torch里面实现神经网络有两种方式

（1）高层API方法：使用torch.nn.****来实现；

（2）低层API方法：使用低层函数方法，torch.nn.functional.****来实现；

其中，我们推荐使用高层API的方法，原因如下：

高层API是使用类的形式来包装的，既然是类就可以存储参数，比如全连接层的权值矩阵、偏置矩阵等都可以作为类的属性存储着，但是低层API仅仅是实现函数的运算功能，没办法保存这些信息，会丢失参数信息，但是高层API是依赖于低层API的计算函数的，比如上面的两个层：

Linear高级层——>低层F.linear() 函数

Conv2d高级层——>低层F.conv2d()函数

1.4 自定义层的步骤

要实现一个自定义层大致分以下几个主要的步骤：

（1）自定义一个类，继承自Module类，并且一定要实现两个基本的函数，第一是构造函数__init__，第二个是层的逻辑运算函数，即所谓的前向计算函数forward函数。

（2）在构造函数_init__中实现层的参数定义。比如Linear层的权重和偏置，Conv2d层的in_channels, out_channels, kernel_size, stride=1,padding=0, dilation=1, groups=1,bias=True, padding_mode=’zeros’这一系列参数；

（3）在前向传播forward函数里面实现前向运算。这一般都是通过torch.nn.functional.***函数来实现，当然很多时候我们也需要自定义自己的运算方式。如果该层含有权重，那么权重必须是nn.Parameter类型，关于Tensor和Variable（0.3版本之前）与Parameter的区别请参阅相关的文档。简单说就是Parameter默认需要求导，其他两个类型则不会。另外一般情况下，可能的话，为自己定义的新层提供默认的参数初始化，以防使用过程中忘记初始化操作。

（4）补充：一般情况下，我们定义的参数是可以求导的，但是自定义操作如不可导，需要实现backward函数。

总结：这里其实和定义一个自定义模型是一样的，核心都是实现最基本的构造函数__init__和前向运算函数forward函数，可以参考我的前一篇文章：pytorch教程之nn.Module类详解——使用Module类来自定义模型

二、自定义层的简单例子
比如我要实现一个简单的层，这个层的功能是y=w*sqrt(x2+bias),

即输入X的平方再加上一个偏执项，再开跟根号，然后再乘以权值矩阵w，那要怎么做呢，按照上面的定义过程，我们先定义一个这样的层（即一个类），代码如下：

2.1 定义一个自定义层MyLayer

# 定义一个 my_layer.py
import torch

class MyLayer(torch.nn.Module):
”’
因为这个层实现的功能是：y=weights*sqrt(x2+bias),所以有两个参数：
权值矩阵weights
偏置矩阵bias
输入 x 的维度是（in_features,)
输出 y 的维度是（out_features,) 故而
bias 的维度是（in_fearures,)，注意这里为什么是in_features,而不是out_features，注意体会这里和Linear层的区别所在
weights 的维度是（in_features, out_features）注意这里为什么是（in_features, out_features）,而不是（out_features, in_features），注意体会这里和Linear层的区别所在
”’
def __init__(self, in_features, out_features, bias=True):
super(MyLayer, self).__init__() # 和自定义模型一样，第一句话就是调用父类的构造函数
self.in_features = in_features
self.out_features = out_features
self.weight = torch.nn.Parameter(torch.Tensor(in_features, out_features)) # 由于weights是可以训练的，所以使用Parameter来定义
if bias:
self.bias = torch.nn.Parameter(torch.Tensor(in_features)) # 由于bias是可以训练的，所以使用Parameter来定义
else:
self.register_parameter(‘bias’, None)

def forward(self, input):
input_=torch.pow(input,2)+self.bias
y=torch.matmul(input_,self.weight)
return y
2.2 自定义模型并且训练

import torch
from my_layer import MyLayer # 自定义层

N, D_in, D_out = 10, 5, 3 # 一共10组样本，输入特征为5，输出特征为3

# 先定义一个模型
class MyNet(torch.nn.Module):
def __init__(self):
super(MyNet, self).__init__() # 第一句话，调用父类的构造函数
self.mylayer1 = MyLayer(D_in,D_out)

def forward(self, x):
x = self.mylayer1(x)

return x

model = MyNet()
print(model)
”’运行结果为：
MyNet(
(mylayer1): MyLayer() # 这就是自己定义的一个层
)
”’
下面开始训练

# 创建输入、输出数据
x = torch.randn(N, D_in) #（10，5）
y = torch.randn(N, D_out) #（10，3）

#定义损失函数
loss_fn = torch.nn.MSELoss(reduction=’sum’)

learning_rate = 1e-4
#构造一个optimizer对象
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

for t in range(10): #

# 第一步：数据的前向传播，计算预测值p_pred
y_pred = model(x)

# 第二步：计算计算预测值p_pred与真实值的误差
loss = loss_fn(y_pred, y)
print(f”第 {t} 个epoch, 损失是 {loss.item()}”)

# 在反向传播之前，将模型的梯度归零，这
optimizer.zero_grad()

# 第三步：反向传播误差
loss.backward()

# 直接通过梯度一步到位，更新完整个网络的训练参数
optimizer.step()
程序的运行结果为：

第 0 个epoch, 损失是 42.33689498901367
第 1 个epoch, 损失是 42.3133430480957
第 2 个epoch, 损失是 42.289825439453125
第 3 个epoch, 损失是 42.26634979248047
第 4 个epoch, 损失是 42.2429084777832
第 5 个epoch, 损失是 42.21950912475586
第 6 个epoch, 损失是 42.19615173339844
第 7 个epoch, 损失是 42.17283248901367
第 8 个epoch, 损失是 42.14955520629883
第 9 个epoch, 损失是 42.12631607055664
总结：

本文的实践说明了如何使用Module父类来拓展实现自定义模型、自定义层，我们发现二者有异曲同工之处，这也是pytorch如此受欢迎的原因之一了。后面还会继续讲解通过Function来自定义一个层。需要注意的是：

Function与Module都可以对pytorch进行自定义拓展，使其满足网络的需求，但这两者还是有十分重要的不同，具体的不同后面再说。