之前对于这一方面了解较少,搭网络就是直接一层一层堆砌,简单粗暴。当然这样做是不对的,写出来的代码不仅难看,集成度不高,迁移困难,而且还容易出错,并且对于一些情况,简单一层一层堆网络是解决不了的。因此,了解一下pytorch container的相关内容还是有必要的。
1. nn.Module
这个是最常用的container,所有其他网络都是这个类的继承。我们在自己定义一个网络或者层时,就需要继承这个类。module允许以树结构进行嵌入,一个module可以包含其他module,这个module就是原有module的submodule。
class MyModule(nn.Module):
def __init__(self):
self.conv1 = nn.Conv2d(16, 32, 3, 1)
self.conv2 = nn.Conv2d(32, 32, 3, 1)
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
return x
上面的例子中,conv1就是MyModule的submodule。当对MyModule的实例进行cuda转换时,conv1作为其submodule,也会被转换为cuda数据格式,并且,这一过程是递归进行的。这很重要,需要认真理解。如果不是nn.Module的实例,就不会被加入到计算图中,也不会被转换为cuda格式。今天我就趟了一个这样的坑,模型怎么都训不出来,最后发现计算图里根本没有层。
-
add_module
除了上面的做法,也可以用add_module添加一个层到网络里,这样做的好处是可以给层命名,这样就可以直接通过层名来找到一个层了。
class MyModule(nn.Module):
def __init__(self):
self.add_module('conv1', nn.Conv2d(16, 32, 3, 1))
self.add_module('conv2', nn.Conv2d(32, 32, 3, 1))
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
return x
m = MyModule()
m1 = m.conv1
2. nn.Sequential
这个在网络中的出现次数也比较频繁,通过这个模块,可以将代码写得更密集一点,可读性也更强。sequential是一个时序模型,根据每个submodule传入的顺序写到计算图里,在forward的时候也会顺序执行。
conv1 = nn.Sequential(
nn.Conv2d(32, 64, 3, 1),
nn.BatchNorm2d(64),
nn.ReLU()
)
当然,也可以传一个OrderedDict来构造网络。
3. nn.ModuleList
如果不想让module按照传入顺序执行,就可以将它们写成一个list,用下标来进行索引。但是哪怕是直接在init函数里定义为成员变量,最后也不会被加到计算图里,真的是很心烦。我今天就跳了一个这样的坑,正打算自己亲自写一个类时,发现pytorch已经帮我们实现好了,就是ModuleList。官方大法好~~
ModuleList也是继承了Module的一个子类,可以像python list一样用下标索引,可以使用append和extend方法,最重要的是,也会被加到计算图里,总之能用就是了。
MyModule = nn.ModuleList()
MyModule.append(nn.Conv2d(32, 64, 3, 1))
MyModule.append(nn.Conv2d(32, 64, 1, 1))
m = MyModule()
x1 = m[0](x)
x2 = m[1](x)