Pytorch 实现 LeNet - cifar10 网络

供个人学习使用，参考 Pytorch 官网教程实现。

源码

# LeNet 
import torch
import torch.utils.tensorboard
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from torchvision import datasets
from torchvision.transforms import ToTensor, Lambda


training_data = datasets.CIFAR10(
    root="cifar10",
    download=True,
    train=True,
    transform=ToTensor()
)

test_data = datasets.CIFAR10(
    root="cifar10",
    train=False,
    download=True,
    transform=ToTensor()
)


train_dataloader = DataLoader(training_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)


class LeNet(nn.Module):
    def __init__(self): # 子类初始化
        super(LeNet, self).__init__() # 父类初始化
        self.LeNet = nn.Sequential(
            nn.Conv2d(3, 6, 5),
            nn.MaxPool2d(2),
            nn.Conv2d(6, 16, 5),
            nn.MaxPool2d(2),
            nn.Flatten(),  # 16*5*5
            nn.Linear(16*5*5, 120),
            nn.Linear(120, 84),
            nn.Linear(84, 10)
        )

    def forward(self, x):
        x = self.LeNet(x)
        return x


model = LeNet()
model = model.cuda() # 将网络模型加入到cuda

loss_fn = nn.CrossEntropyLoss()
loss_fn = loss_fn.cuda()
learning_rate = 1e-2  # 1 x (10)(-2)/100 = 0.01
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)


# 设置训练网络的一些参数
total_train_step = 0 # 记录训练的次数
epoch = 10 # 训练的轮数


# model.train() 用到 Dropout 和 BatchNormal 需要写， 看官网介绍
for i in range(epoch):
    print(f'-------------第 {i+1} 轮训练开始:------------')
    for data in train_dataloader:
        imgs, targets = data
        
        # 将输入和输出加入到cuda
        imgs = imgs.cuda() 
        targets = targets.cuda()
        
        outputs = model(imgs) # 将输入送入网络得出预测结果
        loss = loss_fn(outputs, targets) # 计算预测和标签的loss，这个loss值就是交叉熵的结果，交叉熵越小，结果越好

        # 优化器优化模型
        optimizer.zero_grad() # 梯度清零
        loss.backward() # 损失反向传播，用梯度下降法，求解损失函数最小值。
        optimizer.step() # 执行单个优化步骤

        total_train_step += 1
        if total_train_step % 100 == 0: # 每训练100次，打印一次

            print(f'训练次数：{total_train_step}, Loss: {loss.item()}') # 打印训练次数和当前损失

    test_loss = 0  # 损失总数
    test_loss_rate = 0 # 损失率
    test_current = 0 # 正确的个数
    test_accuracy_rate = 0 # 正确率
    test_data_size = len(test_data)  # 测试数据集大小
    
    with torch.no_grad(): # 测试步骤，保证没有梯度，不会进行优化模型
        for data in test_dataloader:
            imgs, targets = data

            imgs = imgs.cuda() # 图片
            targets = targets.cuda()  # 标签

            outputs = model(imgs) # 预测结果
            loss = loss_fn(outputs, targets) # 计算预测和标签之间的误差
            
            test_loss = test_loss + loss.item() # 累计损失个数
            test_current += (outputs.argmax(1) == targets).type(torch.float).sum().item() # 累计正确个数
            '''
            outputs.argmax(1) 返回该行最大的数的索引
            例如：
            outputs = torch.tensor([[0.1, 0.2],
                       [0.05, 0.4]])
            print(outputs.argmax(1))
            
            输出
            tensor([1, 1])
            '''
            
         
    test_loss_rate = test_loss / test_data_size # 丢失的个数 / 测试集总长 = 损失率
    test_accuracy_rate = test_current / test_data_size # 正确的个数 / 测试集总长 = 正确率
    print(f"Test: \n test_accuracy_rate: {(100*test_accuracy_rate):>0.1f}%, test_loss_rate: {test_loss_rate:>8f} \n")
    
print('训练完成')