pytorch nn.Sequential()动态添加方法

之前我们使用nn.Sequential()都是直接写死的，就如下所示：

# Example of using Sequential
model = nn.Sequential(
          nn.Conv2d(1,20,5),
          nn.ReLU(),
          nn.Conv2d(20,64,5),
          nn.ReLU()
        )

# Example of using Sequential with OrderedDict
model = nn.Sequential(OrderedDict([
          ('conv1', nn.Conv2d(1,20,5)),
          ('relu1', nn.ReLU()),
          ('conv2', nn.Conv2d(20,64,5)),
          ('relu2', nn.ReLU())
        ]))

那如果我们想要根据条件一点点添加进去，那就可以使用其的add_module方法

 

torch.nn.Module.add_module

add_module(name, module)

添加子模块到当前模块中

该添加子模块能够使用给定的名字name来访问

参数：

• name (string)：子模块的名字。该添加子模块能够使用给定的名字name来从该模块中被访问
• module (Module) ：添加到该模块中的子模块

 

例子：

class Encoder(nn.Module):  #输入图片的大小isize、噪声的维度nz=100、输入图片的通道nc=3、ndf=64、
    def __init__(self,isize,nz,nc,ndf,ngpu,n_exter_layers=0,add_final_conv=True):
        super(Encoder,self).__init__()
        self.ngpu=ngpu
        # 必须为16倍数
        assert isize % 16==0,"isize has to be a multiple of 16"

        main=nn.Sequential()
        # 图片的高宽缩小一倍
        main.add_module('initial-conv-{0}-{1}'.format(nc,ndf),nn.Conv2d(nc,ndf,4,2,1,bias=False))
        main.add_module('initial-relu-{0}'.format(ndf),nn.LeakyReLU(0.2,inplace=True))
        csize,cndf=isize/2,ndf

        for t in range(n_exter_layers): #在这里面特征宽高不变，通道数也不变
            main.add_module('extra-layers-{0}-{1}-conv'.format(t,cndf),nn.Conv2d(cndf,cndf,3,1,1,bias=False))
            main.add_module('extra-layers-{0}-{1}-batchnorm'.format(t,cndf),nn.BatchNorm2d(cndf))
            main.add_module('extra-layers-{0}-{1}-relu'.format(t,cndf),nn.LeakyReLU(0.2,inplace=True))

        # 在特征高宽仍大于4时，就添加缩小一倍高宽，通道增加一倍的卷积块
        while csize>4:
            in_feat = cndf

            out_feat = cndf * 2

            main.add_module('pyramid-{0}-{1}-conv'.format(in_feat, out_feat),nn.Conv2d(in_feat, out_feat, 4, 2, 1, bias=False))

            main.add_module('pyramid-{0}-batchnorm'.format(out_feat),nn.BatchNorm2d(out_feat))

            main.add_module('pyramid-{0}-relu'.format(out_feat),nn.LeakyReLU(0.2, inplace=True))

            cndf = cndf * 2

            csize = csize / 2

        # 最后一层卷积，将4*4变为1*1,得到nz = 100的噪声
        if add_final_conv:

            main.add_module('final-{0}-{1}-conv'.format(cndf, 1),nn.Conv2d(cndf, nz, 4, 1, 0, bias=False))
            self.main=main
    
    def forward(self,input):
        if self.ngpu>1:
             output=nn.parallel.data_parallel(self.main,input,range(self.ngpu)) #在多个gpu上运行模型，并行计算
        else:
            output=self.main(input)
        
        return output  #如果输入的大小是3×32×32，最后的输出是100×1×1.