零基础学Pytorch02_图像分类任务_模型的数据导入以及训练

模型的数据导入以及训练

train.py
此是GPU版本，如果是CPU版本则删除相应的转入GPU部分的代码即可

import numpy as np
import torch
import torchvision
import torch.nn as nn
from matplotlib import pyplot as plt

from LetNet import LetNet
import torch.optim as optim
import torchvision.transforms as transforms

# 图像均值化，中心化，去除干扰因素
transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

# 50000张训练图片
# 第一次使用时要将download设置为True才会自动去下载数据集
train_set = torchvision.datasets.CIFAR10(root='./data', train=True,
                                         download=True, transform=transform)
# 这四个参数分别代表导入的训练集 每批训练的样本数 是否打乱训练集 线程数
train_loader = torch.utils.data.DataLoader(train_set, batch_size=36,
                                           shuffle=True, num_workers=0)

# 10000张验证图片
# 第一次使用时要将download设置为True才会自动去下载数据集
val_set = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=False, transform=transform)
val_loader = torch.utils.data.DataLoader(val_set, batch_size=10000,
                                         shuffle=False, num_workers=0)
val_data_iter = iter(val_loader)


val_image, val_label = val_data_iter.next()


classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

# def imshow(img):
#     img = img / 2 + 0.5
#     npimg = img.numpy()  # 将图像反标准化
#     plt.imshow(np.transpose(npimg, (1, 2, 0)))
#     plt.show()
#
# 需要显示图片的话，即把 batch_size=10000中的10000改为4，因为10000太大了，显示不出来
# print(' '.join('%5s' % classes[val_label[j]] for j in range(4)))
# imshow(torchvision.utils.make_grid(val_image))

net = LetNet()  # 实例化

#如果CUDA可以用，让我们首先定义下我们的设备为第一个可见的cuda设备。
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)
#将神经网络转到GPU上
net.to(device)


# 定义一个损失函数和优化器 让我们使用分类交叉熵Cross-Entropy 作损失函数，动量SGD做优化器。
# Define loss (Cross-Entropy)
loss_function = nn.CrossEntropyLoss()  # 定义损失函数

optimizer = optim.Adam(net.parameters(), lr=0.001)  # 优化器 第一个参数为需要训练的参数，lr为学习率

# 将训练集迭代五轮
for epoch in range(5):  # loop over the dataset multiple times

    running_loss = 0.0 # 累加训练过程中的损失
    for step, data in enumerate(train_loader, start=0):
        # get the inputs; data is a list of [inputs, labels]
        inputs, labels = data

        # 必须在每一个步骤向GPU发送输入和目标
        inputs, labels = inputs.to(device), labels.to(device)

        # zero the parameter gradients
        """
        为什么每计算一个batch，就需要调用一次optimizer.zero_grad()？
        因为如果不清楚历史梯度，就会对计算的历史梯度进行累加（通过这个特性你能够变相实现一个很大的batch数值的训练
        """
        optimizer.zero_grad() # 将历史损失梯度清零
        # forward + backward + optimize
        outputs = net(inputs)
        loss = loss_function(outputs, labels)
        loss.backward() # 反向传播
        optimizer.step()# 参数更新

        # print statistics
        running_loss += loss.item()

        if step % 500 == 499:  # print every 500 mini-batches
            with torch.no_grad(): # with是一个上下文管理器
            
                # 将数据转到GPU上
                val_image = val_image.cuda()
                val_label = val_label.cuda()
                images, labels = val_image.to(device), val_label.to(device)

                outputs = net(val_image)  # [batch, 10]
                predict_y = torch.max(outputs, dim=1)[1]
                accuracy = torch.eq(predict_y, val_label).sum().item() / val_label.size(0)

                print('[%d, %5d] train_loss: %.3f  test_accuracy: %.3f' %
                      (epoch + 1, step + 1, running_loss / 500, accuracy))
                running_loss = 0.0

print('Finished Training')


# 模型数据的保存，pth是pytorch的缩写
save_path = './Lenet.pth'
torch.save(net.state_dict(), save_path)