01_CNN界的hello world

ai生成代码+自行理解注释
python 复制代码
import torch    #PyTorch核心
import torch.nn as nn   #神经网络模块（层、模型）
import torch.optim as optim #优化器（怎么更新参数）
from torchvision import datasets, transforms    #现成数据集（MNIST就在这里）和 数据预处理。

# 1. 超参数
batch_size = 64 #一次喂64张图片
learning_rate = 0.001 #每次更新参数的步子大小（lr：学太快会炸，太慢学不会）
epochs = 3  #整个数据集训练3遍

# 2. 数据，让数据分布更稳定 → 更容易训练
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

train_loader = torch.utils.data.DataLoader( #DataLoader 喂数据给模型
    datasets.MNIST('./data', train=True, download=True, transform=transform),   #加载数据
    batch_size=batch_size,  #每次64张
    shuffle=True    #打乱数据（防止模型死记顺序）
)


# 3. 定义CNN模型，所有模型必须继承 nn.Module
class SimpleCNN(nn.Module):                 #定义一个叫 SimpleCNN 的类，它继承自 nn.Module
    def __init__(self):                     #构造函数（初始化函数）。创建对象的时候（model = SimpleCNN())，__init__() 就会自动执行
        super(SimpleCNN, self).__init__()   #调用父类 nn.Module 的初始化函数

        #卷积部分
        self.conv = nn.Sequential(  #按顺序堆层
            #输出通道：16，提取16种特征
            nn.Conv2d(1, 16, 3, 1, 1),  # (1,28,28) -> (16,28,28)   #nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding)
            nn.ReLU(),  #激活函数，会把负值置零 → 输出很多 0
            nn.MaxPool2d(2),  # -> (16,14,14)   #保留重要特征，减少计算量

            nn.Conv2d(16, 32, 3, 1, 1),  # -> (32,14,14)
            nn.ReLU(),
            nn.MaxPool2d(2)  # -> (32,7,7)
        )
        #全连接层
        self.fc = nn.Sequential(
            nn.Flatten(),   #把卷积输出"铺平"以便送到全连接层
            nn.Linear(32 * 7 * 7, 128), #第一层全连接，把一个 1568 维的向量 线性变换 成 128 维向量
            nn.ReLU(),
            nn.Linear(128, 10)
        )

    def forward(self, x):   #forward 是模型的核心计算函数，通过调用 model(x) 时由 PyTorch 内部机制触发；
        x = self.conv(x)
        x = self.fc(x)
        return x


# 4. 初始化
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SimpleCNN().to(device)  #模型放到GPU/CPU

criterion = nn.CrossEntropyLoss()   #分类问题最常用损失函数
#优化器就是神经网络的"调节器"，根据梯度调整模型参数，让 loss 下降。
optimizer = optim.Adam(model.parameters(), lr=learning_rate)    #parameters在模型创建时（__init__）就已经被初始化（通常是随机的）

# 5. 训练
for epoch in range(epochs): #训练epochs轮
    model.train()   #开启训练模式
    total_loss = 0  #统计loss

    #每轮循环（batch_idx）喂64张图
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)   #把数据搬到你指定的设备（CPU/GPU）上

        # 前向
        output = model(data)    #这里触发forward()前向传播，得到预测
        loss = criterion(output, target)    #衡量 output 和 target 的差距，并把这个差距转成一个可以优化的数值

        # 反向
        optimizer.zero_grad()   #清空梯度（不然会累加）
        loss.backward() #计算梯度（链式法则）
        optimizer.step()    #更新参数

        total_loss += loss.item()   #.item()：把tensor变成数值(训练是按 batch 进行的：我们更关心整个 epoch 的整体表现)

        if batch_idx % 100 == 0:
            print(f"Epoch {epoch}, Batch {batch_idx}, Loss {loss.item():.4f}")

    print(f"Epoch {epoch} Avg Loss: {total_loss / len(train_loader):.4f}")

print("训练完成")

#保存整个模型
torch.save(model, "mnist_full.pt")

#-------------------------------------------------------------------------------------
#'''
# 切换为推理模式（重要）
model.eval()

# 构造输入
dummy_input = torch.randn(1, 1, 28, 28).to(device)

# 导出 ONNX
torch.onnx.export(model, dummy_input, "mnist.onnx")

print("ONNX 导出完成")
#'''