图像分类训练设计

$$图像分类训练设计$$

# ==  引入工具包  ==
import torch
from torch import nn
import torch.nn.functional as F
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import torch.optim as optim
# == step 0 参数配置 ==

# == step 1 数据处理 ==
norm_mean = [0.33424968,0.33424437, 0.33428448]
norm_std = [0.24796878, 0.24796101, 0.24801227]

train_transform = transforms.Compose([
    transforms.Resize((32, 32)),
    transforms.ToTensor(), # 0-255 归一化到0-1 转Tensor
    transforms.Normalize(norm_mean, norm_std),
])

train_data_path = r"D:\PycharmProjects\AI_Easy_Demo\MyData\split_data\train"
from MyDataset.Cifar10_Dataset import LoadDataset
train_dataset = LoadDataset(data_dir=r"D:\PycharmProjects\AI_Easy_Demo\MyData\split_data\train",transform=train_transform)
train_loader = DataLoader(dataset=train_dataset, batch_size=8, shuffle=True)  # shuffle训练时打乱样本
# == step 2 模型 ==
from MyNet.ResNet import ResNet34
net = ResNet34(num_classes=10, num_linear=512)

# == step 3 损失函数 ==
criterion = nn.NLLLoss()

# == step 4 优化器 ==
optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.9)                        # 选择优化器
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)     # 设置学习率下降策略，每过step_size个epoch，做一次更新
# == step 5 评测函数==

# == step 6 训练 ==
MAX_EPOCH=200
best = [0]  # 存储最优指标，用于Early Stopping

for i in range(MAX_EPOCH):
    print("当前轮转次数:",i+1)
    for idx, data_info in enumerate(train_loader):
        # print("训练数据索引",idx)
        inputs, labels = data_info
        outputs = net(inputs)
        outputs = F.log_softmax(outputs, dim=1)
        optimizer.zero_grad()  # 梯度置零,设置在loss之前
        loss = criterion(outputs,labels)
        loss.backward()  # loss反向传播，梯度的计算
        print(loss)
        # update weights
        optimizer.step()  # 更新所有的参数 根据误差和梯度进行权重的更新

    # if best_loss < epoch_loss:
    torch.save(net.state_dict(), "best.pth")
    scheduler.step()  # 更新学习率



# == step 7 训练可视化 ==

# == inference ==