4. Pytorch：Module

4.1. Container

nn.Module 是 Pytorch 中提供的一个类，是所有神经网络模块的基类。自定义的模块要继承这个基类。

下面介绍三个容器，它们都继承自 nn.Module ：

  • Sequential

  • ModuleList

  • ModuleDict

Sequential

class torch.nn.Sequential(*args)

当模型的前向计算为简单 串联 各个层的计算时，Sequential 类可以通过更加简单的方式定义模型。这正是 Sequential 类的目的：它可以接收一个子模块的有序字典（OrderedDict）或者一系列子模块作为参数来逐一添加 nn.Module 的实例，而模型的前向计算就是将这些实例 按添加的顺序 逐一计算。

 1# Example of using Sequential
 2model = nn.Sequential(
 3          nn.Conv2d(1,20,5),
 4          nn.ReLU(),
 5          nn.Conv2d(20,64,5),
 6          nn.ReLU()
 7        )
 8
 9# Example of using Sequential with OrderedDict
10model = nn.Sequential(OrderedDict([
11          ('conv1', nn.Conv2d(1,20,5)),
12          ('relu1', nn.ReLU()),
13          ('conv2', nn.Conv2d(20,64,5)),
14          ('relu2', nn.ReLU())
15        ]))

ModuleList

class torch.nn.ModuleList(modules=None)

ModuleList 接收一个子模块的列表作为输入，但 ModuleList 仅仅是一个储存各种模块的列表，这些模块之间没有联系也没有顺序（所以不用保证相邻层的输入输出维度匹配），需要在 forward 中定义前向计算的顺序。

ModuleList 不同于一般的 Python list，加入到 ModuleList 里面的所有模块的参数会被自动添加（注册）到整个网络中。

 1class MyModule(nn.Module):
 2    def __init__(self):
 3        super(MyModule, self).__init__()
 4        self.linears = nn.ModuleList([nn.Linear(10, 10) for i in range(10)])
 5
 6    def forward(self, x):
 7        # ModuleList can act as an iterable, or be indexed using ints
 8        for i, l in enumerate(self.linears):
 9            x = self.linears[i // 2](x) + l(x)
10        return x

方法：

  • append(module)

  • extend(modules)

  • insert(index, module)

参数中的 module 可以使 Pytorch 内建模块，也可以是自定义的继承自 nn.Module 的模块。

ModuleDict

class torch.nn.ModuleDict(modules=None)

ModuleDict 接收一个子模块的字典作为输入。和 ModuleList 一样，ModuleDict 实例仅仅是存放了一些模块的字典，并没有定义 forward 函数需要自己定义。同样，ModuleDict 也与 Python dict 有所不同， ModuleDict 里的所有模块的参数会被自动添加到整个网络中。

 1class MyModule(nn.Module):
 2    def __init__(self):
 3        super(MyModule, self).__init__()
 4        self.choices = nn.ModuleDict({
 5                'conv': nn.Conv2d(10, 10, 3),
 6                'pool': nn.MaxPool2d(3)
 7        })
 8        self.activations = nn.ModuleDict([
 9                ['lrelu', nn.LeakyReLU()],
10                ['prelu', nn.PReLU()]
11        ])
12
13    def forward(self, x, choice, act):
14        x = self.choices[choice](x)
15        x = self.activations[act](x)
16        return x

 1class CombinedROIHeads(nn.ModuleDict):
 2    """
 3    Combines a set of individual heads (for box prediction or masks) into a single head.
 4    -- from maskrcnn_benchmark
 5    """
 6
 7    def __init__(self, heads):
 8        super(CombinedROIHeads, self).__init__(heads)
 9
10    def forward(self):
11        pass
12
13>>> heads = []
14>>> heads.append(("box", nn.Conv2d(256, 512, kernel_size=3, stride=2)))
15>>> heads.append(("mask", nn.Conv2d(256, 512, kernel_size=3, stride=1)))
16>>> heads.append(("keypoint", nn.Linear(256, 512, bias=True)))
17>>> roi_heads = CombinedROIHeads(heads)
18>>> roi_heads
19CombinedROIHeads(
20  (box): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2))
21  (mask): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1))
22  (keypoint): Linear(in_features=256, out_features=512, bias=True)
23)

方法：

  • clear()

  • items()

  • keys()

  • values()

  • pop(key)

  • update(modules) modules (iterable) 是值为 nn.Module 的字典类型或包含 (string, nn.Module) 对的可迭代类型。

4.2. add_module

如果有一个元素是 nn.Module 的列表，直接赋值给一个模型的属性（Attribute），并不会让列表内的 Modules 立即注册（Register）为模型的模块。

1import torch
2import torch.nn as nn
3
4class A(nn.Module):
5    def __init__(self):
6        super(A, self).__init__()
7
8        self.layerList = [nn.Linear(5, 1, bias=False), nn.Linear(2,1, bias=False)]

1>>> a = A()
2>>> print a
3A(
4)

使用 add_module 可以自由地将列表的元素变成模型的模块。 add_module 建立了对 nn.Module 的引用，并不是添加了新的对象。 因此，对引用的修改会直接修改列表内的 nn.Module 。

加入之后，可以通过模型的名字来进行访问：_modules[name] 。_modules 是一个顺序字典（OrderedDict）。

 1import torch
 2import torch.nn as nn
 3
 4class A(nn.Module):
 5    def __init__(self):
 6        super(A, self).__init__()
 7        self.layerList = [nn.Linear(5, 1, bias=False), nn.Linear(2,1, bias=False)]
 8        self.add_module("layer_0", self.layerList[0])
 9        self.add_module("layer_1", self.layerList[1])
10
11        print self.layerList[0].weight
12        print self._modules['layer_0'].weight
13        self._modules['layer_0'].weight.data = self._modules['layer_0'].weight.data + 2 * torch.ones_like(self._modules['layer_0'].weight.data)
14
15    def forward(self):
16        print self.layerList[0].weight
17        print self._modules['layer_0'].weight

 1>>> a = A() ## init
 2Parameter containing:
 3 0.0244 -0.0521 -0.4013 -0.1229  0.0343
 4[torch.FloatTensor of size 1x5]
 5
 6Parameter containing:
 7 0.0244 -0.0521 -0.4013 -0.1229  0.0343
 8[torch.FloatTensor of size 1x5]
 9
10>>> print a
11A(
12  (layer_0): Linear(in_features=5, out_features=1, bias=False)
13  (layer_1): Linear(in_features=2, out_features=1, bias=False)
14)
15
16>>> a() ## forward
17Parameter containing:
18 2.0244  1.9479  1.5987  1.8771  2.0343
19[torch.FloatTensor of size 1x5]
20
21Parameter containing:
22 2.0244  1.9479  1.5987  1.8771  2.0343
23[torch.FloatTensor of size 1x5]
24## 可以看到，上面的参数是同步更新的

4.3. Attribute 索引

除了使用 _modules[name] 访问模块，还可以将 name 转换成属性（Attribute）的索引，通过下标的形式访问。

 1import torch
 2import torch.nn as nn
 3
 4class AttrProxy(object):
 5    """Translates index lookups into attribute lookups."""
 6    def __init__(self, module, prefix):
 7        self.module = module
 8        self.prefix = prefix
 9    def __getitem__(self, index):
10        return getattr(self.module, self.prefix + str(index))
11
12class A(nn.Module):
13    def __init__(self):
14        super(A, self).__init__()
15        self.layerList = [nn.Linear(5, 1, bias=False), nn.Linear(2,1, bias=False)]
16        self.add_module("layer_0", self.layerList[0])
17        self.add_module("layer_1", self.layerList[1])
18
19        self.layer = AttrProxy(self, "layer_")
20
21        print self.layerList[0].weight
22        print self.layer[0].weight
23        self.layer[0].weight.data = self.layer[0].weight.data + 2 * torch.ones_like(self.layer[0].weight.data)
24
25    def forward(self):
26        print self.layerList[0].weight
27        print self.layer[0].weight
28        print self.layer[1].weight

 1>>> a = A() ## init
 2Parameter containing:
 3-0.2655  0.1539 -0.2107  0.0740  0.1922
 4[torch.FloatTensor of size 1x5]
 5
 6Parameter containing:
 7-0.2655  0.1539 -0.2107  0.0740  0.1922
 8[torch.FloatTensor of size 1x5]
 9
10>>> a() ## forward
11Parameter containing:
12 1.7345  2.1539  1.7893  2.0740  2.1922
13[torch.FloatTensor of size 1x5]
14
15Parameter containing:
16 1.7345  2.1539  1.7893  2.0740  2.1922
17[torch.FloatTensor of size 1x5]
18
19Parameter containing:
20 0.0068 -0.1787
21[torch.FloatTensor of size 1x2]

4.4. 附：Module 部分实例方法

add_module(name, module)

向当前模块中加入新的子模块。

apply(fn)

循环地向每个子模块（ .children() ）及其自身施加 fn 函数，典型的用法包括初始化模型参数。

buffers(recurse=True)

返回模块缓存的迭代器，如 BatchNorm 需要缓存 running_mean 和 running_var。

named_buffers(prefix='', recurse=True)

返回模块缓存的迭代器，迭代器会同时产生缓存的名字和缓存的数据。

children()

返回直接子模块的迭代器。

named_children()

返回直接子模块的迭代器，迭代器会同时产生模块的名字和模块自身。

modules()

返回所有子模块的迭代器。

named_modules(memo=None, prefix='')

返回所有子模块的迭代器，迭代器会同时产生模块的名字和模块自身。

cpu()

将所有的模型参数和缓冲移到 CPU。

cuda(device=None)

将所有的模型参数和缓冲移到 GPU，必须在构造优化器之前调用。

double()

将所有的浮点型参数和缓存强制转换为 double 类型。

float()

将所有的参数和缓存强制转换为浮点类型。

half()

将所有的参数和缓存强制转换为半浮点类型。

train(mode=True)

将模块设置为训练模式。

eval()

与 train(False) 等效。

forward(*input)

定义前向传播的计算过程。

state_dict(destination=None, prefix='', keep_vars=False)

返回包含模块完整状态的字典，字典中同时包含参数和持续性缓存。

load_state_dict(state_dict, strict=True)

从状态字典中将参数和缓存复制到该模块及其子模块中；如果 strict=True , 那么 state_dict 的键必须和该模块的 state_dict() 函数返回的键一致。

parameters(recurse=True)

返回模块参数的迭代器，通常传递给优化器。

named_parameters(prefix='', recurse=True)

返回模块参数的迭代器，迭代器将同时产生参数的名称和参数自身。

register_parameter(name, param)

向模块中加入参数。

to(*args, **kwargs)

移动或强制转换参数和缓存。

type(dst_type)

将所有的参数和缓存强制转换为 dst_type（python:type or string）类型。

zero_grad()

将所有模型参数的梯度设置为 0。

4.5. 参考资料

  1. pytorch documentation

  1. List of nn.Module in a nn.Module

  1. 模型构造

  1. TORCH.NN

  1. 看pytorch文档学深度学习——Containers