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

---

前言

当深度学习模型逐渐变得复杂，在编写代码时便会遇到诸多麻烦，此时便需要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 code
 model = 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
 # ParameterList
 class 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 ints
        for i, p in enumerate(self.params):
            x = self.params[i // 2].mm(x) + p.mm(x)
        return x

这是专门为Pytorch中的参数（如权重和偏置）设计的容器。它确保添加到其中的参数会被正确地注册到网络中，以便进行参数管理和梯度更新。与module类似，参数容器中的参数也会被包含在网络的参数列表中，并在模型保存和加载时被正确处理。

---

总结

容器是pytorch框架对网络进行组织管理的实用工具，合理运用可以极大提高代码的可读性与可维护性。