使用 PyTorch 完成图像分类任务是一个常见的机器学习任务，通常包括数据准备、模型设计、训练和测试等步骤。以下是一个完整的流程示例，帮助你快速入门：
- 环境准备
确保你已经安装了 PyTorch 和其他必要的库。如果没有安装，可以通过以下命令安装：
pip install torch torchvision
- 数据准备
使用 torchvision 提供的 datasets 和 transforms 模块来加载和预处理数据集。以 CIFAR-10 数据集为例：
import torch
import torchvision
import torchvision.transforms as transforms
数据预处理
transform = transforms.Compose([
transforms.RandomHorizontalFlip(), # 随机水平翻转
transforms.ToTensor(), # 转换为张量
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) # 归一化
])
加载数据集
train_dataset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
test_dataset = torchvision.datasets.CIFAR10(root='./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=64, shuffle=False)
- 定义模型
定义一个简单的卷积神经网络（CNN）用于图像分类：
import torch.nn as nn
import torch.nn.functional as F
class SimpleCNN(nn.Module):
def init(self):
super(SimpleCNN, self).init()
self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
self.fc1 = nn.Linear(64 * 8 * 8, 128) # CIFAR-10 图像大小为 32x32
self.fc2 = nn.Linear(128, 10) # 输出 10 类
def forward(self, x):
x = F.relu(F.max_pool2d(self.conv1(x), 2))
x = F.relu(F.max_pool2d(self.conv2(x), 2))
x = x.view(-1, 64 * 8 * 8)
x = F.relu(self.fc1(x))
x = self.fc2(x)
return x
model = SimpleCNN()
- 定义损失函数和优化器
使用交叉熵损失函数和随机梯度下降（SGD）优化器：
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
- 训练模型
训练模型并记录训练过程中的损失：
num_epochs = 10
for epoch in range(num_epochs):
model.train()
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()
if i % 100 == 99:
print(f"[{epoch + 1}, {i + 1}] loss: {running_loss / 100:.3f}")
running_loss = 0.0
print("Training finished.")
- 测试模型
在测试集上评估模型的性能：
model.eval()
correct = 0
total = 0
with torch.no_grad():
for inputs, labels in test_loader:
outputs = model(inputs)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print(f"Accuracy: {100 * correct / total:.2f}%")
- 保存和加载模型
保存训练好的模型：
torch.save(model.state_dict(), "simple_cnn.pth")
加载模型：
model = SimpleCNN()
model.load_state_dict(torch.load("simple_cnn.pth"))
model.eval()
总结
以上是一个完整的使用 PyTorch 进行图像分类的流程。你可以根据需要调整模型结构、数据预处理方式和超参数，以获得更好的性能。如果需要处理更复杂的数据集（如 ImageNet），可以考虑使用预训练模型（如 ResNet、VGG 等）来提高效率和性能。