Pytorch入门教程（九）：CIFAR图片分类实战

1. 设计网络结构：

采用最原始的Lenrt5:

2. 代码编写：

import torch
from torch import nn


class Lenet5(nn.Module):

    def __init__(self):
        super(Lenet5, self).__init__()

        # 卷积单元
        self.conv_unit = nn.Sequential(
            # x:[b, 3, 32, 32]
            nn.Conv2d(3, 6, kernel_size=5, stride=1, padding=0),
            nn.MaxPool2d(kernel_size=2, stride=2, padding=0),
            nn.Conv2d(6, 16, kernel_size=5, stride=1, padding=0),
            nn.MaxPool2d(kernel_size=2, stride=2, padding=0),
        )
        # 全连接单元
        self.fc_unit = nn.Sequential(
            nn.Linear(16*5*5, 120),  # 全连接之前有一个打平操作，所以输入为三者之积
            nn.ReLU(),
            nn.Linear(120, 84),
            nn.ReLU(),
            nn.Linear(84, 10)
        )

    def forward(self, x):
        # [b, 3, 32, 32] => [b, 16, 5, 5]
        batchsz = x.size(0)
        x = self.conv_unit(x)
        # [b, 16, 5, 5] => [b, 16*5*5]
        x = x.view(batchsz, -1)
        # [b, 16*5*5] => [b, 10]
        logits = self.fc_unit(x)
        return logits

3.读取数据，进行训练，测试

import torch
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision import transforms
from lenet5 import Lenet5
from torch import nn, optim


def main():
    batchsz = 32
    # 读取数据，位置，是否为训练集，数据转化
    cifar_train = datasets.CIFAR10("cifar", True,  transform=transforms.Compose([
        transforms.Resize((32, 32)),
        transforms.ToTensor()]), download=True)
    # 批次读取, 原数据，批次大小，是否乱序
    cifar_train = DataLoader(cifar_train, batch_size=batchsz, shuffle=True)

    cifar_test = datasets.CIFAR10("cifar", False,  transform=transforms.Compose([
        transforms.Resize((32, 32)),
        transforms.ToTensor()]), download=True)
    cifar_test = DataLoader(cifar_test, batch_size=batchsz, shuffle=True)

    x, label = iter(cifar_train).next()  # 迭代器，一条一条的获取数据
    print("x, label", x.shape, label.shape)

    device = torch.device("cuda")  # 设置加速器
    model = Lenet5().to(device)  # 在载入模型
    print(model)
    criteon = nn.CrossEntropyLoss().to(device)  # 计算交叉熵
    optimizer = optim.Adam(model.parameters(), lr=1e-3)  # 设置优化器
    model.train()
    for epoch in range(1000):
        for batchid, (x, label) in enumerate(cifar_train):
            # x:[b, 3, 32, 32]
            # label:[b]
            x, label = x.to(device), label.to(device)  # 转化到cuda上
            # logits:[b, 10]
            # label:[b]
            logits = model(x)  # 计算输出
            loss = criteon(logits, label)  # 计算损失
            # backprob
            optimizer.zero_grad()  # 梯度清零
            loss.backward()  # loss更新梯度
            optimizer.step()  # 前向传播，更新权重

        print(epoch, loss.item())

        model.eval()
        with torch.no_grad():
            total_crrect = 0
            total_num = 0
            for x, label in cifar_test:
                x, label = x.to(device), label.to(device)
                logits = model(x)
                pred = logits.argmax(dim=1)  # 返回最大值所在位置
                total_crrect += torch.eq(pred, label).float().sum().item()  # 判断与label是否相等，返回正确个数
                total_num += x.size(0)
            acc = total_crrect / total_num  # 计算正确率
            print(epoch, acc)

if __name__ == '__main__':
    main()