基于CNN和GCN的图像分类模型实现

本文介绍了如何使用CNN和GCN构建图像分类模型，并提供了详细的代码示例。模型首先使用CNN对图像进行特征提取，然后使用GCN对提取到的特征进行图卷积操作，最终实现图像分类。

代码实现

import os
import pandas as pd
import torch
import torch.nn as nn
from torch_geometric.data import Data, DataLoader
from torch_geometric.nn import GCNConv
import torch.nn.functional as F
from torchvision import transforms
from PIL import Image
from sklearn.model_selection import train_test_split

# 定义CNN网络
class CNN(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, 16, kernel_size=3, stride=1, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
        self.conv2 = nn.Conv2d(16, out_channels, kernel_size=3, stride=1, padding=1)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = self.pool(x)
        x = F.relu(self.conv2(x))
        x = self.pool(x)
        return x

# 定义GCN模型
class GCN(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(GCN, self).__init__()
        self.conv1 = GCNConv(in_channels, 128)
        self.conv2 = GCNConv(128, out_channels)

    def forward(self, data):
        x, edge_index = data.x, data.edge_index
        # print(x.shape)
        x = F.relu(self.conv1(x, edge_index))
        x = self.conv2(x, edge_index)
        return x

# 读取边的关系数据
edges = pd.read_csv('C:\Users\jh\Desktop\data\input\edges_L.csv', header=None)
edges = edges.values

# 读取节点特征数据
features = []
labels = []
for i in range(1, 43):
    for j in range(37):
        image_path = f'C:\Users\jh\Desktop\data\input\images\{i}.png_{j}.png'
        image = Image.open(image_path).convert('RGB')
        transform = transforms.Compose([transforms.Resize((40, 40)), transforms.ToTensor()])
        image_tensor = transform(image)
        features.append(image_tensor)

        labels_path = f'C:\Users\jh\Desktop\data\input\labels\{i}_{j}.txt'
        with open(labels_path, 'r') as file:
            label = [int(l) for l in file.readline().strip().split()]
            labels.extend(label)

x = torch.stack(features)
x = x.view(-1, 3, 40, 40)

# 划分训练集和验证集的掩码
mask_train = torch.zeros(42, 37, dtype=torch.bool)
mask_val = torch.zeros(42, 37, dtype=torch.bool)
for i in range(42):
    mask_train[i, :30] = 1
    mask_val[i, 30:] = 1
mask_train = mask_train.view(-1)
mask_val = mask_val.view(-1)

# 创建图结构
edge_index = torch.tensor(edges, dtype=torch.long).t().contiguous()
data_list = []
for i in range(42):
    data = Data(x=x, edge_index=edge_index)
    data.mask_train = mask_train[i*37:(i+1)*37]
    data.mask_val = mask_val[i*37:(i+1)*37]
    data.y = torch.tensor(labels[i*37:(i+1)*37])
    data_list.append(data)

# 创建CNN模型实例，降维至8维
cnn_model = CNN(in_channels=3, out_channels=8)

# 使用CNN模型对节点特征进行降维
with torch.no_grad():
    cnn_output = []
    for i in range(42):
        x_i = x[i*37:(i+1)*37].unsqueeze(1)
        x_i = x_i.squeeze(1)
        output_i = cnn_model(x_i)
        output_i = output_i.view(output_i.size(0), -1)
        cnn_output.append(output_i)
    cnn_output = torch.cat(cnn_output, dim=0)

# 将降维后的特征加入data对象
for i in range(42):
    data_list[i].x = cnn_output[i*37:(i+1)*37]

# 创建GCN模型实例
gcn_model = GCN(in_channels=800, out_channels=8)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(gcn_model.parameters(), lr=0.002)

# 训练模型
num_epochs = 2000
for epoch in range(num_epochs):
    gcn_model.train()
    total_loss = 0

    for data in data_list:
        optimizer.zero_grad()
        out = gcn_model(data)
        labels_batch = data.y

        # 使用交叉熵损失函数
        loss = criterion(out, labels_batch)

        loss.backward()
        optimizer.step()
        total_loss += loss.item()

    avg_loss = total_loss / len(data_list)
    print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {avg_loss:.6f}')

# 在验证集上评估模型
gcn_model.eval()
with torch.no_grad():
    total_correct = 0
    total_samples = 0
    for data in data_list:
        out = gcn_model(data)
        labels_batch = data.y
        predicted = out.argmax(dim=-1)
        total_correct += (predicted == labels_batch).sum().item()
        total_samples += labels_batch.size(0)

    accuracy = total_correct / total_samples
    print(f'Validation Accuracy: {accuracy:.4f}')

代码优化

为了改进损失值，可以尝试以下方法：

调整学习率：尝试减小学习率，例如将lr=0.002改为lr=0.0001。较小的学习率可能需要更多的训练迭代次数，但它可能有助于模型收敛到更低的损失值。
增加训练迭代次数：增加num_epochs的值，例如将num_epochs=2000改为num_epochs=5000。更多的训练迭代次数可能有助于模型进一步降低损失值。
使用更复杂的模型：尝试增加GCN模型中的层数或节点数。更复杂的模型可能有更强的拟合能力，可以更好地适应数据集。
数据增强：尝试在图像数据上应用一些数据增强技术，例如旋转、翻转或随机裁剪。数据增强可以增加数据集的多样性，有助于模型学习更丰富的特征表示。
正则化：尝试在GCN模型中添加正则化项，例如L1或L2正则化。正则化可以防止模型过拟合训练数据，并提高泛化能力。
批量归一化：尝试在GCN模型中添加批量归一化层。批量归一化可以加速模型收敛，提高模型的稳定性和泛化能力。

请注意，改进模型的方法是一个迭代过程，需要根据具体问题和数据集进行不断尝试和调整。