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.