当前位置：首页 > news >正文

【PyTorch】使用容器(Containers)进行网络层管理(Module)

news 来源：原创 2024/9/21 5:27:38

文章目录

前言
一、Sequential
二、ModuleList
三、ModuleDict
四、ParameterList & ParameterDict
总结

前言

当深度学习模型逐渐变得复杂，在编写代码时便会遇到诸多麻烦，此时便需要Containers的帮助。Containers的作用是将一部分网络层模块化，从而更方便地管理和调用。本文介绍PyTorch库常用的nn.Sequential，nn.ModuleList，nn.ModuleDict容器以及nn.ParameterList & ParameterDict参数容器。

一、Sequential

Sequential是最为常用的容器，它的功能也十分简单直接-将多个网络层按照固定的顺序连接，从前往后依次执行。比如在AlexNet中，多次需要conv+relu+maxpool的组合，此时便可以将其放入Sequential容器，便于在forward中调用。
下面来看PyTorch官方代码示例：

model = nn.Sequential(nn.Conv2d(1,20,5),nn.ReLU(),nn.Conv2d(20,64,5),nn.ReLU())# Using Sequential with OrderedDict. This is functionally the# same as the above codemodel = nn.Sequential(OrderedDict([('conv1', nn.Conv2d(1,20,5)),('relu1', nn.ReLU()),('conv2', nn.Conv2d(20,64,5)),('relu2', nn.ReLU())]))

示例中展示了两种Sequential使用方法：1，直接串联各个网络层。2，使用OrderedDict为每个module取名。这两种方法是等效的。

二、ModuleList

"顾名思义"ModuleList的作用如同Python的列表，将各个层存入一个类似于List的结构中，从而可以利用索引来进行调用。
注意这里是类似于list的结构，那为什么我们不直接用list呢？
ModuleList是专门为Pytorch中的神经网络模块（即继承自nn.Module的类）设计的容器。它确保所有添加到其中的模块都会正确地注册到网络中，以便进行参数管理和梯度更新。当模型被保存或加载时，nn.ModuleList中的模块也会相应地被保存或加载。而Python的列表是一个通用的容器，可以存储任意类型的对象。它没有专门为神经网络模块设计，因此不会进行参数的自动注册或管理。
代码示例：

class MyModule(nn.Module):def __init__(self):super(MyModule, self).__init__()self.linears = nn.ModuleList([nn.Linear(10, 10) for i in range(10)])# self.linears = [nn.Linear(10, 10) for i in range(10)]    def forward(self, x):for sub_layer in self.linears:x = sub_layer(x)return x

三、ModuleDict

ModuleDict是一个类似python字典的容器，相比于ModuleList，它的优点在于可以利用名字来调用网络层，这就避免了必须记住网络层具体元素才能调用的麻烦。
代码示例：

 class MyModule2(nn.Module):def __init__(self):super(MyModule2, self).__init__()self.choices = nn.ModuleDict({'conv': nn.Conv2d(3, 16, 5),'pool': nn.MaxPool2d(3)})self.activations = nn.ModuleDict({'lrelu': nn.LeakyReLU(),'prelu': nn.PReLU()})def forward(self, x, choice, act):x = self.choices[choice](x)x = self.activations[act](x)return x

四、ParameterList & ParameterDict

除了Module有容器，Parameter也有容器。与ModuleList和ModuleDict类似的，Paramter也有List和Dict，使用方法一样。

class MyModule(nn.Module):def __init__(self):super(MyModule, self).__init__()self.params = nn.ParameterDict({'left': nn.Parameter(torch.randn(5, 10)),'right': nn.Parameter(torch.randn(5, 10))})def forward(self, x, choice):x = self.params[choice].mm(x)return x# ParameterListclass MyModule(nn.Module):def __init__(self):super(MyModule, self).__init__()self.params = nn.ParameterList([nn.Parameter(torch.randn(10, 10)) for i in range(10)])def forward(self, x):# ParameterList can act as an iterable, or be indexed using intsfor i, p in enumerate(self.params):x = self.params[i // 2].mm(x) + p.mm(x)return x