PyTorch使用CIFAR-10数据进行训练

场景

最近在PyTorch中使用CIFAR-10数据进行模型训练。

main.py

import time

import torch
import torchvision
from torch import nn
from torch.utils.data import DataLoader
from torchvision import transforms
from PIL import Image

from MyCNN import MyCNN


def cifar10_go():
    transform = transforms.Compose([
        transforms.RandomResizedCrop((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    cifar10_dataset = torchvision.datasets.CIFAR10(root='./data',
                                                   train=False,
                                                   transform=transform,
                                                   target_transform=None,
                                                   download=True)
    dataloader = DataLoader(dataset=cifar10_dataset,  # 传入的数据集, 必须参数
                            batch_size=32,  # 输出的batch大小
                            shuffle=True,  # 数据是否打乱
                            num_workers=4)  # 进程数, 0表示只有主进程
    model = MyCNN()
    # 交叉熵损失函数
    criterion = nn.CrossEntropyLoss()
    # 定义优化器
    optimizer = torch.optim.SGD(model.parameters(), lr=1e-4, weight_decay=1e-2, momentum=0.9)
    # 下面开始训练
    start = time.time()  # 记时开始
    for epoch in range(3):  # 设置要在全部数据上训练的次数

        for i, data in enumerate(dataloader):
            # data就是我们获取的一个batch size大小的数据

            inputs, labels = data  # 分别得到输入的数据及其对应的类别结果
            # 首先要通过zero_grad()函数把梯度清零，不然PyTorch每次计算梯度会累加，不清零的话第二次算的梯度等于第一次加第二次的
            optimizer.zero_grad()
            # 获得模型的输出结果，也即是当前模型学到的效果
            outputs = model(inputs)
            # 获得输出结果和数据真正类别的损失函数
            loss = criterion(outputs, labels)
            print('Epoch {}, Loss {}'.format(epoch + 1, loss))
            # 算完loss之后进行反向梯度传播，这个过程之后梯度会记录在变量中
            loss.backward()
            # 用计算的梯度去做优化
            optimizer.step()
    end = time.time()  # 计时结束
    print('使用时间: {:.5f} s'.format(end - start))
    # 保存模型训练结果
    torch.save(model, './MyCNN_model.pth')
def test():
    # 测试一波
    mycnn_model_2 = torch.load('./MyCNN_model.pth')
    mycnn_model_2.eval()
    for parameter in mycnn_model_2.named_parameters():
        print(parameter)
    # 开始测试
    im = Image.open('dog.jpg')
    transform = transforms.Compose([transforms.RandomResizedCrop((224, 224)), transforms.ToTensor(),
                                    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
    input_tensor = transform(im).unsqueeze(0)
    print(mycnn_model_2(input_tensor).argmax())

# Press the green button in the gutter to run the script.
if __name__ == '__main__':
    cifar10_go()
    test()

MyCNN.py

from torch import nn
class MyCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3)
        # conv1输出的特征图为222x222大小
        self.fc = nn.Linear(16 * 222 * 222, 10)

    def forward(self, input):
        x = self.conv1(input)
        # 进去全连接层之前，先将特征图铺平
        x = x.view(x.shape[0], -1)
        x = self.fc(x)
        return x

结果

tensor(5)

总结

训练模型主要有如下步骤：

• 1.加载数据；
• 2.确定和使用模型；
• 3.确定和使用损失函数，优化方法；
• 4.训练模型
• 5.测试效果
  效果是弄出来了，但是对于损失函数选取，以及优化参数配置也是一脸懵。

参考：

• The CIFAR-10 dataset
• 14 | 构建网络：一站式实现模型搭建与训练