[pytorch]完整的模型训练套路

model.py

import torch
from torch import nn


class NeuralNode(nn.Module):
    def __init__(self):
        super(NeuralNode, self).__init__()
        # 输出通道数 = 卷积核数
        self.model1 = nn.Sequential(
            nn.Conv2d(3, 32, 5, padding=2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 32, 5, padding=2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 5, padding=2),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(1024, 64),
            nn.Linear(64, 10)
        )



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

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

import torchvision
from torch.utils.tensorboard import SummaryWriter

from model import *
# 准备数据
from torch import nn
from torch.nn import Sequential, Conv2d, MaxPool2d, Flatten, Linear
from torch.utils.data import DataLoader

train_data = torchvision.datasets.CIFAR10(root="./data", train=True, transform=torchvision.transforms.ToTensor(),
                                          download=True)

test_data = torchvision.datasets.CIFAR10(root="./data", train=False, transform=torchvision.transforms.ToTensor(),
                                         download=True)
# length 长度
train_data_size = len(train_data)
test_data_size = len(test_data)

print("train_data_size = ", train_data_size)

print("test_data_size = ", test_data_size)

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

# 搭建神经网络
node = NeuralNode()

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

# 定义优化器
# SGD 随机梯度下降
# learning_rate = 0.01
learning_rate = 1e-2
optimizer = torch.optim.SGD(node.parameters(), lr=learning_rate)

# 设置训练网络的一些参数
# 记录训练的次数
total_train_step = 0

# 记录测试的次数
total_test_step = 0

# 训练的轮数
epoch = 10

# 添加tensorboard
writer = SummaryWriter("./logs_train")

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

    # 训练步骤开始
    # node.train()
    for data in train_dataloader:
        imgs, targets = data
        outputs = node(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()))
            writer.add_scalar("train_loss", loss.item(), total_train_step)

    # 测试步骤开始
    # node.eval()
    total_test_loss = 0
    total_accuracy = 0
    with torch.no_grad():
        for data in test_dataloader:
            imgs, targets = data
            outputs = node(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("整体测试集上的正确率:{}".format(total_accuracy/test_data_size))
    writer.add_scalar("test_accuracy",total_accuracy/test_data_size,total_test_step)
    writer.add_scalar("test_loss", total_test_loss, total_test_step)
    total_test_step += 1

    torch.save(node,"node_{}.pth".format(i))
    # torch.save(node.state_dict(),"node_{}".format(i))
    print("模型已保存")

writer.close()

运行结果：

C:\Users\Tom\.conda\envs\yolov8\python.exe D:/pythonProject/pytorch16.py
Files already downloaded and verified
Files already downloaded and verified
train_data_size =  50000
test_data_size =  10000
-----------------------第1轮训练开始-------------------
训练次数:100,Loss:2.2813925743103027
训练次数:200,Loss:2.2742743492126465
训练次数:300,Loss:2.2269670963287354
训练次数:400,Loss:2.134626865386963
训练次数:500,Loss:2.0348572731018066
训练次数:600,Loss:2.018798589706421
训练次数:700,Loss:2.000631332397461
整体测试集上的Loss:313.22724294662476
整体测试集上的正确率:0.2825999855995178
模型已保存
-----------------------第2轮训练开始-------------------
训练次数:800,Loss:1.879097819328308
训练次数:900,Loss:1.847396731376648
训练次数:1000,Loss:1.9424469470977783
训练次数:1100,Loss:1.966457724571228
训练次数:1200,Loss:1.7198699712753296
训练次数:1300,Loss:1.6470344066619873
训练次数:1400,Loss:1.7444041967391968
训练次数:1500,Loss:1.8047219514846802
整体测试集上的Loss:301.84073090553284
整体测试集上的正确率:0.3086000084877014
模型已保存
-----------------------第3轮训练开始-------------------
训练次数:1600,Loss:1.7345367670059204
训练次数:1700,Loss:1.6733229160308838
训练次数:1800,Loss:1.9394733905792236
训练次数:1900,Loss:1.702191710472107
训练次数:2000,Loss:1.8698254823684692
训练次数:2100,Loss:1.499967336654663
训练次数:2200,Loss:1.4865745306015015
训练次数:2300,Loss:1.758401870727539
整体测试集上的Loss:270.03703331947327
整体测试集上的正确率:0.3743000030517578
模型已保存
-----------------------第4轮训练开始-------------------
训练次数:2400,Loss:1.7163890600204468
训练次数:2500,Loss:1.3585485219955444
训练次数:2600,Loss:1.5691684484481812
训练次数:2700,Loss:1.6807634830474854
训练次数:2800,Loss:1.474306344985962
训练次数:2900,Loss:1.5900392532348633
训练次数:3000,Loss:1.328945279121399
训练次数:3100,Loss:1.5269112586975098
整体测试集上的Loss:266.39168417453766
整体测试集上的正确率:0.38600000739097595
模型已保存
-----------------------第5轮训练开始-------------------
训练次数:3200,Loss:1.341350793838501
训练次数:3300,Loss:1.4943385124206543
训练次数:3400,Loss:1.4554413557052612
训练次数:3500,Loss:1.555479884147644
训练次数:3600,Loss:1.5840874910354614
训练次数:3700,Loss:1.3194756507873535
训练次数:3800,Loss:1.3042551279067993
训练次数:3900,Loss:1.4639265537261963
整体测试集上的Loss:251.46541500091553
整体测试集上的正确率:0.42089998722076416
模型已保存
-----------------------第6轮训练开始-------------------
训练次数:4000,Loss:1.4055039882659912
训练次数:4100,Loss:1.4073412418365479
训练次数:4200,Loss:1.4922577142715454
训练次数:4300,Loss:1.2540905475616455
训练次数:4400,Loss:1.1579914093017578
训练次数:4500,Loss:1.3563917875289917
训练次数:4600,Loss:1.3859093189239502
整体测试集上的Loss:237.89875304698944
整体测试集上的正确率:0.4487999975681305
模型已保存
-----------------------第7轮训练开始-------------------
训练次数:4700,Loss:1.29132878780365
训练次数:4800,Loss:1.5408289432525635
训练次数:4900,Loss:1.3593134880065918
训练次数:5000,Loss:1.4323735237121582
训练次数:5100,Loss:1.000712513923645
训练次数:5200,Loss:1.3162072896957397
训练次数:5300,Loss:1.1935902833938599
训练次数:5400,Loss:1.3839805126190186
整体测试集上的Loss:224.08775210380554
整体测试集上的正确率:0.4832000136375427
模型已保存
-----------------------第8轮训练开始-------------------
训练次数:5500,Loss:1.2033913135528564
训练次数:5600,Loss:1.1681513786315918
训练次数:5700,Loss:1.1975551843643188
训练次数:5800,Loss:1.18670654296875
训练次数:5900,Loss:1.3126091957092285
训练次数:6000,Loss:1.564153790473938
训练次数:6100,Loss:1.037826418876648
训练次数:6200,Loss:1.0969135761260986
整体测试集上的Loss:210.79971551895142
整体测试集上的正确率:0.5163000226020813
模型已保存
-----------------------第9轮训练开始-------------------
训练次数:6300,Loss:1.3967925310134888
训练次数:6400,Loss:1.138830304145813
训练次数:6500,Loss:1.535596489906311
训练次数:6600,Loss:1.0663063526153564
训练次数:6700,Loss:1.0792276859283447
训练次数:6800,Loss:1.1497656106948853
训练次数:6900,Loss:1.1112866401672363
训练次数:7000,Loss:0.9238876700401306
整体测试集上的Loss:201.00618183612823
整体测试集上的正确率:0.541100025177002
模型已保存
-----------------------第10轮训练开始-------------------
训练次数:7100,Loss:1.2247889041900635
训练次数:7200,Loss:0.9504277110099792
训练次数:7300,Loss:1.0882889032363892
训练次数:7400,Loss:0.8527622818946838
训练次数:7500,Loss:1.2434874773025513
训练次数:7600,Loss:1.236551284790039
训练次数:7700,Loss:0.8347920775413513
训练次数:7800,Loss:1.281840205192566
整体测试集上的Loss:193.47922903299332
整体测试集上的正确率:0.5615000128746033
模型已保存

进程已结束，退出代码 0

[pytorch]完整的模型验证(测试、demo)套路

利用已经训练好的模型，然后给它提供输入

CIFAR 10 model结构

import torch
import torchvision
from PIL import Image
from torch import nn

image_path = "./images/002.png"

image = Image.open(image_path)

# 因为png格式是四个通道，除了RGB三通道外，还有一个透明度通道。所以我们调用image = image.convert("RGB"),保留其颜色通道。
image = image.convert("RGB")
print(image)

transform = torchvision.transforms.Compose([torchvision.transforms.Resize((32,32)),torchvision.transforms.ToTensor()])

image = transform(image)

print(image.shape)


class NeuralNode(nn.Module):
    def __init__(self):
        super(NeuralNode, self).__init__()
        # 输出通道数 = 卷积核数
        self.model1 = nn.Sequential(
            nn.Conv2d(3, 32, 5, padding=2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 32, 5, padding=2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 5, padding=2),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(1024, 64),
            nn.Linear(64, 10)
        )



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

model = torch.load("node_9.pth")

print(model)

image = torch.reshape(image,(1,3,32,32))
model.eval()
with torch.no_grad():
    output = model(image)

print(output)

print(output.argmax(1))

运行结果：

C:\Users\Tom\.conda\envs\yolov8\python.exe D:/pythonProject/test.py
<PIL.Image.Image image mode=RGB size=257x166 at 0x1BCA5488BB0>
torch.Size([3, 32, 32])
NeuralNode(
  (model1): Sequential(
    (0): Conv2d(3, 32, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
    (1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (2): Conv2d(32, 32, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
    (3): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (4): Conv2d(32, 64, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
    (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (6): Flatten(start_dim=1, end_dim=-1)
    (7): Linear(in_features=1024, out_features=64, bias=True)
    (8): Linear(in_features=64, out_features=10, bias=True)
  )
)
tensor([[-1.2682, -6.1971,  4.5997,  3.2537,  3.8549,  5.1549, -0.5000,  2.2357,
         -3.8257, -4.5925]])
tensor([5])

进程已结束，退出代码 0