将 LeNet 模型中的全连接层替换为 nn.AdaptiveAvgPool2d() 可以按照以下方式实现:

class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(6),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )

        self.conv2 = nn.Sequential(
            nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5, stride=1),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )

        self.conv3 = nn.Sequential(
            nn.Conv2d(in_channels=16, out_channels=120, kernel_size=5),
            nn.BatchNorm2d(120),
            nn.ReLU()
        )

        self.avg_pool = nn.AdaptiveAvgPool2d((1, 1))  # 替换全连接层

        self.fc = nn.Linear(120, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = self.avg_pool(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

这里使用 nn.AdaptiveAvgPool2d((1, 1)) 将最后一个卷积层的输出转换为 1x1 的特征图,然后使用 x.view(x.size(0), -1) 将特征图展平,最后使用 nn.Linear(120, 10) 作为全连接层替换原有的两个全连接层。


原文地址: https://www.cveoy.top/t/topic/nC3r 著作权归作者所有。请勿转载和采集!

免费AI点我,无需注册和登录