PyTorch Geometric 图神经网络模型训练：基于图像特征的节点分类

本项目使用 PyTorch Geometric 框架，构建了一个 GCN 模型，并结合 CNN 模型提取图像特征，用于节点分类任务。代码包含数据加载、模型构建、训练和验证等部分。

数据加载和 PyG 数据集类

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 sklearn.model_selection import train_test_split
from torchvision import transforms
from PIL import Image

# 加载数据并创建 PyG 数据集类：
class MyDataset(torch.utils.data.Dataset):
    def __init__(self, root, transform=None, pre_transform=None):
        self.edges = pd.read_csv(os.path.join(root, 'input', 'edges_L.csv'), header=None)
        self.transform = transform
        self.pre_transform = pre_transform

        # 读取特征和标签数据
        self.features = []
        self.labels = []
        for i in range(1, 43):
            for j in range(37):
                # 读取特征
                img_name = os.path.join(root, 'input', 'images', '{}.png_{}.png'.format(i, j))
                img = Image.open(img_name).convert('RGB').resize((224, 224), resample=Image.BILINEAR)
                img_tensor = transforms.ToTensor()(img)
                self.features.append(img_tensor)

                # 读取标签
                label_name = os.path.join(root, 'input', 'labels', '{}_{}.txt'.format(i, j))
                with open(label_name, 'r') as f:
                    labels = [int(x) for x in f.readline().strip().split()]
                self.labels.append(labels)

        # 将特征调整维度为[batch_size, num_node_features, width, height]
        self.features = torch.stack(self.features, dim=0)
        self.labels = torch.tensor(self.labels)

        # Calculate the total number of nodes
        self.num_nodes = len(self.labels)

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        network_id = idx // 37  # Calculate the network index
        node_id = idx % 37  # Calculate the node index within the network

        # Determine the edges for the current network
        network_edges = self.edges[self.edges[0].isin(range(network_id * 37, (network_id + 1) * 37)) & self.edges[1].isin(range(network_id * 37, (network_id + 1) * 37))]

        # Construct the edge_index tensor for undirected edges
        source_nodes = network_edges[0].values
        target_nodes = network_edges[1].values
        edge_index = torch.tensor([source_nodes, target_nodes], dtype=torch.long)

        x = self.features[idx]  # 获取节点特征
        y = self.labels[idx]  # 获取标签

        # Define graph-wide train_mask and val_mask
        train_mask = torch.zeros(self.num_nodes, dtype=torch.bool)
        val_mask = torch.zeros(self.num_nodes, dtype=torch.bool)

        # Set train_mask for the first 30 nodes in each network, and val_mask for the last 7 nodes
        if node_id < 30:
            train_mask[node_id + network_id * 37] = 1
        else:
            val_mask[node_id + network_id * 37] = 1

        data = Data(x=x, edge_index=edge_index, y=y, train_mask=train_mask, val_mask=val_mask)

        if self.pre_transform is not None:
            data = self.pre_transform(data)
        if self.transform is not None:
            data.x = self.transform(data.x)
        return data

模型构建

CNN 模型

# 定义CNN模型
class CNN(torch.nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1)
        self.fc = nn.Linear(32 * 56 * 56, 40 * 40 * 3)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 32 * 56 * 56)
        x = self.fc(x)
        return x

GCN 模型

# 创建GCN模型
class GCN(torch.nn.Module):
    def __init__(self, num_node_features, num_classes):
        super(GCN, self).__init__()

        self.cnn = CNN()  # 添加CNN模型

        self.conv1 = GCNConv(num_node_features, 8)
        self.conv2 = GCNConv(8, 16)
        self.conv3 = GCNConv(16, num_classes)

    def forward(self, data):
        x, edge_index = data.x, data.edge_index
        print (x.shape)
        print (edge_index.shape)
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = self.conv2(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, training=self.training)
        x = self.conv3(x, edge_index)
        x = F.softmax(x, dim=1)
        x = x.view(-1, 37, -1)  # 调整输出维度
        return x

训练和验证模型

# 创建训练和验证模型
def train_model(dataset, model, optimizer, device):
    model.train()
    total_loss = 0.0

    for data in dataset:
        data = data.to(device)
        optimizer.zero_grad()

        features = model.cnn(data.x)  # 提取图像特征
        data.x = features.view(features.size(0), -1)

        output = model(data)
        loss = F.cross_entropy(output[data.train_mask], data.y[data.train_mask])
        loss.backward()
        optimizer.step()
        total_loss += loss.item()

    return total_loss / len(dataset)


def validate_model(dataset, model, device):
    model.eval()
    correct = 0
    total = 0

    for data in dataset:
        data = data.to(device)

        features = model.cnn(data.x)  # 提取图像特征
        data.x = features.view(features.size(0), -1)

        output = model(data)
        _, predicted = torch.max(output[data.val_mask], 1)
        total += data.val_mask.sum().item()
        correct += (predicted == data.y[data.val_mask]).sum().item()

    return correct / total

if __name__ == '__main__':
    dataset = MyDataset(root="C:\Users\jh\Desktop\data")
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    cnn_model = CNN().to(device)
    model = GCN(num_node_features=40 * 40 * 3, num_classes=8).to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

    train_dataset, val_dataset = train_test_split(dataset, test_size=0.1)
    train_loader = DataLoader(train_dataset, batch_size=1, shuffle=False)
    val_loader = DataLoader(val_dataset, batch_size=1, shuffle=False)

    epochs = 2
    for epoch in range(epochs):
        train_loss = train_model(train_loader, model, optimizer, device)
        print(f'Epoch {epoch+1}/{epochs}, Train Loss: {train_loss:.4f}')

        val_accuracy = validate_model(val_loader, model, device)
        print(f'Val_Acc: {val_accuracy:.4f}')

错误信息分析

根据代码和错误信息来看，可能是因为数据集中的边的索引存在问题。在代码中，网络图的边是根据节点的索引计算得到的，而且每个网络图中有 37 个节点。但是根据错误信息，边的索引的维度为 [2, 0]，这说明没有正确地计算出边的索引。

你可以检查以下几个问题：

确保 edges_L.csv 文件中的边的索引是正确的，并且包含了所有的边。
确保 network_edges 的计算是正确的。你可以在 __getitem__ 函数中打印出 network_edges 的值，检查它是否包含了正确的边的索引。
确保 edge_index 的维度是正确的。你可以在 __getitem__ 函数中打印出 edge_index 的维度，检查它是否为 [2, num_edges] 的形状。

如果你仔细检查了以上几点仍然无法解决问题，可以尝试打印出 network_edges 的值，以及 source_nodes 和 target_nodes 的值，看是否符合预期。另外，确保你的数据集中的边的索引是从 0 开始的。

希望这些提示能帮助你解决问题！