参考
4.4 自定义层
深度学习的一个魅力在于神经网络中各式各样的层,例如全连接层和后面章节将要用介绍的卷积层、池化层与循环层。虽然PyTorch提供了大量常用的层,但有时候我们依然希望自定义层。本节将介绍如何使用Module来自定义层,从而可以被重复调用。
4.4.1 不含模型参数的自定义层
我们先介绍如何定义一个不含模型参数的自定义层。
import torchfrom torch import nnclass CenteredLayer(nn.Module):def __init__(self, **kwargs):super(CenteredLayer, self).__init__(**kwargs)def forward(self, x):return x - x.mean()
layer = CenteredLayer()layer(torch.tensor([1, 2, 3, 4, 5], dtype=torch.float))
我们也可以用它来构造更复杂的模型。
net = nn.Sequential(nn.Linear(8, 128), CenteredLayer())
y = net(torch.rand(4, 8))y.mean().item()
4.4.2 含模型参数的自定义层
我们还可以自定义含模型参数的自定义层。其中的模型参数可以通过训练学习。
Parameter
类其实是Tensor
的子类,如果一个Tensor是Parameter
,那么它会自动被添加到模型的参数列表里。所以在自定义含模型参数的层时,我们应该将参数定义成Parameter
,除了像4.2.1节那样直接定义成Parameter
类外,还可以使用ParameterList
和ParameterDict
分别定义参数的列表和字典。
ParameterList
接收一个Parameter
实例的列表作为输入然后得到一个参数列表,使用的时候可以用索引来访问某个参数,另外也可以使用append
和extend
在列表后面新增参数。
class MyDense(nn.Module):def __init__(self):super(MyDense, self).__init__()self.params = nn.ParameterList([nn.Parameter(torch.randn(4, 4)) for i in range(3)])self.params.append(nn.Parameter(torch.randn(4, 1)))def forward(self, x):for i in range(len(self.params)):x = torch.mm(x, self.params[i])return xnet = MyDense()print(net)
而ParameterDict
接收一个Parameter
实例的字典作为输入然后得到一个参数字典,然后可以按照字典的规则使用了。
class MyDictDense(nn.Module):def __init__(self):super(MyDictDense, self).__init__()self.params = nn.ParameterDict({'linear1': nn.Parameter(torch.randn(4, 4)),'linear2': nn.Parameter(torch.randn(4, 1))})self.params.update({'linear3': nn.Parameter(torch.randn(4, 2))})def forward(self, x, choice='linear1'):return torch.mm(x, self.params[choice])net = MyDictDense()print(net)
x = torch.ones(1, 4)print(net(x, 'linear1'))print(net(x, 'linear2'))print(net(x, 'linear3'))
我们也可以使用自定义层构造模型。它和PyTorch的其他层在使用上很类似。
net = nn.Sequential(MyDictDense(),MyDictDense())print(net)print(net(x))
如果觉得《[pytorch 学习] - 4.4 自定义层》对你有帮助,请点赞、收藏,并留下你的观点哦!