安装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'))