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 训练核心步骤总结
步骤 | 内容 | 说明 |
---|---|---|
1️⃣ | 导入依赖 | torch , torchvision 等 |
2️⃣ | 加载数据 | 使用 DataLoader 加载 MNIST 等数据集 |
3️⃣ | 定义模型 | 使用 nn.Conv2d , nn.Linear 组合模型结构 |
4️⃣ | 定义损失函数 和优化器 | 常用 NLLLoss + SGD 或 Adam |
5️⃣ | 训练模型 | 循环:前向传播、计算损失、 反向传播、优化更新 |
6️⃣ | 测试模型 | 评估精度、计算平均损失 |