基于PyTorch和PyG的图神经网络多标签分类任务代码示例

本示例使用PyTorch和PyG库，构建一个GCN网络，实现对图像节点像素特征的降维和多标签分类任务，并使用torch.nn模块中的MultiLabelSoftMarginLoss作为损失函数。代码示例包括数据预处理、模型构建、训练和评估等步骤。

数据说明:

一共有42个时刻的图，而且边的连接关系相同。
每个图都有37个节点，节点特征文件是'C:\Users\jh\Desktop\data\input\images\i.png_j.png'的所有图片的RGB像素值，其中i表示图，i从1到42，j表示节点，j从0到36，特征图片的尺寸为40 x 40。
每个节点有8个标签，储存在'C:\Users\jh\Desktop\data\input\labels\i_j.txt'文本文件中，标签用空格隔开。
边的关系储存在'C:\Users\jh\Desktop\data\input\edges_L.csv'csv文件中，表格中没有header，第一列为源节点，第二列为目标节点，共有61条无向边。

任务目标:

建立一个CNN网络对节点像素特征x进行降维。
将后面3个图的某些节点作为验证集。
使用PYG库建立GCN网络实现多标签分类任务。
使用torch.nn模块中的MultiLabelSoftMarginLoss来实现损失函数。

代码示例:

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

# 定义节点特征数据集类
class NodeFeatureDataset(Dataset):
    def __init__(self, root_dir):
        self.root_dir = root_dir
        self.transform = transforms.Compose([
            transforms.ToTensor()
        ])
    
    def __len__(self):
        return 42
    
    def __getitem__(self, idx):
        img_path = os.path.join(self.root_dir, f"{idx+1}.png")
        image = Image.open(img_path).convert("RGB")
        image = self.transform(image)
        return image

# 定义边关系数据集类
class EdgeDataset(Dataset):
    def __init__(self, csv_file):
        self.edge_df = pd.read_csv(csv_file, header=None)
    
    def __len__(self):
        return len(self.edge_df)
    
    def __getitem__(self, idx):
        src_node = self.edge_df.iloc[idx, 0]
        tgt_node = self.edge_df.iloc[idx, 1]
        return src_node, tgt_node

# 定义GCN模型
class GCN(nn.Module):
    def __init__(self, in_features, hidden_dim, num_classes):
        super(GCN, self).__init__()
        self.conv1 = GCNConv(in_features, hidden_dim)
        self.conv2 = GCNConv(hidden_dim, num_classes)
    
    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index)
        x = torch.relu(x)
        x = F.dropout(x, training=self.training)
        x = self.conv2(x, edge_index)
        return torch.sigmoid(x)

# 定义训练函数
def train(model, device, train_loader, optimizer, criterion):
    model.train()
    running_loss = 0.0
    for data in train_loader:
        data = data.to(device)
        optimizer.zero_grad()
        output = model(data.x, data.edge_index)
        loss = criterion(output[data.train_mask], data.y[data.train_mask])
        loss.backward()
        optimizer.step()
        running_loss += loss.item() * data.num_graphs
    return running_loss / len(train_loader.dataset)

# 定义验证函数
def evaluate(model, device, val_loader, criterion):
    model.eval()
    val_loss = 0.0
    with torch.no_grad():
        for data in val_loader:
            data = data.to(device)
            output = model(data.x, data.edge_index)
            val_loss += criterion(output[data.val_mask], data.y[data.val_mask]).item() * data.num_graphs
    return val_loss / len(val_loader.dataset)

# 设置参数
root_dir = "C:\Users\jh\Desktop\data\input\images"
label_dir = "C:\Users\jh\Desktop\data\input\labels"
edge_file = "C:\Users\jh\Desktop\data\input\edges_L.csv"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
in_features = 3
hidden_dim = 16
num_classes = 8
batch_size = 8
lr = 0.01
epochs = 20
val_nodes = [37, 38, 39]  # 验证集节点索引

# 创建数据加载器
node_feature_dataset = NodeFeatureDataset(root_dir)
edge_dataset = EdgeDataset(edge_file)
train_edges, val_edges = train_test_split(edge_dataset, test_size=0.1, random_state=42)

train_loader = DataLoader(train_edges, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_edges, batch_size=batch_size, shuffle=False)

# 创建GCN模型和优化器
model = GCN(in_features, hidden_dim, num_classes).to(device)
optimizer = optim.Adam(model.parameters(), lr=lr)
criterion = nn.MultiLabelSoftMarginLoss()

# 训练和验证
for epoch in range(epochs):
    train_loss = train(model, device, train_loader, optimizer, criterion)
    val_loss = evaluate(model, device, val_loader, criterion)
    print(f"Epoch {epoch+1}: Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}")

代码说明:

数据预处理:
- NodeFeatureDataset类用于读取图像节点特征数据，使用transforms.ToTensor()将图像转换为张量。
- EdgeDataset类用于读取边关系数据，并返回源节点和目标节点的索引。
- 使用train_test_split函数将边关系数据集分成训练集和验证集。
模型构建:
- GCN类定义了一个两层GCN模型，使用GCNConv层进行图卷积操作。
- forward方法定义了模型的前向传播过程，使用ReLU激活函数和Dropout层进行非线性化。
训练和验证:
- train函数定义了训练过程，使用optimizer更新模型参数，并使用criterion计算损失。
- evaluate函数定义了验证过程，使用验证集评估模型性能。

注意:

请确保将代码中的文件路径更改为您实际的文件路径。
此代码假设您已经安装了PyTorch、PyG和Pandas等必要的库。
您可以根据实际情况调整模型参数、训练参数和验证集的选取。

运行代码:

将代码保存为.py文件。
使用以下命令运行代码：

python your_file.py

代码将输出每个epoch的训练损失和验证损失，帮助您评估模型的训练效果。