GCN训练报错：批量数据索引越界问题的排查与解决

训练GCN模型过程中，按批次选取节点进行训练，节点ID范围为0-999，但原代码直接将节点ID作为索引获取节点相关数据并传入GCN，导致触发索引越界错误，报错信息如下：

RuntimeError: index 332 is out of bounds for dimension 0 with size 64

曾尝试通过节点ID查找其在批次中的索引来获取所需数据，但未找到可行的实现方式，最终通过重构代码解决了问题：构建节点ID到全局索引的映射生成邻接矩阵，改用邻接矩阵与特征矩阵作为GCN输入。

原问题代码 #

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.data import Data
from torch_geometric.nn import GCNConv
from torch_geometric.utils import train_test_split_edges
from torch_geometric.data import DataLoader as loader
import pandas as pd
import torch.optim as optim
from torch.utils.data import Dataset
from torch.nn.utils.rnn import pad_sequence

# feature
class GCN(nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super(GCN, self).__init__()
        self.conv1 = GCNConv(in_channels, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, out_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.conv2(x, edge_index)
        return x

class MyDataset(Dataset):
    def __init__(self, nodes_file, edges_file, communities_file, features_file):
        self.nodes_data = pd.read_csv(nodes_file)
        self.edges_data = pd.read_csv(edges_file)
        self.communities_data = pd.read_csv(communities_file)
        self.features_data = pd.read_csv(features_file)

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

    def __getitem__(self, idx):
        node_id = self.nodes_data.iloc[idx]['ID']
        node_category = self.nodes_data.iloc[idx]['Category']
        node_community = self.communities_data[self.communities_data['ID'] == node_id]['Community'].values[0]
        edge_source = self.nodes_data.iloc[idx]['ID']
        edge_targets = eval(self.edges_data.iloc[idx]['neighbor'])

        # Get the corresponding features for the current node_id
        node_features = self.features_data[self.features_data['NodeID'] == node_id].iloc[:, 1:].values
        node_features = torch.tensor(node_features, dtype=torch.float)

        # You can process the node, edge, and community data as per your requirements
        # and return them as tensors
        return node_id, node_category, node_community, edge_source, edge_targets, node_features


def custom_collate_fn(batch):
    node_ids, node_categories, node_communities, edge_sources, edge_targets, node_features = zip(*batch)

    # Convert to PyTorch tensors
    node_ids = torch.tensor(node_ids, dtype=torch.long)
    node_categories = torch.tensor(node_categories, dtype=torch.long)
    node_communities = torch.tensor(node_communities, dtype=torch.long)
    edge_targets = [torch.tensor(targets) for targets in edge_targets]

    # Replicate edge_sources to have the same length as edge_targets
    edge_sources_replicated = [torch.tensor([source] * len(targets), dtype=torch.long) for source, targets in zip(edge_sources, edge_targets)]

    return node_ids, node_categories, node_communities, edge_sources_replicated, edge_targets, node_features

nodes_file = 'nodes.csv'
edges_file = 'RWR_features.csv'
communities_file = 'community.csv'
feature_file = 'node_features.csv'

dataset = MyDataset(nodes_file, edges_file, communities_file, feature_file)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=64, shuffle=True, collate_fn=custom_collate_fn)


input_dim = 1000  # Replace with the actual input dimension
hidden_dim = 64
output_dim = 32
num_categories = 3  # Replace with the actual number of node categories
num_communities = 5  # Replace with the actual number of communities
num_epochs = 300

# GCN
gcn_model = GCN(input_dim, hidden_dim, output_dim)

# classify
...

# train loop
for epoch in range(num_epochs):
    for batch_data in dataloader:
        node_ids, node_categories, node_communities, edge_sources_list, edge_targets_list, node_feature = batch_data
        
        # print('node_ids:', node_ids)
        # print('len:', len(node_ids))
        # print('node_categories:', node_categories)
        # print('len:', len(node_categories))
        # print('node_communities:', node_communities)
        # print('len:', len(node_communities))
        # print('edge_sources:', edge_sources)
        # print('len:', len(edge_sources))
        # print('edge_targets', edge_targets)
        # print('len:', len(edge_targets))
        
        node_feature = torch.cat(node_feature, dim=0)
        
        # print('node_feature', node_feature)
        # print('len:', len(node_feature))

        
        edge_sources = torch.cat(edge_sources_list, dim=0)
        edge_targets = torch.cat(edge_targets_list, dim=0)
        # print(edge_sources)
        # print(len(edge_sources))
        # print(edge_targets)
        # print(len(edge_targets))

        edge_index = torch.stack([edge_sources, edge_targets], dim=0)

        # print(edge_index)


        x = gcn_model(node_feature, edge_index)

修复后的代码 #

import pandas as pd
import numpy as np

def load_data():
    # Load nodes.csv
    nodes_df = pd.read_csv('nodes.csv')
    
    # Load community.csv
    community_df = pd.read_csv('community.csv')
    
    # Load node_features.csv
    node_features_df = pd.read_csv('node_features.csv')
    
    # Load RWR_features.csv
    rwr_features_df = pd.read_csv('RWR_features.csv', usecols=['node', 'neighbor'])
    rwr_features_df['neighbor'] = rwr_features_df['neighbor'].apply(eval)
    rwr_features_df['neighbor'] = rwr_features_df['neighbor'].apply(lambda x: [int(i) for i in x])
    
    # Create mapping from node ID to index
    node_to_index = {node_id: index for index, node_id in enumerate(nodes_df['ID'])}
    
    # Create adjacency matrix
    num_nodes = len(nodes_df)
    adj = np.zeros((num_nodes, num_nodes))
    for _, row in rwr_features_df.iterrows():
        node_id = row['node']
        neighbor_ids = row['neighbor']
        for neighbor_id in neighbor_ids:
            adj[node_to_index[node_id]][node_to_index[neighbor_id]] = 1
            adj[node_to_index[neighbor_id]][node_to_index[node_id]] = 1  # Assuming undirected graph
    
    # Create feature matrix
    feature_columns = node_features_df.columns[1:]  # Exclude the first column (ID)
    feature = node_features_df[feature_columns].values
    
    # Create class labels
    class_label = nodes_df['Category'].values
    
    # Create community labels
    com_label = community_df['Community'].values
    
    # Split indices for train, val, and test sets (You can modify this split ratio as needed)
    num_train = int(num_nodes * 0.6)
    num_val = int(num_nodes * 0.2)
    num_test = num_nodes - num_train - num_val
    
    # Generate random indices for train, val, and test sets
    indices = np.random.permutation(num_nodes)
    train_idx = indices[:num_train]
    val_idx = indices[num_train:num_train + num_val]
    test_idx = indices[num_train + num_val:]
    
    return adj, feature, class_label, com_label, train_idx, val_idx, test_idx

# Usage
adj, feature, class_label, com_label, train_idx, val_idx, test_idx = load_data()

class MyDataset(Dataset):
    def __init__(self, adj, feature, class_label, com_label):
        self.adj = adj
        self.feature = feature
        self.class_label = class_label
        self.com_label = com_label
        
    def __len__(self):
        return len(self.class_label)
    
    def __getitem__(self, idx):
        return idx, self.adj[idx] , self.feature[idx], self.class_label[idx], self.com_label[idx]

将邻接矩阵与特征矩阵传入GCN：

gcn_feature = gcn_model(batch_adj, batch_feature)

内容的提问来源于stack exchange，提问作者Da-qiong