基于PYG库的GCN多标签分类任务代码示例

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
from torch_geometric.utils import to_networkx

# 定义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 torch.sigmoid(x)


# 构建数据集
class GraphDataset(torch.utils.data.Dataset):
    def __init__(self, images_dir, labels_dir, edges_file):
        self.images_dir = images_dir
        self.labels_dir = labels_dir
        self.edges_file = edges_file

        self.num_graphs = 42
        self.num_nodes = 37
        self.num_labels = 8

        self.graphs = self.load_graphs()
        self.labels = self.load_labels()
        self.edges = self.load_edges()

    def load_graphs(self):
        graphs = []
        for i in range(1, self.num_graphs + 1):
            graph = self.load_graph(i)
            graphs.append(graph)
        return graphs

    def load_graph(self, i):
        node_features = []
        for j in range(self.num_nodes):
            image_path = f'{self.images_dir}/i{j+1}.png'
            image = self.load_image(image_path)
            node_features.append(image)
        node_features = torch.stack(node_features)
        return Data(x=node_features)

    def load_image(self, image_path):
        # 从图片路径加载RGB像素值
        # 这里需要根据你的具体数据加载方式进行实现
        return image

    def load_labels(self):
        labels = []
        for i in range(1, self.num_graphs + 1):
            graph_labels = self.load_graph_labels(i)
            labels.append(graph_labels)
        return labels

    def load_graph_labels(self, i):
        labels_path = f'{self.labels_dir}/{i}.txt'
        with open(labels_path, 'r') as f:
            graph_labels = f.readline().split()
        graph_labels = torch.tensor([float(label) for label in graph_labels])
        return graph_labels

    def load_edges(self):
        edges_path = f'{self.edges_file}'
        with open(edges_path, 'r') as f:
            edges = []
            for line in f:
                node1, node2 = line.split(',')
                edges.append([int(node1), int(node2)])
        return torch.tensor(edges)

    def __len__(self):
        return self.num_graphs

    def __getitem__(self, idx):
        graph = self.graphs[idx]
        labels = self.labels[idx]
        return graph, labels, self.edges


# 定义训练函数
def train(model, device, train_loader, optimizer, criterion):
    model.train()

    for data in train_loader:
        data = data.to(device)
        optimizer.zero_grad()
        output = model(data.x, data.edge_index)
        loss = criterion(output, data.y)
        loss.backward()
        optimizer.step()


# 定义测试函数
def test(model, device, test_loader):
    model.eval()

    correct = 0
    total = 0

    with torch.no_grad():
        for data in test_loader:
            data = data.to(device)
            output = model(data.x, data.edge_index)
            predicted_labels = torch.round(output)
            total += data.y.size(0)
            correct += (predicted_labels == data.y).sum().item()

    accuracy = 100 * correct / total
    print('Test Accuracy: {:.2f}%'.format(accuracy))


# 设置训练参数
input_dim = 40 * 40 * 3  # 输入特征维度
hidden_dim = 64  # 隐藏层维度
output_dim = 8  # 输出维度
lr = 0.01  # 学习率
epochs = 50  # 迭代次数
batch_size = 8  # 批次大小
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  # 检查是否有GPU可用

# 加载数据集
dataset = GraphDataset('C:/Users/jh/Desktop/data/input/images', 'C:/Users/jh/Desktop/data/input/labels', 'C:/Users/jh/Desktop/data/input/edges_L.csv')
train_dataset = dataset[:38]
test_dataset = dataset[38:]

# 创建数据加载器
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

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

# 开始训练
for epoch in range(epochs):
    train(model, device, train_loader, optimizer, criterion)
    test(model, device, test_loader)

请根据自己的实际数据和环境进行必要的修改，尤其是加载图片和实现特征的部分。