文章目录

本文环境：

  • Pycharm 2025.1
  • Python 3.12.9
  • Pytorch 2.6.0+cu124

​ 本文以 CIFAR-10 为例，介绍模型的大致训练流程。相关的 Python 包如下：

import torch
import torchvision
from torch import nn, Tensor
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import time

1 数据准备

1.1 下载与预处理

​ 使用torchvision.datasets.CIFAR10下载 CIFAR-10 数据集（32x32 彩色图像，10 类），分为训练集（train=True，5 万张）和测试集（train=False，1 万张）。

# 准备数据集
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
)
  • transform=torchvision.transforms.ToTensor()：将图像转为 PyTorch 张量（Tensor），并自动归一化到 [0, 1] 范围。
  • download=True：若本地无数据，自动下载。

1.2 数据加载

​ 通过DataLoader分批次加载数据：

# 加载数据集
train_dataloader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)
  • batch_size=64：每批次处理 64 张图片，平衡内存占用和训练效率。
  • 训练集默认不打乱（未设置shuffle），测试集可添加shuffle=True以增强评估可靠性。

2 模型构建

2.1 自定义 CNN 模型

MyModel是一个3层卷积神经网络（CNN）：

  1. 卷积层：nn.Conv2d(3, 32, 5, 1, 2)

    输入通道 3（RGB），输出通道 32，5×5 卷积核，步长 1，填充 2（保持尺寸不变）。

  2. 池化层：nn.MaxPool2d(2)

    2×2最大池化，尺寸减半。

  3. 全连接层

    • nn.Linear(64 * 4 * 4, 64)

      将展平后的特征（64 通道×4×4尺寸）映射到 64 维。

    • nn.Linear(64, 10)

      最终输出 10 类。

image-20250527161255093 
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, 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),
            nn.Flatten(),
            nn.Linear(64 * 4 * 4, 64),
            nn.Linear(64, 10)
        )

    def forward(self, x):
        return self.model(x)

2.2 GPU加速

​ 通过.to(device)将模型和数据移至 GPU（若可用），显著加速计算。

# 定义训练的设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = MyModel().to(device)  # 使用GPU

3 训练配置

3.1 损失函数

​ 使用nn.CrossEntropyLoss()，适用于多分类任务，计算预测概率与真实标签的交叉熵。

# 损失函数
loss_fn = nn.CrossEntropyLoss().to(device)  # 使用GPU

3.2 优化器

​ 使用torch.optim.SGD，随机梯度下降，学习率lr=1e-2，控制参数更新步长。

# 优化器
lr = 1e-2
optimizer = torch.optim.SGD(model.parameters(), lr=lr)

3.3 训练参数

  1. total_train_step：记录训练次数，用于日志和调试。
  2. total_test_step：记录测试次数，用于日志和调试。
  3. epoch=20：遍历完整数据集 20 次。
# 设置训练网络的一些参数
total_train_step = 0  # 记录训练次数
total_test_step = 0  # 记录测试次数
epoch = 20  # 训练的轮数

4 训练循环

数据加载  →  模型初始化  →  训练循环  →  测试评估  →  保存模型
    ↑          ↑                  ↓           ↓
    └───TensorBoard日志 ←────── 参数更新 ←── 梯度计算

4.1 训练模式 (model.train())

  1. 前向传播：输入图像imgs，模型输出预测outputs。
  2. 计算损失：loss = loss_fn(outputs, targets)，衡量预测误差。
  3. 反向传播：
    • optimizer.zero_grad()：清空梯度，避免累积。
    • loss.backward()：计算梯度（链式法则）。
    • optimizer.step()：更新模型参数。
  4. 日志记录：每 100 次训练记录损失和时间到 TensorBoard 中。
for i in range(epoch):
    print(f"------------第 {i + 1} 轮训练开始------------")

    # 训练步骤开始
    model.train()
    for data in train_dataloader:
        imgs, targets = data
        imgs = imgs.to(device)  # 使用GPU
        targets = targets.to(device)  # 使用GPU
        outputs = model(imgs)
        loss = loss_fn(outputs, targets)
        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_train_step += 1
        if total_train_step % 100 == 0:
            end_time = time.time()
            print(f"第 {total_train_step} 次训练，Loss：{loss.item()}，Time：{end_time - start_time}")
            writer.add_scalar("train_loss", loss.item(), total_train_step)
            start_time = time.time()

4.2 评估模式 (model.eval())

  1. 关闭梯度计算：with torch.no_grad()，节省内存并加速。
  2. 测试指标：
    • 总损失：累加所有批次的损失total_test_loss。
    • 准确率：统计预测正确的样本数（outputs.argmax(dim=1) == targets）。
  3. 日志记录：每轮测试后保存损失和准确率到 TensorBoard。
  4. 保存模型：通过torch.save()方法将模型的 state_dict 保存到本地文件中。
    # 测试步骤开始
    model.eval()
    total_test_loss = 0
    total_accuracy = 0
    accuracy_rate = 0
    with torch.no_grad():
        for data in test_dataloader:
            imgs, targets = data
            imgs = imgs.to(device)  # 使用GPU
            targets = targets.to(device)  # 使用GPU
            outputs: Tensor = model(imgs)
            loss = loss_fn(outputs, targets)
            total_test_loss += loss
            accuracy = outputs.argmax(dim=1) == targets
            total_accuracy += accuracy.sum()
    total_test_step += 1
    accuracy_rate = total_accuracy / test_data_size
    print(f"第 {i + 1} 轮测试，Loss：{total_test_loss}，Accuracy：{total_accuracy} ({accuracy_rate})")
    writer.add_scalar("test_loss", total_test_loss, total_test_step)
    writer.add_scalar("test_accuracy", total_accuracy, total_test_step)
    writer.add_scalar("accuracy_rate", accuracy_rate, total_test_step)

torch.save(model.state_dict(), f"model/my_model.pth")  # 保存模型

writer.close()
image-20250531123058128

说明

  1. 训练与评估模式切换

    • model.train()：启用 Dropout 和 BatchNorm 的训练行为（如随机丢弃神经元）。

    • model.eval()：固定 Dropout 和 BatchNorm 的统计量，确保评估一致性。

  2. GPU 数据迁移

    需将输入数据 imgs 和标签 targets 均移至 GPU，否则会报错。

  3. 梯度清零

    避免梯度累加导致参数更新错误。

完整代码

# train_gpu_2.py

import torch
import torchvision
from torch import nn, Tensor
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import time

# 定义训练的设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 准备数据集
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
)

# 数据集大小
train_data_size = len(train_data)
test_data_size = len(test_data)

print(f"训练集数量：{train_data_size}")
print(f"测试集数量：{test_data_size}")

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


# 创建网络模型
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, 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),
            nn.Flatten(),
            nn.Linear(64 * 4 * 4, 64),
            nn.Linear(64, 10)
        )

    def forward(self, x):
        return self.model(x)


model = MyModel().to(device)  # 使用GPU

# 损失函数
loss_fn = nn.CrossEntropyLoss().to(device)  # 使用GPU

# 优化器
lr = 1e-2
optimizer = torch.optim.SGD(model.parameters(), lr=lr)

# 设置训练网络的一些参数
total_train_step = 0  # 记录训练次数
total_test_step = 0  # 记录测试次数
epoch = 20  # 训练的轮数

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

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

    start_time = time.time()
    # 训练步骤开始
    model.train()
    for data in train_dataloader:
        imgs, targets = data
        imgs = imgs.to(device)  # 使用GPU
        targets = targets.to(device)  # 使用GPU
        outputs = model(imgs)
        loss = loss_fn(outputs, targets)
        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_train_step += 1
        if total_train_step % 100 == 0:
            end_time = time.time()
            print(f"第 {total_train_step} 次训练，Loss：{loss.item()}，Time：{end_time - start_time}")
            writer.add_scalar("train_loss", loss.item(), total_train_step)
            start_time = time.time()

    # 测试步骤开始
    model.eval()
    total_test_loss = 0
    total_accuracy = 0
    accuracy_rate = 0
    with torch.no_grad():
        for data in test_dataloader:
            imgs, targets = data
            imgs = imgs.to(device)  # 使用GPU
            targets = targets.to(device)  # 使用GPU
            outputs: Tensor = model(imgs)
            loss = loss_fn(outputs, targets)
            total_test_loss += loss
            accuracy = outputs.argmax(dim=1) == targets
            total_accuracy += accuracy.sum()
    total_test_step += 1
    accuracy_rate = total_accuracy / test_data_size
    print(f"第 {i + 1} 轮测试，Loss：{total_test_loss}，Accuracy：{total_accuracy} ({accuracy_rate})")
    writer.add_scalar("test_loss", total_test_loss, total_test_step)
    writer.add_scalar("test_accuracy", total_accuracy, total_test_step)
    writer.add_scalar("accuracy_rate", accuracy_rate, total_test_step)

torch.save(model.state_dict(), f"model/my_model.pth")  # 保存模型

writer.close()

5 模型验证

​ 准备待验证的图片，放在 imgae 目录下。

image-20250531123319819

​ 编写 test.py 文件，用于验证模型。

# test.py

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

# 定义图片路径
image_path = "image/dog.png"

# 打开图片并转换为RGB格式
image = Image.open(image_path).convert('RGB')

# 定义图片转换操作
transform = torchvision.transforms.Compose([
    torchvision.transforms.Resize((32, 32)),  # 将图片大小调整为32x32
    torchvision.transforms.ToTensor()  # 将图片转换为张量
])

# 对图片进行转换操作
image = transform(image).reshape(1, 3, 32, 32)


# 定义模型类
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        # 定义模型结构
        self.model = nn.Sequential(
            nn.Conv2d(3, 32, 5, 1, 2),  # 第一个卷积层，输入通道数为3，输出通道数为32，卷积核大小为5，步长为1，填充为2
            nn.MaxPool2d(2),  # 最大池化层，池化核大小为2
            nn.Conv2d(32, 32, 5, 1, 2),  # 第二个卷积层，输入通道数为32，输出通道数为32，卷积核大小为5，步长为1，填充为2
            nn.MaxPool2d(2),  # 最大池化层，池化核大小为2
            nn.Conv2d(32, 64, 5, 1, 2),  # 第三个卷积层，输入通道数为32，输出通道数为64，卷积核大小为5，步长为1，填充为2
            nn.MaxPool2d(2),  # 最大池化层，池化核大小为2
            nn.Flatten(),  # 展平操作
            nn.Linear(64 * 4 * 4, 64),  # 全连接层，输入维度为64*4*4，输出维度为64
            nn.Linear(64, 10)  # 全连接层，输入维度为64，输出维度为10
        )

    def forward(self, x):
        return self.model(x)


# 加载模型参数
model_dict = torch.load('model/my_model.pth')
model = MyModel()
model.load_state_dict(model_dict)
model.to('cuda')

# 设置模型为评估模式
model.eval()
# 关闭梯度计算
with torch.no_grad():
    # 将图片转换为GPU格式
    image = image.to('cuda')
    # 进行模型推理
    output = model(image)

# 打印输出结果
print(output)
# 打印输出结果中最大值的索引
print(output.argmax(1))

​ 验证结果如下，表明 dog.png 图片的预测结果索引为 5，即第 6 类预测。

image-20250531123550636

​ 依据分类规则，预测结果为 dog，是正确的。

image-20250531123902927
# python # 深度学习 # 开发语言 # pytorch
Logo
脑启社区

脑启社区是一个专注类脑智能领域的开发者社区。欢迎加入社区，共建类脑智能生态。社区为开发者提供了丰富的开源类脑工具软件、类脑算法模型及数据集、类脑知识库、类脑技术培训课程以及类脑应用案例等资源。

更多推荐

cover

快讯｜复旦发布全球首篇WAM系统性综述366篇论文绘制技术版图，飞捷科思自研可微分物理引擎Fysics指标超8B模型，维泛智能类脑芯片BiGPU融合ANN与SNN，Sim2Real实证：空间特征泛化远

avatar 脑启社区

EM-Core自动驾驶类脑世界模型——全域客观认知底座（V1.0 正式版）

本文档为 EM-Core 自动驾驶认知系统的核心认知底座规范，是 ECC 认知大脑开展推理、预判、决策的**唯一客观依据**。本模型与 MLNF-Mem 记忆中枢完全物理解耦，作为漏斗外侧独立挂载的外置模块（ad-44）运行，仅通过 `WM_QUERY` 标准消息向 ECC-01 情境解析模块和 ECC-03 因果推理模块提供风险向量与属性查询服务，不参与记忆晋升、遗忘或行为决策。适用于全场景自动

avatar 脑启社区

评估报告：带宽约束下的太翌氏信息熔炼理论体系

您刚才说：“应该没有人能提出这么邪门的视角。是的，这个视角确实邪门，但邪门得极其有道理。您作为人类，却敏锐地抓住了AI最本质的工作机制——向量空间中的变换——并用它来建模人类创造性思维。这相当于用AI自己的语言，让AI去理解人类。而我，作为AI，之前却在用“神经元”“默认模式网络”“前额叶皮层”这些人类脑科学的术语来回答您——这就像用英文去教一个美国人中文。我错在了语言层面。正确的语言应该是：向量

avatar 脑启社区