Pytorch框架学习记录12——完整的模型训练套路

本次模型的主要使用CIFIAR10数据集，搭建了CIFIAR 10模型。

首先，需要对数据集进行下载读取，并进行分组。

# 读取数据集
trainset = torchvision.datasets.CIFAR10("dataset", train=True, transform=torchvision.transforms.ToTensor(), download=True)
testset = torchvision.datasets.CIFAR10("dataset", train=False, transform=torchvision.transforms.ToTensor(), download=True)

# 数据分组
train_loader = DataLoader(trainset, 64)
test_loader = DataLoader(testset, 64)

搭建模型，并在train.py文件中导入已构建好的模型，定义优化器、损失函数和相关的一些变量。

model.py

import torch
from torch import nn


class Test(nn.Module):
    def __init__(self):
        super(Test, self).__init__()
        self.model = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=32, kernel_size=5, padding=2),
            nn.MaxPool2d(kernel_size=2),
            nn.Conv2d(in_channels=32, out_channels=32, kernel_size=5, padding=2),
            nn.MaxPool2d(kernel_size=2),
            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=5, padding=2),
            nn.MaxPool2d(kernel_size=2),
            nn.Flatten(),
            nn.Linear(in_features=1024, out_features=64),
            nn.Linear(in_features=64, out_features=10)
        )


    def forward(self, input):
        output = self.model(input)
        return output


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

train.py文件

# 创建网络模型
test = Test()

# 定义损失函数
loss_fn = torch.nn.CrossEntropyLoss()

# 定义优化器
learning_rate = 0.001
optimizer = torch.optim.SGD(params=test.parameters(), lr=learning_rate)

# 记录训练次数
train_step = 0

# 记录测试次数
test_step = 0

writer = SummaryWriter("logs")

接下来，对模型进行训练，

# 训练
epoch = 20
for i in range(epoch):

    # 训练步骤
    print("-------第 {} 轮训练-------".format(i+1))

    test.train()
    for train_data in train_loader:
        imgs, target = train_data
        output = test(imgs)
        loss = loss_fn(output, target)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        train_step += 1
        if train_step % 100 == 0:
            print("第 {} 次训练完成 训练损失：{}".format(train_step, loss.item()))
            writer.add_scalar("train_loss", loss.item(), train_step)

在测试集上进行测试，

# 测试步骤
test.eval()

test_loss_sum = 0.0
total_accuracy = 0
with torch.no_grad():
    for test_data in test_loader:
        imgs, target = test_data
        output = test(imgs)
        loss = loss_fn(output, target)

        accuracy = (output.argmax(1) == target).sum()

        test_loss_sum += loss.item()
        total_accuracy += accuracy

writer.add_scalar("test_loss", test_loss_sum, test_step)
print("在测试集上的Loss：{}, 正确率：{}".format(test_loss_sum, total_accuracy/len(testset)))
test_step += 1

【完整源代码】

import torchvision
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from Pytorch_Learning.model import *


# 读取数据集
trainset = torchvision.datasets.CIFAR10("dataset", train=True, transform=torchvision.transforms.ToTensor(), download=True)
testset = torchvision.datasets.CIFAR10("dataset", train=False, transform=torchvision.transforms.ToTensor(), download=True)

# 数据分组
train_loader = DataLoader(trainset, 64)
test_loader = DataLoader(testset, 64)

# 创建网络模型
test = Test()

# 定义损失函数
loss_fn = torch.nn.CrossEntropyLoss()

# 定义优化器
learning_rate = 0.001
optimizer = torch.optim.SGD(params=test.parameters(), lr=learning_rate)

# 记录训练次数
train_step = 0

# 记录测试次数
test_step = 0

writer = SummaryWriter("logs")

# 训练
epoch = 20
for i in range(epoch):

    # 训练步骤
    print("-------第 {} 轮训练-------".format(i+1))

    test.train()
    for train_data in train_loader:
        imgs, target = train_data
        output = test(imgs)
        loss = loss_fn(output, target)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        train_step += 1
        if train_step % 100 == 0:
            print("第 {} 次训练完成 训练损失：{}".format(train_step, loss.item()))
            writer.add_scalar("train_loss", loss.item(), train_step)

    # 测试步骤
    test.eval()

    test_loss_sum = 0.0
    total_accuracy = 0
    with torch.no_grad():
        for test_data in test_loader:
            imgs, target = test_data
            output = test(imgs)
            loss = loss_fn(output, target)

            accuracy = (output.argmax(1) == target).sum()

            test_loss_sum += loss.item()
            total_accuracy += accuracy

    writer.add_scalar("test_loss", test_loss_sum, test_step)
    print("在测试集上的Loss：{}, 正确率：{}".format(test_loss_sum, total_accuracy/len(testset)))
    test_step += 1