pytorch 实现 lenet5 并用于 mnist 分类需要注意哪些问题？

首先，结合下面这个图，使用 pytorch 实现 letnet5 模型。
计算每一层的 channel、kernel_size、stride 数值，得到如下所示：

卷积层	in_channel	out_channel	kernel_size	stride	padding
Conv	1	6	5	1	2
MaxPool	6	6	2	2
Conv	6	16	5	1
MaxPool	16	16	2	2
Conv	16	120	5	1

全连接层	输入	输出	bias
fc1	120	84	True
fc2	84	10	True

然后开始编码，下面是需要注意的点，因为仅仅从上述模型图中看不出来这些 trick，参考 vgg.py

1、每一个卷积层 Conv 后面都需要跟一个激活函数，默认使用ReLu

output = self.conv1(input)
output = self.relu(output)

2、第一个全连接层 Linear 后面，需要跟上一个激活函数，默认使用 ReLu。

output = self.fc1(output)
output = self.relu(output)

3、最后一个全连接层 Linear 后面，不能跟上激活函数 ReLu，因为 ReLu 会将小于0 的数重置为 0。可以跟上 softmax 激活函数，将输出限制为 [0,1] 的概率值。也可以不添加 softmax 激活函数，此时经过交叉熵公式后变为 $loss=y\mathrm{l}\mathrm{n}\hat{y}=y\mathrm{l}\mathrm{n}\max (y)$。参考 vgg.py

4、初始化参数

for m in self.modules():
	if isinstance(m, nn.Conv2d):
	    nn.init.kaiming_normal_(m.weight)
	elif isinstance(m, nn.Linear):
	    nn.init.normal_(m.weight, 0, 0.01)
	    nn.init.constant_(m.bias, 0)

主要是初始化 卷积层 和 全连接层 中的权重参数

经过以上配置后，经过以下步骤训练，可以很快得收敛。

net = Lenet()
optimizer = torch.optim.SGD(net.parameters(), 0.01)
criterion = nn.CrossEntropyLoss()

for epoch in range(0, 10):
    
    epoch_loss = 0.0
    
    for i, (image, target) in enumerate(zip(train_images, train_targets)):
        
        optimizer.zero_grad()
        preds = net(image)
        loss = criterion(preds, target)
        epoch_loss += loss.item()
        
        loss.backward()
        optimizer.step()
        
    print("epoch - {}, loss - {:.4f}".format(epoch, epoch_loss / iter_count / batch_size))