图神经网络 (GCN) 多标签节点分类模型代码优化

本代码使用图神经网络 (GCN) 模型对包含 42 个图、37 个节点的图数据进行多标签节点分类。每个节点具有 8 个标签,标签数据存储在单独的文本文件中。代码经过优化,使用单个图作为输入,避免批处理操作,并使用 BCEWithLogitsLoss 函数进行多标签分类。

数据描述:

  • num_graphs = 42: 图的数量
  • num_nodes = 37: 每个图的节点数量
  • image_size = 40: 节点特征图像大小
  • num_labels = 8: 每个节点的标签数量
  • num_edges = 61: 图中边的数量
  • 节点特征文件是 'C:\Users\jh\Desktop\data\input\images{i}.png_{j}.png' 的所有图片的像素值,其中 i 表示图序号,i 从 1 到 42,j 表示节点序号,j 从 0 到 36。
  • 每个节点有 8 个标签,储存在 'C:\Users\jh\Desktop\data\input\labels{i}_{j}.txt' 文本文件中,标签用空格隔开。
  • 边的关系储存在 'C:\Users\jh\Desktop\data\input\edges_L.csv' csv 文件中,表格中没有 header,第一列为源节点,第二列为目标节点,共有 61 条无向边。
  • 每个节点的标签文件内容都是 '2 3 1 1 3 2 2 1' 这样的格式,只是数字不同,标签类别总共有:0, 1, 2, 3, 4。
  • 进入 GCN 模型的特征是一个图一个图进入的(即每次进入 37 个节点的特征)。

代码:

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
        x = F.relu(self.conv1(x, edge_index))
        x = self.conv2(x, edge_index)
        return x

# 读取边的关系数据
edges = pd.read_csv(r"C:\Users\jh\Desktop\data\input\edges_L.csv", header=None)
edges = edges.values  # 转换为NumPy数组

# 读取节点特征数据
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"
        label = torch.tensor(list(map(int, open(labels_path).read().strip().split())))
        labels.append(label)

# 将节点特征和标签转换为PyTorch的Tensor
x = torch.stack(features)
x = x.view(-1, 3, 40, 40)  # 调整数据的维度
labels = torch.stack(labels)

# 创建图结构
edge_index = torch.tensor(edges, dtype=torch.long).t().contiguous()
data_list = []
for i in range(42):
    data = Data(x=x[i*37:(i+1)*37], edge_index=edge_index)
    data.y = 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 data in data_list:
        output = cnn_model(data.x)
        cnn_output.append(output)
    cnn_output = torch.cat(cnn_output, dim=0)

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

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

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

# 训练模型
num_epochs = 50
for epoch in range(num_epochs):
    gcn_model.train()
    total_loss = 0
    for data in data_list:
        optimizer.zero_grad()
        out = gcn_model(data)
        loss = criterion(out, data.y.float())
        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:.4f}")

# 在验证集上评估模型
gcn_model.eval()
with torch.no_grad():
    total_correct = 0
    total_samples = 0
    for data in data_list:
        out = gcn_model(data)
        predicted = torch.round(torch.sigmoid(out))
        total_correct += (predicted == data.y).sum().item()
        total_samples += data.y.size(0) * data.y.size(1)
    accuracy = total_correct / total_samples
    print(f"Validation Accuracy: {accuracy:.2f}")

修改说明:

  • 使用 BCEWithLogitsLoss() 作为损失函数,适用于多标签分类任务。
  • 将节点特征和标签作为单个图输入模型,而不是使用批处理。
  • 在验证集上评估模型时,使用的是每个节点的真实标签,并计算准确率。

注意:

  • 代码中的数据路径需要根据实际情况进行修改。
  • 该代码仅提供一个示例,可以根据实际需求进行调整。

代码功能:

  • 读取图数据,包括节点特征、标签和边关系。
  • 使用 CNN 模型对节点特征进行降维。
  • 使用 GCN 模型对图进行学习。
  • 使用 BCEWithLogitsLoss() 函数进行多标签分类。
  • 在验证集上评估模型的准确率。

原文地址: https://www.cveoy.top/t/topic/pham 著作权归作者所有。请勿转载和采集!

免费AI点我,无需注册和登录