Graph Neural Network (GCN) for Multi-Class Node Classification with Image Features

This code demonstrates a Graph Convolutional Network (GCN) implementation for multi-class node classification using image features. The dataset consists of 42 time-step graph data, each with 37 nodes represented by 40x40 images. The model leverages edge connections, node features, and labels to predict the classes of nodes.

Dataset Structure:

• Images: Images are stored in 'C:\Users\jh\Desktop\data\input\images\i.png_j.png', where 'i' represents the time-step (1 to 42) and 'j' denotes the node index (0 to 36). Each image is 40x40 pixels in size.
• Labels: Labels for each node are stored in 'C:\Users\jh\Desktop\data\input\labels\i_j.txt', where 'i' and 'j' represent the time-step and node index, respectively. Each file contains 8 labels separated by spaces.
• Edges: Edge connections are stored in 'C:\Users\jh\Desktop\data\input\edges_L.csv'. The first column represents the source node, and the second column represents the target node. Edges are undirected.

Model:

The GCN model uses three layers of GCNConv to learn node representations from the image features and edge connections. The model architecture is defined in the 'GCN' class.

Training and Evaluation:

The code trains the GCN model using the Adam optimizer and the cross-entropy loss function. It splits the dataset into training and validation sets (90% for training, 10% for validation). Training and validation accuracies are reported for each epoch.

Troubleshooting:

The code includes a section addressing potential issues that could lead to zero loss values and low accuracy, including data preprocessing errors, label encoding issues, model architecture problems, training parameters, loss function selection, overfitting, and randomness.

Code:

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 PIL import Image

class MyDataset(torch.utils.data.Dataset):
    def __init__(self, root, transform=None, pre_transform=None):
        self.img_dir = os.path.join(root, 'images')
        self.label_dir = os.path.join(root, 'labels')
        self.edge_file = os.path.join(root, 'edges_L.csv')
        self.transform = transform
        self.pre_transform = pre_transform
        self.dataset = []
        self.create_dataset()

    def create_dataset(self):
        edges = None
        edge_index, num_nodes = self.read_edges(self.edge_file)
        for i in range(1, 43):
            image_list = []
            label_list = []
            for j in range(37):
                image_path = os.path.join(self.img_dir, f'{i}.png_{j}.png')
                label_path = os.path.join(self.label_dir, f'{i}_{j}.txt')
                features = self.read_image_features(image_path)
                labels = self.read_labels(label_path)
                labels = torch.tensor(labels, dtype=torch.float)
                features = torch.tensor(features).unsqueeze(0)
                features = features.float()
                data = Data(x=features, edge_index=edge_index, y=labels)
                if j < 30:
                    data.train_mask = torch.tensor([True] * num_nodes, dtype=torch.bool)
                    data.val_mask = torch.tensor([False] * num_nodes, dtype=torch.bool)
                else:
                    data.train_mask = torch.tensor([False] * num_nodes, dtype=torch.bool)
                    data.val_mask = torch.tensor([True] * num_nodes, dtype=torch.bool)
                image_list.append(data)
                label_list.append(data)

        self.dataset.extend(image_list)
        return self.dataset, edges

    def read_edges(self, edge_path):
        edges = []
        with open(edge_path, 'r') as file:
            for line in file:
                src, tgt = line.strip().split(',')
                edges.append((int(src), int(tgt)))

        max_node_idx = max(max(edges, key=lambda x: max(x)))
        num_nodes = max_node_idx + 1
        edge_index = torch.tensor(edges, dtype=torch.long).t().contiguous()
        return edge_index, num_nodes

    def read_image_features(self, image_path):
        img = Image.open(image_path)
        img = img.resize((40, 40))
        rgb_img = img.convert('RGB')
        features = []

        for i in range(40):
            for j in range(40):
                r, g, b = rgb_img.getpixel((i, j))
                features.append([r, g, b])

        return features

    def read_labels(self, label_path):
        with open(label_path, 'r') as file:
            labels = [int(label) for label in file.read().strip().split()]

        return labels

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

    def __getitem__(self, idx):
        data = self.dataset[idx]

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

        return data

class GCN(torch.nn.Module):
    def __init__(self, num_node_features, num_classes):
        super(GCN, self).__init__()
        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
        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)
        return x

def train_model(data_loader, model, optimizer, device):
    model.train()
    train_loss = 0
    for data in data_loader:
        data = data.to(device)
        optimizer.zero_grad()
        output = model(data).view(-1)
        train_mask = data.train_mask
        labels = data.y
        train_indices = torch.nonzero(train_mask, as_tuple=True)[0]

        if len(train_indices) == 0:
            continue

        max_index = torch.max(train_indices)
        if max_index >= len(output) or max_index >= len(labels):
            continue

        loss = F.cross_entropy(output[train_indices], labels[train_indices])
        loss.backward()
        optimizer.step()
        train_loss += loss.item() * data.num_graphs

    train_loss /= len(data_loader.dataset)
    return train_loss

def validate_model(data_loader, model, device):
    model.eval()
    correct = 0
    total = 0
    for data in data_loader:
        data = data.to(device)
        output = model(data).view(-1)
        val_mask = data.val_mask
        labels = data.y
        val_indices = torch.nonzero(val_mask, as_tuple=True)[0]

        if len(val_indices) == 0:
            continue

        val_indices = val_indices[val_indices < len(labels)]  # Exclude out-of-bounds indices
        predicted = output[val_indices].argmax()
        total += len(val_indices)
        correct += (predicted == labels[val_indices]).sum().item()

    val_accuracy = correct / total
    return val_accuracy

if __name__ == '__main__':
    dataset = MyDataset(root="C:\Users\jh\Desktop\data\input")
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = GCN(num_node_features=3, num_classes=8).to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
    train_dataset, val_dataset = train_test_split(dataset, test_size=0.1)
    train_loader = DataLoader(train_dataset, batch_size=1, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=1, shuffle=False)
    epochs = 10
    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}')

Explanation:

• The 'MyDataset' class loads the data from the specified directories and creates 'Data' objects for each node. The 'read_image_features' and 'read_labels' functions extract features and labels from the images and text files, respectively.
• The 'GCN' class defines the model architecture using three layers of GCNConv.
• The 'train_model' function trains the model for one epoch using the specified data loader, optimizer, and device.
• The 'validate_model' function evaluates the model on the validation set and returns the accuracy.
• The main script instantiates the dataset, model, and optimizer. It then trains the model for the specified number of epochs and reports the training loss and validation accuracy for each epoch.

Possible Issues:

1. Data Preprocessing:
   • Verify that the specified paths for images, labels, and edges are correct.
   • Ensure that the 'read_image_features' and 'read_labels' functions are reading the data correctly and converting it to the appropriate format.
2. Label Encoding:
   • Check that the labels are encoded as integers or floats compatible with the number of classes in the model.
3. Model Architecture:
   • Verify that the number of input features ('num_node_features') and output classes ('num_classes') in the 'GCN' class match the data.
   • Adjust the number and size of the GCNConv layers based on the complexity of the dataset.
4. Training Parameters:
   • Experiment with different learning rates, batch sizes, and number of epochs to find optimal values.
5. Loss Function:
   • If the labels are not one-hot encoded, use 'F.cross_entropy' for multi-class classification.
6. Overfitting:
   • Add regularization techniques like dropout or weight decay to prevent overfitting.
   • Increase the size of the validation set or use cross-validation to assess model performance.
7. Randomness:
   • Set a random seed to ensure reproducibility and compare performance across different runs.