P17：完整的模型训练套路（以CIFAR10数据集为例）

1、导包：

新建一个 train.py 文件：

import torch
import torchvision
from torch import nn
from torch.nn import Conv2d, MaxPool2d, Flatten, Linear, ReLU, CrossEntropyLoss
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

from model import *

2、准备数据集：

# 1、准备数据集
train_data = torchvision.datasets.CIFAR10('dataset', train = True, transform=torchvision.transforms.ToTensor(), download=True)
test_data = torchvision.datasets.CIFAR10('dataset', train = False, transform=torchvision.transforms.ToTensor(), download=True)

train_data_size = len(train_data)
test_data_size = len(test_data)
print('训练数据集的长度为：{}'.format(train_data_size))
print('测试数据集的长度为：{}'.format(test_data_size))

3、利用Dataloader来加载数据集：

# 2、利用 DataLoader 来加载数据集
train_dataloader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)

4、搭建神经网络：

新建一个 model.py 文件：

import torch
from torch import nn
from torch.nn import Conv2d, MaxPool2d, Flatten, Linear, ReLU


class Tudui(nn.Module):
    def __init__(self):
        super(Tudui, self).__init__()
        self.model = nn.Sequential(
            Conv2d(3, 32, kernel_size=5, stride=1, padding=2),
            MaxPool2d(kernel_size=2),
            Conv2d(32, 32, kernel_size=5, stride=1, padding=2),
            MaxPool2d(kernel_size=2),
            Conv2d(32, 64, kernel_size=5, stride=1, padding=2),
            MaxPool2d(kernel_size=2),
            Flatten(),
            Linear(1024, 64),
            ReLU(),
            Linear(64, 10),
        )

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

if __name__ == '__main__':
    tudui = Tudui()
    input = torch.ones((64,3,32,32))
    output = tudui(input)
    print(output.shape)

5、创建网络模型：

在 train.py 文件中：

# 4、创建网络模型
tudui = Tudui()

6、创建损失函数：

# 5、创建损失函数
loss_fn = CrossEntropyLoss()

7、创建优化器：

# 6、优化器
# learning_rate = 0.01
learning_rata = 1e-2
optimizer = torch.optim.SGD(tudui.parameters(), lr = learning_rata)

8、设置训练网络的一些参数：

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

9、添加tensorboard：

# 添加tensorboard
writer = SummaryWriter('logs_train')

10、训练网络：

for i in range(epoch):
    print('-------第{}轮训练开始-------'.format(i+1))

    # 训练步骤开始
    tudui.train()
    for data in train_dataloader:
        imgs, targets = data
        outputs = tudui(imgs)
        loss = loss_fn(outputs, targets)

        # 优化器优化模型
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_train_step += 1
        if total_train_step % 100 == 0:
            print('训练次数: {}，loss: {}'.format(total_train_step, loss.item())) # item表示只是数字
            writer.add_scalar('train_loss', loss.item(), total_train_step)

    # 测试步骤开始
    tudui.eval()
    total_test_loss = 0
    total_accuracy = 0
    with torch.no_grad():
        for data in test_dataloader:
            imgs, targets = data
            outputs = tudui(imgs)
            loss = loss_fn(outputs, targets)
            total_test_loss += loss.item()
            accuracy = (outputs.argmax(1) == targets).sum()
            total_accuracy += accuracy
    print('整体测试集上的loss: {}'.format(total_test_loss))
    print('整体测试集上的accuracy: {}'.format((total_accuracy / test_data_size)))
    writer.add_scalar('test_loss', total_test_loss, total_test_step)
    writer.add_scalar('test_accuracy', (total_accuracy / test_data_size), total_test_step)
    total_test_step += 1

    # 保存每一轮训练的模型
    torch.save(tudui, 'tudui_{}.pth'.format(i))
    print('模型已保存')

writer.close()