安装PyTorch环境

确保已安装Python 3.7+，通过以下命令安装PyTorch（根据CUDA版本选择适配的安装命令）：

# 无CUDA版本
pip install torch torchvision

# 有CUDA 11.7版本
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu117

数据准备与加载

使用torchvision.datasets加载标准数据集（如CIFAR-10），或自定义数据集通过Dataset类实现：

from torchvision import datasets, transforms

# 数据增强与归一化
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

# 加载CIFAR-10
train_data = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_data, batch_size=64, shuffle=True)

模型架构设计

继承nn.Module类构建卷积神经网络（以LeNet为例）：

import torch.nn as nn

class LeNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)  # 输入通道3，输出通道6，卷积核5x5
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16*5*5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(nn.functional.relu(self.conv1(x)))
        x = self.pool(nn.functional.relu(self.conv2(x)))
        x = torch.flatten(x, 1)
        x = nn.functional.relu(self.fc1(x))
        x = nn.functional.relu(self.fc2(x))
        x = self.fc3(x)
        return x

训练流程实现

定义损失函数与优化器，编写训练循环：

import torch.optim as optim

model = LeNet()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

for epoch in range(10):
    running_loss = 0.0
    for i, (inputs, labels) in enumerate(train_loader):
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print(f'Epoch {epoch+1}, Loss: {running_loss/len(train_loader):.3f}')

模型验证与测试

加载测试集评估模型准确率：

test_data = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
test_loader = torch.utils.data.DataLoader(test_data, batch_size=64, shuffle=False)

correct = 0
total = 0
with torch.no_grad():
    for (images, labels) in test_loader:
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'Test Accuracy: {100 * correct / total:.2f}%')

模型保存与加载

使用torch.save保存模型权重或整个模型：

# 保存权重
torch.save(model.state_dict(), 'lenet.pth')

# 加载权重
loaded_model = LeNet()
loaded_model.load_state_dict(torch.load('lenet.pth'))