利用简单卷积神经网络实现图像分类

创作时间:

作者:

@小白创作中心

利用简单卷积神经网络实现图像分类

引用

CSDN

https://m.blog.csdn.net/qq_44646352/article/details/144427436

本文将详细介绍如何使用简单的卷积神经网络（CNN）实现图像分类任务。我们将以CIFAR-10数据集为例，通过两个卷积层和三个全连接层构建模型，最终实现67.79%的测试准确率。

前言

技术不断发展，日新月异。在学习过程中，我们经常会遇到一些基础操作需要重新回顾。本文将以CIFAR-10数据集为例，利用简单的卷积层和全连接层完成图像分类任务。

传统的图像分类任务通常通过人工设计的特征函数来获取图像的特征信息，通过对特征信息的逻辑判断完成图像的分类任务。然而，设计有效的特征提取器是一个极其困难的过程，需要对样本数据有深入的了解，并具备很强的数据敏感度。幸运的是，我们可以利用深度网络的特征提取能力和拟合能力来完成这一任务。

环境准备

本项目基于PyTorch框架实现，需要安装以下库：

torch
torchvision
sklearn

import torch
import torchvision.transforms as transforms
import numpy as np

数据加载

我们将使用公开的CIFAR-10数据集。数据预处理包括随机水平翻转、转换为Tensor以及标准化处理。

# 数据预处理
transform = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])

# 加载CIFAR-10数据集
train_dataset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True)

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

模型设计

我们将搭建一个简单的卷积神经网络（CNN），包含M个卷积层和L个全连接层。

class SimpleCNN(nn.Module):
    def __init__(self, M, L, num_classes):
        super(SimpleCNN, self).__init__()
        self.convs = nn.ModuleList()
        in_channels = 3  # 假设输入为RGB图像
        
        # 添加 M 个卷积层
        for _ in range(M):
            self.convs.append(nn.Conv2d(in_channels, 32, kernel_size=3, padding=1))  # 输出通道数为32
            in_channels = 32
        
        self.pool = nn.MaxPool2d(2, 2)  # 2x2 最大池化
        
        # 添加 L-1 个全连接层
        self.fc_layers = nn.ModuleList()
        input_size = 32 * 16 * 16  # 假设输入图像为64x64
        for _ in range(L - 1):
            self.fc_layers.append(nn.Linear(input_size, 128))  # 隐藏层大小为128
            input_size = 128
        
        self.fc_layers.append(nn.Linear(input_size, num_classes))  # 输出层
    
    def forward(self, x):
        for conv in self.convs:
            x = F.relu(conv(x))
        x = self.pool(x)  # 池化
        
        x = x.view(x.size(0), -1)  # 展平
        
        for fc in self.fc_layers[:-1]:
            x = F.relu(fc(x))
        
        x = self.fc_layers[-1](x)  # 输出层
        return F.log_softmax(x, dim=1)  # Softmax

模型训练

在模型训练过程中，我们将使用TensorBoard进行训练过程可视化，并设置验证集以防止模型过拟合，同时采用提前终止策略。

def train_model(model, device, writer, train_loader, val_loader, criterion, optimizer, num_epochs=10):
    model.to(device)  # 将模型移动到GPU
    model.train()
    best_val_loss = float('inf')
    patience = 10  # 允许的无改进 epoch 数
    trigger_times = 0
    for epoch in range(num_epochs):
        train_loss = 0.0
        epoch = epoch + 1
        for inputs, labels in train_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            optimizer.zero_grad()  # 清空梯度
            outputs = model(inputs)  # 前向传播
            loss = criterion(outputs, labels)  # 计算损失
            loss.backward()  # 反向传播
            optimizer.step()  # 更新参数
            # 记录损失
            train_loss += loss.item()
        train_loss /= len(train_loader)
        writer.add_scalar('Loss/train', train_loss, epoch)  # 记录损失
        print(f'Epoch [{epoch}/{num_epochs}], Loss: {train_loss:.4f}')
        # 在每个 epoch 结束后进行验证
        val_loss = validate_model(model, val_loader, writer, epoch)
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            trigger_times = 0
            # 可以保存模型
            save_model(model, optimizer, epoch)  # 每个epoch保存模型
        else:
            trigger_times += 1
            if trigger_times >= patience:
                print("Early stopping...")
                break

模型测试

测试模型使用测试集，利用钩子函数获取中间层的特征，进行数据可视化。

def test_model(model, test_loader, device, features_List, features_Label):
    model.to(device)  # 将模型移动到GPU
    register_hooks(model)
    model.eval()  # 设置为评估模式
    correct = 0
    total = 0
    with torch.no_grad():
        for inputs, labels in test_loader:
            inputs, labels = inputs.to(device), labels.to(device)  # 将数据移动到GPU
            outputs = model(inputs)
            # 特征可视化
            features = features_hook[-2].cpu()
            features_List.append(features)
            features_Label.append(labels.cpu())
            _, predicted = torch.max(outputs.data, 1)  # 取最大值的索引
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
            features_hook.clear()
    print(f'Accuracy of the model on the 10000 test images: {100 * correct / total:.2f}%')
    return features_List,features_Label

主函数

将各个模块的逻辑进行串连，进行训练。

if __name__ == '__main__':
    # 创建模型实例
    M = 2  # 2个卷积层
    L = 3  # 3个全连接层
    num_classes = 10  # 假设有10个类别
    model = SimpleCNN(M, L, num_classes)
    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    # 数据预处理
    transform = transforms.Compose([
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
    ])
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    writer = create_tensorboard('./SampleCNN/log/sample_CNN')
    
    # 假设我们已经定义了一个 DataLoader train_loader
    # 示例：data_loader_train= DataLoader(dataset, batch_size=32, shuffle=True)
    # 示例：data_loader_val= DataLoader(dataset, batch_size=32, shuffle=True)
    # 训练模型
    train_model(model, device, writer, data_loader_train, data_loader_val, criterion, optimizer, num_epochs=50)
    
    # 关闭 TensorBoard 记录器
    writer.close()

结果可视化

训练集上损失不断下降，验证集中准确率先升后降，对应的损失先降后升，因此提前终止了训练。通过中间层提取的特征信息，对特征分布进行可视化，这里采用的T-SEN的方法。

if __name__ == '__main__':
    x = torch.rand((12, 200))
    y = torch.tensor([0,2,3,1,2,3,1,2,3,1,2,3])
    standardized_data = StandardScaler().fit_transform(x)
    tsne = TSNE(random_state = 42, n_components=2,verbose=0, perplexity=5, n_iter=250).fit_transform(x)
    plt.subplot(121)
    plt.scatter(tsne[:, 0], tsne[:, 1], s= 5, c=y, cmap='Spectral')
    plt.gca().set_aspect('equal', 'datalim')
    plt.colorbar(boundaries=np.arange(5)-0.5).set_ticks(np.arange(4))
    plt.title('Visualizing Kannada MNIST through t-SNE', fontsize=24)
    
    plt.show()
    from sklearn.decomposition import PCA
    pca_50 = PCA(n_components=2)
    pca_result_50 = pca_50.fit_transform(x)
    pca_tsne = TSNE(random_state = 42, n_components=2, verbose=0, perplexity=5, n_iter=250).fit_transform(pca_result_50)
    plt.subplot(122)
    plt.scatter(pca_tsne[:, 0], pca_tsne[:, 1], s= 5, c=y, cmap='Spectral')
    plt.gca().set_aspect('equal', 'datalim')
    plt.colorbar(boundaries=np.arange(5)-0.5).set_ticks(np.arange(4))
    plt.title('Visualizing Kannada MNIST through t-SNE', fontsize=24)
    plt.show()

模型输出特征分布。其中一些类的特征分布还是比较集中，不利于后续分类面的决策，因此特征提取部分还可以继续优化。