1、加载依赖

在 PyTorch 中训练卷积神经网络前，第一步加载依赖的常用模块。

class Neural_Network(nn.Module):
   def __init__(self, ):
      super(Neural_Network, self).__init__()
      self.inputSize = 2
      self.outputSize = 1
      self.hiddenSize = 3
      # weights
      self.W1 = torch.randn(self.inputSize, 
      self.hiddenSize) # 3 X 2 tensor
      self.W2 = torch.randn(self.hiddenSize, self.outputSize) 

2、准备数据

PyTorch 中准备数据通常包括 下载数据集、数据预处理（transform）、和 使用 DataLoader 加载数据。示例是使用 MNIST 手写数字数据集的典型流程。

transform = transforms.Compose([
    transforms.ToTensor(),  # 转为Tensor
    transforms.Normalize((0.1307,), 
    (0.3081,))  # 标准化
])

train_dataset = datasets.MNIST('./data', train=True,
download=True, transform=transform)
test_dataset  = datasets.MNIST('./data', train=False, 
download=True, transform=transform)

train_loader = torch.utils.data.DataLoader(train_dataset, 
batch_size=64, shuffle=True)
test_loader  = torch.utils.data.DataLoader(test_dataset,
batch_size=1000, shuffle=False)

3、定义 CNN 模型

PyTorch 中定义卷积神经网络（CNN）模型，通常继承 nn.Module 并实现 __init__ 和 forward 方法。

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # 输入1通道，输出10通道，5x5卷积核
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)  
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        # 展平后接全连接
        self.fc1 = nn.Linear(320, 50)  
        # 输出为10分类
        self.fc2 = nn.Linear(50, 10)   

    def forward(self, x):
        # 卷积 -> ReLU -> 池化
        x = F.relu(F.max_pool2d(self.conv1(x), 2))  
        x = F.relu(F.max_pool2d(self.conv2(x), 2))
        x = x.view(-1, 320)  # 展平
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

model = Net()

4、定义损失函数和优化器

 PyTorch 中定义损失函数和优化器是训练过程的关键步骤，常用于监督学习任务。

import torch.optim as optim
import torch.nn.functional as F

# 定义损失函数：NLLLoss 要与模型输出 log_softmax 配套使用
criterion = F.nll_loss  # 也可以写成 nn.NLLLoss()

# 定义优化器：SGD 优化器，带动量
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

5、训练和预测模型

训练阶段主要包括将数据输入模型进行前向传播、计算损失、执行反向传播并更新参数；预测阶段则在不计算梯度的情况下对新数据进行前向传播，输出模型预测结果。训练过程中模型使用 model.train() 模式启用 Dropout 和 BatchNorm 等机制，预测则使用 model.eval() 模式以确保稳定的推理行为。

1）训练模型

def train(epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % 100 == 0:
            print(f'Train Epoch: {epoch}')
            print(f'[{batch_idx*len(data)}/')
            print(f'{len(train_loader.dataset)}]')  
            print(f'Loss: {loss.item():.6f}')

for epoch in range(1, 5):
    train(epoch)

2）测试模型

def test():
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            output = model(data)
            test_loss += F.nll_loss(output, target,
            reduction='sum').item()  # 求和
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()
    test_loss /= len(test_loader.dataset)
    print(f'Test set: Avg loss: {test_loss:.4f}, ')
    print(f'Accuracy: {correct}/{len(test_loader.dataset)} ')
    print(f'({100. * correct / len(test_loader.dataset):.0f}%)')

test()

6、CNN 训练核心步骤总结