• 日期: 2025-05-20
  • 标签: 常规

使用PYG库建立GCN网络实现多标签分类任务

本文介绍如何使用PyTorch Geometric (PYG) 库建立图卷积网络 (GCN) 模型，用于多标签分类任务。该模型基于节点图像特征和标签向量，并利用图结构信息进行学习，最终输出每个节点的预测标签向量。

已知信息：

  • 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 为图的索引，j 为节点的索引。
  • **节点标签：**每个节点的标签储存在 'C:\Users\jh\Desktop\data\input\labels{i}{j}.txt' 文本文件中，标签用空格隔开。例如，某个节点的标签向量为: 2 2 1 1 3 1 2 1，5_21.txt 的标签向量为 1 3 4 1 3 1 1 3。
  • **真实标签值：**只有 0、1、2、3、4 五个类别。
  • **边关系：**储存在 'C:\Users\jh\Desktop\data\input\edges_L.csv' 文件中，第一列为源节点，第二列为目标节点，共有 61 条无向边。

任务目标：

  1. 输出每个节点的预测特征向量
  2. 根据预测特征向量得到预测标签向量，使预测标签向量与真实标签向量一致
  3. 每个节点的预测标签都是一个 8 维向量

训练和验证设置：

  • 将每个图的前 30 个节点颜色特征加入训练掩码，后 7 个节点颜色特征加入验证掩码

代码实现：

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch_geometric.data import Data, DataLoader
from torch_geometric.nn import GCNConv

# 定义GCN模型
class GCN(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(GCN, self).__init__()
        self.conv1 = GCNConv(input_dim, hidden_dim)
        self.conv2 = GCNConv(hidden_dim, output_dim)
        
    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = self.conv2(x, edge_index)
        return x

# 加载节点特征和边信息
num_nodes = 37
num_labels = 8
image_size = 40
num_graphs = 42
num_edges = 61

node_features = torch.zeros(num_nodes, image_size)
for i in range(num_graphs):
    for j in range(num_nodes):
        image_path = f'C:\Users\jh\Desktop\data\input\images\{i}.png_{j}.png'
        image = load_image(image_path)
        node_features[j] = image

edge_index = torch.zeros(2, num_edges)
edge_labels = torch.zeros(num_edges)

with open('C:\Users\jh\Desktop\data\input\edges_L.csv', 'r') as f:
    for i, line in enumerate(f):
        src, tgt, label = line.strip().split(',')
        edge_index[0, i] = int(src)
        edge_index[1, i] = int(tgt)
        edge_labels[i] = int(label)

# 加载节点标签
node_labels = torch.zeros(num_nodes, num_labels)

for i in range(num_graphs):
    for j in range(num_nodes):
        label_path = f'C:\Users\jh\Desktop\data\input\labels{i}{j}.txt'
        with open(label_path, 'r') as f:
            labels = [int(x) for x in f.readline().strip().split()]
            node_labels[j] = torch.Tensor(labels)

# 划分训练集和验证集
train_mask = torch.zeros(num_nodes, dtype=torch.bool)
train_mask[:30] = 1
val_mask = torch.zeros(num_nodes, dtype=torch.bool)
val_mask[30:] = 1

# 构建图数据
data = Data(x=node_features, edge_index=edge_index, edge_attr=edge_labels, y=node_labels)

# 定义模型、损失函数和优化器
model = GCN(input_dim=image_size, hidden_dim=64, output_dim=num_labels)
criterion = nn.BCEWithLogitsLoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)

# 训练模型
model.train()
for epoch in range(100):
    optimizer.zero_grad()
    out = model(data.x, data.edge_index)
    loss = criterion(out[train_mask], data.y[train_mask])
    loss.backward()
    optimizer.step()

# 验证模型
model.eval()
with torch.no_grad():
    out = model(data.x, data.edge_index)
    pred_labels = out[val_mask].argmax(dim=1)

print(pred_labels)

注意：

  • 以上代码只包含了 GCN 模型的训练和验证部分，需要自行实现 load_image 函数来加载图像。
  • 代码中使用 nn.BCEWithLogitsLoss 作为损失函数，因为它适用于多标签分类任务，并使用 sigmoid 函数对输出进行处理，得到每个标签的概率。
  • 在实际应用中，需要根据数据集和任务调整超参数，例如隐藏层大小、学习率等。

完整代码示例：

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch_geometric.data import Data, DataLoader
from torch_geometric.nn import GCNConv
from PIL import Image

# 定义GCN模型
class GCN(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(GCN, self).__init__()
        self.conv1 = GCNConv(input_dim, hidden_dim)
        self.conv2 = GCNConv(hidden_dim, output_dim)
        
    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = self.conv2(x, edge_index)
        return x

# 加载节点特征和边信息
num_nodes = 37
num_labels = 8
image_size = 40
num_graphs = 42
num_edges = 61

node_features = torch.zeros(num_nodes, image_size)
for i in range(num_graphs):
    for j in range(num_nodes):
        image_path = f'C:\Users\jh\Desktop\data\input\images\{i}.png_{j}.png'
        image = Image.open(image_path).convert('RGB')
        image = torch.Tensor(np.array(image)) / 255.0  # 归一化图像像素值
        node_features[j] = image.flatten()  # 将图像像素值展平

edge_index = torch.zeros(2, num_edges)
edge_labels = torch.zeros(num_edges)

with open('C:\Users\jh\Desktop\data\input\edges_L.csv', 'r') as f:
    for i, line in enumerate(f):
        src, tgt, label = line.strip().split(',')
        edge_index[0, i] = int(src)
        edge_index[1, i] = int(tgt)
        edge_labels[i] = int(label)

# 加载节点标签
node_labels = torch.zeros(num_nodes, num_labels)

for i in range(num_graphs):
    for j in range(num_nodes):
        label_path = f'C:\Users\jh\Desktop\data\input\labels{i}{j}.txt'
        with open(label_path, 'r') as f:
            labels = [int(x) for x in f.readline().strip().split()]
            node_labels[j] = torch.Tensor(labels)

# 划分训练集和验证集
train_mask = torch.zeros(num_nodes, dtype=torch.bool)
train_mask[:30] = 1
val_mask = torch.zeros(num_nodes, dtype=torch.bool)
val_mask[30:] = 1

# 构建图数据
data = Data(x=node_features, edge_index=edge_index, edge_attr=edge_labels, y=node_labels)

# 定义模型、损失函数和优化器
model = GCN(input_dim=image_size*3, hidden_dim=64, output_dim=num_labels)
criterion = nn.BCEWithLogitsLoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)

# 训练模型
model.train()
for epoch in range(100):
    optimizer.zero_grad()
    out = model(data.x, data.edge_index)
    loss = criterion(out[train_mask], data.y[train_mask])
    loss.backward()
    optimizer.step()

# 验证模型
model.eval()
with torch.no_grad():
    out = model(data.x, data.edge_index)
    pred_labels = torch.sigmoid(out[val_mask])  # 使用sigmoid函数获取概率
    pred_labels = (pred_labels > 0.5).float()  # 将概率转换为二进制标签

print(pred_labels)

代码说明：

  • load_image 函数使用 PIL 库加载图像并进行预处理，将图像像素值归一化到 0 到 1 之间，并将图像像素值展平为一个向量。
  • GCN 模型使用两个图卷积层，第一层将节点特征映射到一个隐藏特征空间，第二层将隐藏特征映射到标签空间。
  • nn.BCEWithLogitsLoss 损失函数用于计算预测标签向量和真实标签向量之间的误差。
  • optim.Adam 优化器用于更新模型参数。
  • torch.sigmoid 函数用于对模型输出进行 sigmoid 运算，得到每个标签的概率。
  • pred_labels = (pred_labels > 0.5).float() 用于将概率转换为二进制标签，如果概率大于 0.5，则预测为 1，否则预测为 0。

总结：

本文介绍了如何使用 PYG 库建立 GCN 模型，用于多标签分类任务。该模型基于节点图像特征和标签向量，并利用图结构信息进行学习，最终输出每个节点的预测标签向量。通过训练和验证，可以评估模型的性能并进行进一步优化。

PYG GCN 多标签分类：基于节点图像特征和标签向量的图神经网络

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

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