pytorch 模型容器 - 模块化构建深度学习网络
作者：互联网
pytorch 模型容器 - 模块化构建深度学习网络
目录
pytorch 模型容器总结
pytorch 提供的模型容器包括：
nn.Sequential:按顺序包含多个网络层
nn.ModuleList:类似列表(list)的形式包含多个网络层
nn.ModuleDict:类似字典(dict)的形式包含多个网络层
一、nn.Sequential
nn.Sequential中的模块按顺序排列，输入数据从第一个block输入，依次经过后续的block后得到输出结果
代码示例：
class Net(nn.Module):
def __init__(self):
super(CoarseNet_v1, self).__init__()
self.seq1=nn.Sequential(
nn.Conv2d(16,32,3,1),
Block1(32, 32),
nn.ReLU(inplace=True)
)
# Block1为自定义模块
def forward(self,inputs)
output=self.seq1(inputs) # 这样用一行代码就可以实现网络的前向传输
return output
这样在定义前向传播时要便捷很多，但是 nn.Sequential 内部模块执行的顺序是固定的，这是应用时存在的不便之处。
二、nn.ModuleList
nn.ModuleList中的模块以列表的形式进行存储，在使用时，可以采用列表索引的形式实现模块的引用
代码示例：
class NetList(nn.Module):
def __init__(self):
super(NetList, self).__init__()
self.list1=nn.ModuleList(
[
nn.Conv2d(16, 32,3,1),
Block1(32, 32),
nn.Conv2d(32, 64, 3, 1),
nn.Conv2d(96, 64, 3, 1),
Block2(64, 32),
]
)
# Block1 和 Block2 是自定义模块
def forward(self,inputs):
x= self.list1[0](inputs) # 这里引用了nn.Conv2d(16,32,3,1)
x1= self.list1[1](x) # 这里引用了Block1(32, 64)
x2= self.list1[2](x) # 这里引用了 nn.Conv2d(32, 64, 3, 1)
x3=torch.cat((x1,x2),1)
x4=self.list1[3](x3) # 这里引用了 nn.Conv2d(96, 64, 3, 1)
x5=self.list1[4](x4) # 这里引用了 Block2(64, 32)
使用ModuleList对模块进行包装，可以实现List内模块的索引，灵活性要好一些。但是也存在一定的问题，但模块数增加的时候，没有对应的注释，程序将会变得不方便阅读，List结构也不利于后续的修改。
三、nn.ModuleDict
nn.ModuleDict中的模块以字典的形式进行存储，在使用时，可以采用字典键值索引的形式实现模块的引用
代码示例：
class NetDict(nn.Module):
def __init__(self):
super(NetDict, self).__init__()
self.net1=nn.ModuleDict({
'block11':Block1(1, 32),
'downsample11':ConvLRelu(32, 32, 2, 1),
'block12':Block2(32, 64),
'downsample12': ConvLRelu(64, 64, 2, 1),
"block13": ConvLRelu(64, 128),
"block21": Block2(1, 32),
"downsample21":ConvLRelu(32, 32, 2, 2),
"block22":Block2(32, 64),
"downsample22":ConvLRelu(64, 64, 2, 2),
'block23':Block2(64, 128),
})
self.net2 = nn.ModuleDict({
'block_r1' : Block1(128, 128),
'deconv1' : DeConvLRelu(128, 128),
'block_r2' : Block1(128, 64),
'deconv2' : DeConvLRelu(64, 32),
'clr' :ConvLRelu(64, 32, 1, 1)
})
# Block1,Block2 为自定义模块
def forward(self, inputs, testflag=0):
# net1 部分
x11 = self.net1['block11'](inputs)
x21 = self.net1["block21"](inputs)
xd11 = self.net1['downsample11'](x11)
xd21 = self.net1['downsample21'](x21)
x12 = self.net1['block12'](xd11)
x22 = self.net1['block22'](xd21)
xd12 = self.net1['downsample12'](x12)
xd22 = self.net1['downsample22'](x22)
x_cat = torch.cat((xd12 , xd22 ), 1)
# net2 部分
xr1 = self.net2['block_r1'](x_cat) # 128
xdc1 = self.net2['deconv1'](x12) # 128
xadd=torch.add(xr1,xdc1 )
xr2 = self.net1['block_r2'](xadd)
xdc2 = self.net1['deconv2'](xr2)
output= self.net1['clr'](xdc2 )
return output
使用ModuleDict对模块进行包装，可以实现Dict内模块的索引，在定义前向时，可以使用键值索引，灵活性好，也使得程序可读性较好。
后记
使用pytorch 模型容器，主要是考虑对整个网络进行其他操作，所以将网络模块包装起来，方便操作。最近运行一个模型较大，考虑采用模型并行，需要将模型拆开到两个GPU上运行，逐个对网络进行GPU位置的设置比较麻烦，使用容器的话便能很方便地实现网络模块GPU位置的设定。
使用多显卡加速网络训练（数据并行和模型并行）可参考：
标签：容器,nn,模块化,self,pytorch,64,模块,32,net1 来源： https://blog.csdn.net/weixin_43537379/article/details/120403454