用 nn.AdaptiveAvgPool2d() 替换 LeNet 模型中的全连接层
将 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 著作权归作者所有。请勿转载和采集!