### 使用 PyTorch 训练 AI 模型的方法 #### 准备数据集 为了有效地训练模型，首先需要准备好合适的数据集。这通常涉及到加载、预处理以及划分训练集和测试集的操作。可以利用 `torchvision` 或者自定义的 DataLoader 来完成这项工作[^1]。 ```python from torchvision import datasets, transforms import torch.utils.data as data_utils transform = transforms.Compose([ transforms.ToTensor(), ]) train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True) test_dataset = datasets.MNIST(root='./data', train=False, transform=transform) train_loader = data_utils.DataLoader(dataset=train_dataset, batch_size=64, shuffle=True) test_loader = data_utils.DataLoader(dataset=test_dataset, batch_size=64, shuffle=False) ``` #### 定义网络结构 接着要设计并初始化一个神经网络架构。这里可以选择现有的经典模型或是自己搭建新的模型。对于初学者来说，可以从简单的全连接层或多层感知器入手尝试。 ```python import torch.nn as nn import torch.nn.functional as F class SimpleNet(nn.Module): def __init__(self): super(SimpleNet, self).__init__() self.fc1 = nn.Linear(784, 500) self.fc2 = nn.Linear(500, 10) def forward(self, x): x = x.view(-1, 784) x = F.relu(self.fc1(x)) x = self.fc2(x) return F.log_softmax(x, dim=1) model = SimpleNet() ``` #### 设置损失函数与优化器 选择合适的损失函数（Loss Function），比如交叉熵适用于分类任务；同时挑选一种有效的优化算法如SGD或Adam等用于参数更新过程中的梯度下降操作。 ```python criterion = nn.CrossEntropyLoss() optimizer = torch.optim.Adam(model.parameters(), lr=0.001) ``` #### 编写训练循环 编写一段代码实现完整的训练流程，在每一轮epoch内遍历整个训练集，并执行前向传播、计算loss值、反向传播调整权重等一系列动作。注意区分训练阶段(`model.train()`) 和验证/推理阶段 (`model.eval()`)[^3]。 ```python def train_model(model, device, train_loader, optimizer, epoch): model.train() # 将模型设置为训练模式 for batch_idx, (data, target) in enumerate(train_loader): data, target = data.to(device), target.to(device) output = model(data) loss = criterion(output, target) optimizer.zero_grad() loss.backward() optimizer.step() for epoch in range(1, epochs + 1): train_model(model, 'cpu', train_loader, optimizer, epoch) ``` #### 测试模型性能 最后一步是对已经训练完毕后的模型进行评估，查看其在未见过的新样本上的表现情况。可以通过准确率或其他指标衡量效果的好坏。 ```python def test_model(model, device, test_loader): model.eval() # 切换至评估模式 correct = 0 with torch.no_grad(): for data, target in test_loader: data, target = data.to(device), target.to(device) output = model(data) pred = output.argmax(dim=1, keepdim=True) correct += pred.eq(target.view_as(pred)).sum().item() accuracy = 100 * correct / len(test_loader.dataset) print(f'Test Accuracy: {accuracy:.2f}%') test_model(model, 'cpu', test_loader) ```