使用Python实现一个简单的图像分类器

在深度学习和计算机视觉领域，图像分类是一项基础而重要的任务。图像分类的目标是将一张图片分配到一个或多个预定义的类别中。随着卷积神经网络（CNN）的发展，图像分类的准确率得到了极大的提升。

本文将介绍如何使用 Python 和深度学习框架 PyTorch 来构建一个简单的图像分类器。我们将使用经典的 CIFAR-10 数据集进行训练和测试，并展示完整的代码实现流程，包括数据加载、模型构建、训练、评估以及预测。

环境准备

在开始之前，请确保你已经安装了以下库：

你可以通过以下命令安装这些库：

pip install torch torchvision matplotlib

导入必要的库

首先，我们导入项目所需的库：

import torchimport torch.nn as nnimport torch.optim as optimimport torchvisionimport torchvision.transforms as transformsfrom torch.utils.data import DataLoaderimport matplotlib.pyplot as pltimport numpy as np

数据加载与预处理

我们使用 torchvision.datasets.CIFAR10 加载 CIFAR-10 数据集。该数据集包含 60000 张 32x32 的彩色图像，分为 10 个类别，每个类别 6000 张图。

# 定义图像预处理操作transform = transforms.Compose([    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 = DataLoader(train_dataset, batch_size=64, shuffle=True)test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)classes = ('plane', 'car', 'bird', 'cat', 'deer',           'dog', 'frog', 'horse', 'ship', 'truck')

我们可以先看一下部分训练数据的样子：

# 显示部分图像def imshow(img):    img = img / 2 + 0.5  # 反归一化    npimg = img.numpy()    plt.imshow(np.transpose(npimg, (1, 2, 0)))    plt.show()# 获取一批图像dataiter = iter(train_loader)images, labels = next(dataiter)# 显示图像imshow(torchvision.utils.make_grid(images))print(' '.join(f'{classes[labels[j]]}' for j in range(4)))

构建卷积神经网络模型

我们将构建一个简单的 CNN 模型，包含两个卷积层和三个全连接层。

class SimpleCNN(nn.Module):    def __init__(self):        super(SimpleCNN, self).__init__()        self.features = nn.Sequential(            nn.Conv2d(3, 16, kernel_size=3, padding=1),            nn.ReLU(),            nn.MaxPool2d(kernel_size=2),            nn.Conv2d(16, 32, kernel_size=3, padding=1),            nn.ReLU(),            nn.MaxPool2d(kernel_size=2),        )        self.classifier = nn.Sequential(            nn.Linear(32 * 8 * 8, 256),            nn.ReLU(),            nn.Linear(256, 128),            nn.ReLU(),            nn.Linear(128, 10)        )    def forward(self, x):        x = self.features(x)        x = x.view(-1, 32 * 8 * 8)        x = self.classifier(x)        return x# 实例化模型model = SimpleCNN()device = torch.device("cuda" if torch.cuda.is_available() else "cpu")model.to(device)

定义损失函数和优化器

我们使用交叉熵损失函数和随机梯度下降优化器。

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

训练模型

接下来，我们编写训练循环：

num_epochs = 5for epoch in range(num_epochs):    running_loss = 0.0    for i, data in enumerate(train_loader, 0):        inputs, labels = data        inputs, labels = inputs.to(device), labels.to(device)        # 前向传播        outputs = model(inputs)        loss = criterion(outputs, labels)        # 反向传播和优化        optimizer.zero_grad()        loss.backward()        optimizer.step()        running_loss += loss.item()        if i % 200 == 199:  # 每200个batch打印一次            print(f'[Epoch {epoch+1}, Batch {i+1}] Loss: {running_loss / 200:.3f}')            running_loss = 0.0print('Finished Training')

测试模型性能

我们使用测试集来评估模型的准确率：

correct = 0total = 0with torch.no_grad():    for data in test_loader:        images, labels = data        images, labels = images.to(device), labels.to(device)        outputs = model(images)        _, predicted = torch.max(outputs.data, 1)        total += labels.size(0)        correct += (predicted == labels).sum().item()print(f'Accuracy of the network on the 10000 test images: {100 * correct / total:.2f}%')

查看每个类别的准确率

我们还可以分别查看每个类别的识别准确率：

class_correct = list(0. for i in range(10))class_total = list(0. for i in range(10))with torch.no_grad():    for data in test_loader:        images, labels = data        images, labels = images.to(device), labels.to(device)        outputs = model(images)        _, predicted = torch.max(outputs, 1)        c = (predicted == labels).squeeze()        for i in range(4):            label = labels[i]            class_correct[label] += c[i].item()            class_total[label] += 1for i in range(10):    print(f'Accuracy of {classes[i]} : {100 * class_correct[i] / class_total[i]:.2f}%')

保存和加载模型

训练完成后，我们可以将模型保存下来以备后续使用：

# 保存模型PATH = './simple_cnn.pth'torch.save(model.state_dict(), PATH)# 加载模型model = SimpleCNN()model.load_state_dict(torch.load(PATH))model.to(device)

预测单张图像

最后，我们来对一张图像进行预测：

dataiter = iter(test_loader)images, labels = next(dataiter)image = images[0].unsqueeze(0).to(device)model.eval()with torch.no_grad():    output = model(image)    _, predicted = torch.max(output, 1)print(f'Predicted: {classes[predicted.item()]}')imshow(image.cpu().squeeze())

总结

本文介绍了如何使用 PyTorch 构建一个简单的图像分类器。我们完成了从数据加载、模型定义、训练、评估到预测的完整流程。虽然本模型较为简单，但可以作为深入学习图像分类任务的良好起点。

如果你希望进一步提高模型性能，可以尝试以下方法：

希望这篇文章对你理解图像分类任务有所帮助！

字数统计：约 1700 字

如需获取完整源码，请复制上述代码片段并运行即可。