使用PYG库建立GCN网络实现多标签分类任务 - 代码示例 - 常规

使用PYG库建立GCN网络实现多标签分类任务 - 代码示例

本文提供使用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'，所有图片的像素值，每个节点有8个标签。
标签文件：'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五个类别，但是每个节点的标签是一个8维的标签向量。
边关系文件：'C:\Users\jh\Desktop\data\input\edges_L.csv'，csv文件中没有header，第一列为源节点，第二列为目标节点，共有61条无向边。

目标：

输出每个节点的预测特征向量。
根据预测特征向量得到预测标签向量，使预测标签向量与真实标签向量一致。
每个节点的预测标签都是一个8维向量，而不是输出概率向量。

数据划分：

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

代码示例：

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

# Define the GCN model
class GCN(nn.Module):
    def __init__(self, num_features, hidden_size, num_labels):
        super(GCN, self).__init__()
        self.conv1 = GCNConv(num_features, hidden_size)
        self.conv2 = GCNConv(hidden_size, num_labels)
    
    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = self.conv2(x, edge_index)
        return x

# Load node features and labels
num_graphs = 42
num_nodes = 37
image_size = 40
num_labels = 8

node_features = []
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'
        # Load and process the image to get node feature vector
        # ...
        node_features.append(node_feature_vector)

node_features = torch.tensor(node_features, dtype=torch.float)

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 file:
            label_vector = [int(label) for label in file.read().split()]
        labels.append(label_vector)

labels = torch.tensor(labels, dtype=torch.float)

# Load edge relations
edge_data = []
with open('C:\Users\jh\Desktop\data\input\edges_L.csv', 'r') as file:
    for line in file:
        src_node, tgt_node = map(int, line.split(','))
        edge_data.append((src_node, tgt_node))

edge_index = torch.tensor(edge_data, dtype=torch.long).t().contiguous()

# Split node features and labels into train and validation sets
train_mask = torch.zeros(num_graphs * num_nodes, dtype=torch.bool)
train_mask[:30*num_nodes] = 1  # Set first 30 nodes per graph as training nodes
val_mask = ~train_mask

train_features = node_features[train_mask]
val_features = node_features[val_mask]
train_labels = labels[train_mask]
val_labels = labels[val_mask]

# Create the GCN model
num_features = node_features.size(1)
hidden_size = 64
model = GCN(num_features, hidden_size, num_labels)

# Define the loss function
criterion = nn.BCEWithLogitsLoss()

# Define the optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

# Training loop
num_epochs = 100

for epoch in range(num_epochs):
    model.train()
    optimizer.zero_grad()
    
    # Forward pass
    train_logits = model(train_features, edge_index)
    train_loss = criterion(train_logits, train_labels)
    train_loss.backward()
    optimizer.step()
    
    # Validation pass
    model.eval()
    val_logits = model(val_features, edge_index)
    val_loss = criterion(val_logits, val_labels)
    
    print(f'Epoch {epoch+1}: Train Loss = {train_loss.item()}, Val Loss = {val_loss.item()}')

注意：

上述代码中的图像加载和处理以及标签加载是简化的示例，您需要根据您的实际情况进行适当的修改。
此外，还需要根据您的需求调整模型的超参数和优化器的设置。
此代码仅提供了一个基本的框架，您可能需要根据您的具体任务进行更多的定制。