torch.nn的核心数据结构是Module,它是一个抽象概念,既可以表示神经网络中的某个层(layer),也可以表示一个包含很多层的神经网络。在实际使用中,最常见的做法是继承nn.Module,撰写自己的网络/层。下面先来看看如何用nn.Module实现自己的全连接层。全连接层,又名仿射层,输出 y \textbf{y} y和输入 x \textbf{x} x满足 y=Wx+b \textbf{y=Wx+b} y=Wx+b, W \textbf{W} W和 b \textbf{b} b是可学习的参数。
import torch as t
import torch.nn as nn
class Linear(nn.Module): # 继承nn.Module
def __init__(self, in_features, out_features):
super(Linear, self).__init__() # 等价于nn.Module.__init__(self)
self.w = nn.Parameter(t.randn(in_features, out_features))
self.b = nn.Parameter(t.randn(out_features))
def forward(self, x):
x = x.mm(self.w) # x.@(self.w)
return x + self.b.expand_as(x)
layer = Linear(4,3)
input = t.randn(2,4)
output = layer(input)
print(output)
结果:
tensor([[-2.2626, 1.9162, -6.1646],
[-1.4542, -0.1752, 1.1621]], grad_fn=<AddBackward0>)
查看w和b
for name, parameter in layer.named_parameters():
print(name, parameter) # w and b
结果:
w Parameter containing:
tensor([[ 5.5451e-01, 1.0340e+00, 1.3668e-01],
[-2.6401e-01, 4.8274e-01, 1.3362e+00],
[-2.6138e-01, 1.0158e+00, -1.7952e+00],
[ 4.8215e-01, -4.1644e-01, 7.1060e-04]], requires_grad=True)
b Parameter containing:
tensor([-1.3932, 0.7958, -1.1102], requires_grad=True)
可见,全连接层的实现非常简单,但需注意以下几点:
- 自定义层Linear必须继承nn.Module,并且在其构造函数中需调用nn.Module的构造函数,即super(Linear,
self).init() 或nn.Module.init(self),推荐使用第一种用法,尽管第二种写法更直观。 - 在构造函数__init__中必须自己定义可学习的参数,并封装成Parameter,如在本例中我们把w和b封装成parameter。parameter是一种特殊的Tensor,但其默认需要求导(requires_grad
= True),感兴趣的读者可以通过nn.Parameter查看Parameter类的源代码。 - forward函数实现前向传播过程,其输入可以是一个或多个tensor。
- 无需写反向传播函数,nn.Module能够利用autograd自动实现反向传播,这点比Function简单许多。
- 使用时,直观上可将layer看成数学概念中的函数,调用layer(input)即可得到input对应的结果。它等价于layers.call(input),在__call__函数中,主要调用的是 layer.forward(x),另外还对钩子做了一些处理。所以在实际使用中应尽量使用layer(x)而不是使用layer.forward(x)。
- Module中的可学习参数可以通过named_parameters()或者parameters()返回迭代器,前者会给每个parameter都附上名字,使其更具有辨识度。
Module能够自动检测到自己的Parameter,并将其作为学习参数。除了parameter之外,Module还包含子Module,主Module能够递归查找子Module中的parameter。下面再来看看稍微复杂一点的网络,多层感知机。
多层感知机的网络结构如下图所示,它由两个全连接层组成,采用
s
i
g
m
o
i
d
sigmoid
sigmoid函数作为激活函数,图中没有画出。
import torch as t
import torch.nn as nn
class Linear(nn.Module): # 继承nn.Module
def __init__(self, in_features, out_features):
super(Linear, self).__init__() # 等价于nn.Module.__init__(self)
self.w = nn.Parameter(t.randn(in_features, out_features))
self.b = nn.Parameter(t.randn(out_features))
def forward(self, x):
x = x.mm(self.w) # x.@(self.w)
return x + self.b.expand_as(x)
class Perceptron(nn.Module):
def __init__(self, in_features, hidden_features, out_features):
nn.Module.__init__(self)
self.layer1 = Linear(in_features, hidden_features)
self.layer2 = Linear(hidden_features, out_features)
def forward(self,x):
x = self.layer1(x)
x = t.sigmoid(x)#激活函数
return self.layer2(x)
perceptron = Perceptron(3,4,1)
for name, param in perceptron.named_parameters():#查看每层w和b的大小
print(name, param.size())
输出结果:
layer1.w torch.Size([3, 4])
layer1.b torch.Size([4])
layer2.w torch.Size([4, 1])
layer2.b torch.Size([1])
module中parameter的命名规范:
- 对于类似self.param_name = nn.Parameter(t.randn(3, 4)),命名为param_name
- 对于子Module中的parameter,会其名字之前加上当前Module的名字。如对于self.sub_module = SubModel(),SubModel中有个parameter的名字叫做param_name,那么二者拼接而成的parameter name 就是sub_module.param_name。
注意:自定义layer对输入形状都有假设:输入的不是单个数据,而是一个batch。输入只有一个数据,则必须调用tensor.unsqueeze(0) 或 tensor[None]将数据伪装成batch_size=1的batch。