完整的模型训练套路

# 准备数据集
import torch
from torch import nn
from torch.nn.modules.flatten import Flatten

# Building the neural network
class Cifar(nn.Module):

    def __init__(self) -> None:
        super(Cifar, self).__init__()
        self.model = nn.Sequential(
            nn.Conv2d(3, 32, 5, 1, 2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 32, 5, 1, 2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 5, 1, 2),
            nn.MaxPool2d(2),
            Flatten(),
            nn.Linear(1024, 64),
            nn.Linear(64, 10)
        )

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

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

# 准备数据集
import torchvision
from torch import nn
from torch.nn.modules.flatten import Flatten
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

from model import *

train_data = torchvision.datasets.CIFAR10(root="../datasets", train=True, transform=torchvision.transforms.ToTensor(),
                                          download=True)
test_data = torchvision.datasets.CIFAR10(root="../datasets", 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))

# using dataloader load the dataset
train_dataloader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)

# build network model
model = Cifar()

# build loss
loss_fn = nn.CrossEntropyLoss()

# optimizer
learnRate = 0.01
optimizer = torch.optim.SGD(model.parameters(), lr=learnRate, )

# Setting train network's parameter
# Recording training sessions
totalTrainStep = 0

# Recording testing sessions
totalTestStep = 0

# Number of training rounds
epoch = 10

# add tensorboard
writer = SummaryWriter("trainmodel")

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

    # train start
    model.train()
    for data in train_dataloader:
        imgs, targets = data
        output = model(imgs)
        loss = loss_fn(output, targets)
        # The optimizer optimizes the model
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        totalTrainStep = totalTrainStep + 1
        if totalTrainStep % 100 == 0:
            print("训练次数：{}, Loss:{}".format(totalTrainStep, loss.item()))
            writer.add_scalar("trainLoss", loss.item(), totalTrainStep)

    # Test start
    model.eval()

    # Test step start
    totalAccuracy = 0
    totalTestLoss = 0

    # test needn't adjust grad
    with torch.no_grad():
        for data in test_dataloader:
            imgs, targets = data
            output = model(imgs)
            loss = loss_fn(output, targets)
            totalTestLoss = totalTestStep + loss
            accuracy = (output.argmax(1) == targets).sum()
            totalAccuracy = totalAccuracy + accuracy
    print("整体测试集上的Loss：{}".format(totalTestLoss.item()))
    print("整体测试集上的正确率：{}".format(totalAccuracy/test_data_size))
    writer.add_scalar("testLoss", totalTestLoss, totalTestStep)
    writer.add_scalar("testAccuracy", totalAccuracy/test_data_size, totalTestStep)
    totalTestStep = totalTestStep + 1

    # Recording model result
    torch.save(model, "theRounds{}.pth".format(i))
    torch.save(model.state_dict(), "model{}.pth".format(i))

writer.close()