小土堆pytorch教程学习笔记P29

P29.完整的模型训练套路（三）

• training start

tudui.train()

train(mode=True)

Sets the module in training mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected, e.g. Dropout, BatchNorm, etc.

• testing start

tudui.eval()

eval()

Sets the module in evaluation mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected, e.g. Dropout, BatchNorm, etc.

• 步骤

准备数据集

利用DataLoader加载数据集

创建网络模型

损失函数

优化器

设置训练网络的一些参数（训练次数、测试次数、训练轮数、tensorboard）

训练步骤开始（tudui.train()，从dataloader中取数据，计算误差，优化器优化模型，展示输出）

测试步骤开始（tudui.eval()，with torch.no_grad()只需要测试，不需要使用梯度调整，从dataloader中取数据，计算误差，设置并展示指标，保存模型）

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))
    # torch.save(tudui.state_dict(), "tudui_{}.pth".format(i))
    print("model saved")

writer.close()