小土堆pytorch教程学习笔记P28

原创已于 2022-04-18 09:17:56 修改 · 159 阅读

1 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch

于 2022-04-05 20:23:26 首次发布

小土堆pytorch教程学习笔记专栏收录该内容

22 篇文章

订阅专栏

本文详细介绍了使用PyTorch库在CIFAR10数据集上进行深度学习模型（Tudui）的完整训练流程，包括数据预处理、DataLoader加载、模型构建、损失函数和优化器的选择，以及TensorBoard的实时监控。通过多个训练轮次，展示了训练和验证过程中的损失和精度变化。

P28.完整的模型训练套路（二）

import torchvision
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from P27_model import *

# prepare dataset
train_data = torchvision.datasets.CIFAR10(root="dataset", train=True, transform=torchvision.transforms.ToTensor(),
                                          download=True)
test_data = torchvision.datasets.CIFAR10(root="dataset", train=False, transform=torchvision.transforms.ToTensor(),
                                         download=True)

# length
train_data_size = len(train_data)
test_data_size = len(test_data)
# if train_data_size=10,The length of train dataset is 10
print("The length of train dataset is:{}".format(train_data_size))
print("The length of test dataset is:{}".format(test_data_size))

# load dataset with DataLoder
train_dataloader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)

# build a network model
tudui = Tudui()

# loss function
loss_fn = nn.CrossEntropyLoss()

# optimizer
# learning_rate = 0.01
# 1e-2 = 1 x (10)^(-2) = 1/100 = 0.01
learning_rate = 1e-2
optimizer = torch.optim.SGD(tudui.parameters(), lr=learning_rate)

# set some parameters of training network
# record the number of training
total_train_step = 0
# record the number of testing
total_test_step = 0
# number of training rounds
epoch = 10

# add tensorboard
writer = SummaryWriter("logs_train")

for i in range(epoch):
    print("------The {} round training start------".format(i+1))

    # training start
    for data in train_dataloader:
        imgs, targets = data
        outputs = tudui(imgs)
        loss = loss_fn(outputs, targets)

        # use optimizer to optimize the model
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_train_step = total_train_step + 1
        if total_train_step % 100 == 0:
            print("training round:{}, Loss:{}".format(total_train_step, loss.item()))
            writer.add_scalar("train_loss", loss.item(), total_train_step)

    # testing start
    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 = total_test_loss + loss.item()
            accuracy = (outputs.argmax(1) == targets).sum()
            total_accuracy = total_accuracy + accuracy
    print("The total loss on testing:{}".format(total_test_loss))
    print("The accuracy of testing:{}".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 = total_test_step + 1

    torch.save(tudui, "tudui_{}.pth".format(i))
    print("model saved")

writer.close()