PyTorch 的容器 Container 包括 ( nn.ModuleList, nn.ModuleDict, nn.Sequential )
PyTorch 几种构建网络的容器类
class MyModel1(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.conv1 = nn.Conv2d(1, 6, 5)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(6, 16, 5)
self.flatten = nn.Flatten()
self.linears = [nn.Linear(1024, 512), nn.Linear(512, 128), nn.Linear(128, 10)]
def forward(self, x):
x = F.relu(self.conv1(x))
x = self.pool(x)
x = F.relu(self.conv2(x))
x = self.flatten(x)
x = F.relu(self.linears[0](x))
x = F.relu(self.linears[1](x))
x = self.linears[2](x)
return x
在继承 nn.Modules 来构造自定义 module 时，会通过父类的__setattr__ 来识别 module 的成员属性是否为 nn.Modules 的子类来将其注册为模块的子模块。但是，如果子模块没用直接赋值给对象属性，而是通过 Python 的 list 封装的，那么子模块就不会被注册。
>>> net = MyModel1()
>>> net
MyModel1(
(conv1): Conv2d(1, 6, kernel_size=(5, 5), stride=(1, 1))
(pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))
(flatten): Flatten(start_dim=1, end_dim=-1)
)
1. nn.ModuleList
为了解决上述的问题，PyTorch 中定义了 nn.ModuleList 对象。nn.ModuleList 也继承 nn.Module 类，并且实现了和 Python 内置list 相同的功能。
class ModuleList(nn.Module):
# ModuleList的特殊方法。
def __init__(self, modules):
r"""
parameter:
modules: iterable, optional, 要添加的modules的迭代器
"""
pass
def append(self, module):
r"""
将module添加到列表的最后。
parameter:
modules: iterable, optional, 要添加的module
"""
pass
def extend(self, modules):
r"""
将从modules的迭代器添加modules到列表的最后。
parameter:
modules: iterable, optional, 要添加的modules的迭代器
"""
pass
def insert(self, index, module):
r"""
parameter:
index: int, 插入的位置索引
modules: iterable, optional, 要添加的module
"""
pass
nn.ModuleList 可以被识别添加为定义模块的子模块，extend、append 和 insert 方法和 list 中的用法类似。
class MyModel2(nn.Module):
def __init__(self):
super(MyModel2, self).__init__()
self.conv1 = nn.Conv2d(1, 6, 5)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(6, 16, 5)
self.flatten = nn.Flatten()
self.linears = nn.ModuleList([nn.Linear(1024, 512)])
self.linears.extend([nn.Linear(512, 128)])
self.linears.append(nn.Linear(128, 10))
def forward(self, x):
x = F.relu(self.conv1(x))
x = self.pool(x)
x = F.relu(self.conv2(x))
x = self.flatten(x)
x = F.relu(self.linears[0](x))
x = F.relu(self.linears[1](x))
x = self.linears[2](x)
return x
net = MyModel2()
net
输出：
MyModel2(
(conv1): Conv2d(1, 6, kernel_size=(5, 5), stride=(1, 1))
(pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))
(flatten): Flatten(start_dim=1, end_dim=-1)
(linears): ModuleList(
(0): Linear(in_features=1024, out_features=512, bias=True)
(1): Linear(in_features=512, out_features=128, bias=True)
(2): Linear(in_features=128, out_features=10, bias=True)
)
)
nn.ParameterList 使用和 nn.ModuleList 类似，可以将 nn.Parameter 添加为模块的参数。
2. nn.ModuleDict
class ModuleDict(nn.Module):
def __init__(self, modules):
r"""
parameter:
modules: iterable, optional, string:module的字典，或(string, module)的迭代器
"""
pass
def clear(self):
r"""移除ModuleDict中的全部item"""
pass
def items(self):
r"""返回一个key/value对的迭代器"""
pass
def keys(self):
r"""返回一个key的迭代器"""
pass
def values(self):
r"""返回一个ModuleDict值的迭代器"""
pass
def pop(self, key):
r"""从ModuleDict中移除key，并返回对应的module"""
pass
def update(self, modules):
r"""
更新dict
parameters:
modules: iterable, string 和 module的映射类型
"""
pass
定义自定义网络：
class MyModel3(nn.Module):
def __init__(self):
super(MyModel3, self).__init__()
self.conv1 = nn.Conv2d(1, 6, 5)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(6, 16, 5)
self.flatten = nn.Flatten()
self.linears = nn.ModuleDict({'linear1': nn.Linear(1024, 512),
'linear2': nn.Linear(512, 128)})
self.linears.update({'linear3': nn.Linear(128, 10)})
def forward(self, x):
x = F.relu(self.conv1(x))
x = self.pool(x)
x = F.relu(self.conv2(x))
x = self.flatten(x)
x = F.relu(self.linears['linear1'](x))
x = F.relu(self.linears['linear2'](x))
x = self.linears['linear3'](x)
return x
net = MyModel3()
print(net)
输出：
MyModel2(
(conv1): Conv2d(1, 6, kernel_size=(5, 5), stride=(1, 1))
(pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))
(flatten): Flatten(start_dim=1, end_dim=-1)
(linears): ModuleDict(
(linear1): Linear(in_features=1024, out_features=512, bias=True)
(linear2): Linear(in_features=512, out_features=128, bias=True)
(linear3): Linear(in_features=128, out_features=10, bias=True)
)
)
3. nn.Sequential
nn.Sequential 通用继承了 nn.Module 类，与Module不同的是，Sequential 已经默认定义了forward函数，按照顺序依次输入输出。
例1：
class MyModel4(nn.Module):
def __init__(self):
super(MyModel4, self).__init__()
self.conv1 = nn.Conv2d(1, 6, 5)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(6, 16, 5)
self.flatten = nn.Flatten()
self.linears = nn.Sequential(nn.Linear(1024, 512),
nn.Linear(512, 128),
nn.Linear(128, 10))
def forward(self, x):
x = F.relu(self.conv1(x))
x = self.pool(x)
x = F.relu(self.conv2(x))
x = self.flatten(x)
x = F.relu(self.linears['linear1'](x))
x = F.relu(self.linears['linear2'](x))
x = self.linears['linear3'](x)
return x
net = MyModel4()
print(net)
输出：
MyModel4(
(conv1): Conv2d(1, 6, kernel_size=(5, 5), stride=(1, 1))
(pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))
(flatten): Flatten(start_dim=1, end_dim=-1)
(linears): Sequential(
(0): Linear(in_features=1024, out_features=512, bias=True)
(1): Linear(in_features=512, out_features=128, bias=True)
(2): Linear(in_features=128, out_features=10, bias=True)
)
)
例2：
class MyModel5(nn.Module):
def __init__(self):
super(MyModel5, self).__init__()
self.conv1 = nn.Conv2d(1, 6, 5)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(6, 16, 5)
self.flatten = nn.Flatten()
self.linears = nn.Sequential(OrderedDict(
[('1', nn.Linear(1024, 512)),
('2', nn.Linear(512, 128)),
('3', nn.Linear(128, 10))])
)
def forward(self, x):
x = F.relu(self.conv1(x))
x = self.pool(x)
x = F.relu(self.conv2(x))
x = self.flatten(x)
x = F.relu(self.linears['linear1'](x))
x = F.relu(self.linears['linear2'](x))
x = self.linears['linear3'](x)
return x
net = MyModel5()
print(net)
输出：
MyModel5(
(conv1): Conv2d(1, 6, kernel_size=(5, 5), stride=(1, 1))
(pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))
(flatten): Flatten(start_dim=1, end_dim=-1)
(linears): Sequential(
(1): Linear(in_features=1024, out_features=512, bias=True)
(2): Linear(in_features=512, out_features=128, bias=True)
(3): Linear(in_features=128, out_features=10, bias=True)
)
)