PyTorch什么时候使用数组形式创建多个卷积层

Posted 2022-08-23 算法与编程之美

问题

PyTorch实现CNN卷积网络的时候，使用数组的方式创建层是一种十分常见的方式，如下所示：

        classifier = []
        classifier.append(nn.Linear(64, 512)) # in_cha
        classifier.append(nn.Linear(512, 2))
        self.classifier = nn.Sequential(*classifier)

然而，什么情况下应该使用这种方式呢？

方法

本文将通过案例详细阐述。

方法1

import torch
from torch import nn


class Net(nn.Module):
    
    def __init__(self) -> None:
        super().__init__()
        
        # 这种情况下,使用数组来构建层毫无意义
        layers = []
        layers.append(nn.Conv2d(3, 32, 3, padding=1))
        layers.append(nn.AdaptiveAvgPool2d(1))
        layers.append(nn.Flatten())
        layers.append(nn.Linear(32, 2))

        self.classifier = nn.Sequential(*layers)
        
    def forward(self, x):
        
        return self.classifier(x)
    

if __name__ == '__main__':
    
    x = torch.rand(size=(1, 3, 224, 224))
    net = Net()
    
    print(net(x).shape)

方法2

import torch
from torch import nn


class Net(nn.Module):
    
    def __init__(self) -> None:
        super().__init__()
        
        self.classifier = nn.Sequential(
            nn.Conv2d(3, 32, 3, padding=1),
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(32, 2)
        )
        
    def forward(self, x):
        
        return self.classifier(x)
    

if __name__ == '__main__':
    
    x = torch.rand(size=(1, 3, 224, 224))
    net = Net()
    
    print(net(x).shape)

方法3

方法1和方法2相比，没有使用数组的方法2更加简洁、清晰。

假如构建10个卷积层，卷积层的输入通道数每次增加5个，应该如何实现呢？

import torch
from torch import nn


class Net(nn.Module):
    
    def __init__(self) -> None:
        super().__init__()

        in_channels = 3
        grow = 5
        layers = []
        
        for i in range(10):
            # 每次增加grow个通道,并保持特征图大小不改变
            layers.append(nn.Conv2d(in_channels, in_channels+grow, 3, padding=1))
            
            in_channels += grow
            
        self.block = nn.Sequential(*layers)
        
        self.classifier = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(in_channels, 2)
        )
        
    def forward(self, x):
        x = self.block(x)
        return self.classifier(x)
    

if __name__ == '__main__':
    from torchsummary import summary
    
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    
    x = torch.rand(size=(1, 3, 224, 224)).to(device)
    net = Net().to(device)
    print(net(x).shape)

    summary(net, (3, 224, 224))

网络输出结果：

结语

当网络层的输入通道数呈现某种规律性的变化时，使用数组创建层的方式将会变得简单，而通常情况下，则无需使用数组，使用Sequential的方式创建会使代码变得更加清晰、易懂。