圖神經網路（Graph Neural Networks, GNNs）完整入門指南

import torch
import torch.nn as nn
import torch.nn.functional as F
# Ensure you have torch_geometric installed: pip install torch_geometric
try:
    from torch_geometric.nn import GCNConv
    from torch_geometric.datasets import Planetoid
    from torch_geometric.transforms import NormalizeFeatures
except ImportError:
    print("PyTorch Geometric not found. Please install it: pip install torch_geometric")
    # You might also need to install dependencies based on your PyTorch+CUDA version
    # See: https://pytorch-geometric.readthedocs.io/en/latest/install/installation.html
    exit()

# --- 1. Load Dataset ---
try:
    # Load a standard citation network dataset (Cora)
    # NormalizeFeatures standardizes node features to have zero mean and unit variance
    dataset = Planetoid(root='./data/Cora', name='Cora', transform=NormalizeFeatures())
    data = dataset[0]  # Get the single graph object from the dataset
except Exception as e:
    print(f"Error loading dataset: {e}")
    print("Please ensure the dataset can be downloaded or accessed.")
    exit()

# --- 2. Print Dataset Information ---
print(f'Dataset: {dataset}:')
print('===================')
print(f'Number of graphs: {len(dataset)}')
print(f'Number of nodes: {data.num_nodes}') # Use data.num_nodes
print(f'Number of node features: {dataset.num_node_features}') # Use dataset.num_node_features
print(f'Number of edge features: {dataset.num_edge_features}') # Use dataset.num_edge_features (usually 0 for Cora)
print(f'Number of classes: {dataset.num_classes}')
print(f'Number of edges: {data.num_edges}') # Use data.num_edges
print(f'Training nodes: {data.train_mask.sum().item()}') # Show count of training nodes
print(f'Validation nodes: {data.val_mask.sum().item()}')   # Show count of validation nodes
print(f'Test nodes: {data.test_mask.sum().item()}')     # Show count of test nodes
print(f'Graph has isolated nodes: {data.has_isolated_nodes()}')
print(f'Graph has self-loops: {data.has_self_loops()}')
print(f'Graph is undirected: {data.is_undirected()}')
print('=============================================================')

# --- 3. Define the GCN Model ---
class GCN(nn.Module):
    """
    A simple two-layer Graph Convolutional Network.
    """
    def __init__(self, in_channels, hidden_channels, out_channels):
        super(GCN, self).__init__()
        torch.manual_seed(12345) # for reproducibility
        # First Graph Convolutional layer: maps node features to hidden dimensions
        self.conv1 = GCNConv(in_channels, hidden_channels)
        # Second Graph Convolutional layer: maps hidden dimensions to output classes
        self.conv2 = GCNConv(hidden_channels, out_channels)

    def forward(self, x, edge_index):
        """
        Defines the forward pass of the GCN.
        Args:
            x (Tensor): Node feature matrix [num_nodes, in_channels]
            edge_index (LongTensor): Graph connectivity in COO format [2, num_edges]
        Returns:
            Tensor: Log-softmax probabilities for each node [num_nodes, out_channels]
        """
        # --- Layer 1 ---
        # Apply graph convolution
        x = self.conv1(x, edge_index)
        # Apply ReLU activation function
        x = F.relu(x)
        # Apply dropout for regularization (only during training)
        x = F.dropout(x, p=0.5, training=self.training)

        # --- Layer 2 ---
        # Apply second graph convolution
        x = self.conv2(x, edge_index)

        # --- Output ---
        # Apply log_softmax for classification (works well with NLLLoss)
        return F.log_softmax(x, dim=1)

# --- 4. Initialize Model and Optimizer ---
model = GCN(in_channels=dataset.num_node_features,
            hidden_channels=16, # A common choice for hidden layer size
            out_channels=dataset.num_classes)

print("nModel Architecture:")
print(model)
print('=============================================================')

# Optimizer: Adam is a popular choice
# weight_decay adds L2 regularization
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)

# --- 5. Training Function ---
def train():
    """
    Performs a single training step.
    """
    model.train() # Set model to training mode (enables dropout)
    optimizer.zero_grad() # Clear gradients from previous iteration

    # Perform forward pass
    out = model(data.x, data.edge_index)

    # Calculate loss using Negative Log Likelihood Loss
    # Only calculate loss on the nodes in the training set (using train_mask)
    loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])

    # Perform backward pass: compute gradients
    loss.backward()

    # Update model parameters
    optimizer.step()

    return loss.item() # Return the loss value for this step

# --- 6. Testing Function ---
def test(mask):
    """
    Evaluates the model on a given node mask (train, val, or test).
    Args:
        mask (BoolTensor): Mask indicating which nodes to evaluate.
    Returns:
        float: Accuracy on the specified node set.
    """
    model.eval() # Set model to evaluation mode (disables dropout)
    with torch.no_grad(): # Disable gradient calculation for efficiency
        out = model(data.x, data.edge_index)
        # Get predictions by finding the class with the highest log-probability
        pred = out.argmax(dim=1)
        # Compare predictions with true labels for the nodes in the mask
        correct = pred[mask] == data.y[mask]
        # Calculate accuracy
        acc = int(correct.sum()) / int(mask.sum())
    return acc

# --- 7. Training Loop ---
print("Starting training...")
for epoch in range(1, 201):
    loss = train()
    if epoch % 10 == 0:
        # Evaluate on training, validation, and test sets
        train_acc = test(data.train_mask)
        val_acc = test(data.val_mask) # Also test on validation set
        test_acc = test(data.test_mask)
        print(f'Epoch: {epoch:03d}, Loss: {loss:.4f}, Train Acc: {train_acc:.4f}, Val Acc: {val_acc:.4f}, Test Acc: {test_acc:.4f}')
print("Training finished.")
print('=============================================================')

# --- 8. Final Evaluation ---
final_train_acc = test(data.train_mask)
final_val_acc = test(data.val_mask)
final_test_acc = test(data.test_mask)
print(f'Final Train Accuracy: {final_train_acc:.4f}')
print(f'Final Validation Accuracy: {final_val_acc:.4f}')
print(f'Final Test Accuracy: {final_test_acc:.4f}')
print('=============================================================')

# --- 9. Examine Node Embeddings (Optional) ---
print("nExamining node embeddings...")
model.eval() # Ensure model is in eval mode
with torch.no_grad():
    # --- Get embeddings after the first layer (before dropout and second layer) ---
    # Re-run the first part of the forward pass to get intermediate embeddings
    hidden_embeddings = model.conv1(data.x, data.edge_index)
    hidden_embeddings = F.relu(hidden_embeddings) # Apply activation

    print(f"nShape of hidden embeddings: {hidden_embeddings.shape}") # [num_nodes, hidden_channels]

    # Look at the embedding of the first few nodes
    print("nSample node embeddings after first GCN layer + ReLU:")
    print(hidden_embeddings[:3])

    # --- Analyze similarity between connected nodes ---
    if data.num_edges > 0:
        # Find a pair of connected nodes (using the first edge in the list)
        edge_index_sample = data.edge_index[:, 0]  # Take the first edge [source_node, target_node]
        node1_idx = edge_index_sample[0].item()
        node2_idx = edge_index_sample[1].item()

        print(f"nComparing embeddings for connected nodes {node1_idx} and {node2_idx}:")

        # Get their embeddings from the hidden layer
        emb1 = hidden_embeddings[node1_idx]
        emb2 = hidden_embeddings[node2_idx]

        # Calculate Cosine Similarity
        # unsqueeze(0) adds a batch dimension, required by cosine_similarity
        similarity = F.cosine_similarity(emb1.unsqueeze(0), emb2.unsqueeze(0))

        print(f"Embedding Node {node1_idx}: {emb1.numpy().round(2)}") # Print rounded numpy array
        print(f"Embedding Node {node2_idx}: {emb2.numpy().round(2)}")
        print(f"Cosine similarity: {similarity.item():.4f}")

        # --- Analyze similarity between potentially unconnected nodes (Optional) ---
        # Find two nodes that are likely not connected (e.g., node 0 and node 100)
        if data.num_nodes > 100:
             node3_idx = 0
             node4_idx = 100
             print(f"nComparing embeddings for potentially unconnected nodes {node3_idx} and {node4_idx}:")
             emb3 = hidden_embeddings[node3_idx]
             emb4 = hidden_embeddings[node4_idx]
             similarity_unconnected = F.cosine_similarity(emb3.unsqueeze(0), emb4.unsqueeze(0))
             print(f"Embedding Node {node3_idx}: {emb3.numpy().round(2)}")
             print(f"Embedding Node {node4_idx}: {emb4.numpy().round(2)}")
             print(f"Cosine similarity: {similarity_unconnected.item():.4f}")

    else:
        print("nGraph has no edges, cannot compare connected node embeddings.")

print("nAnalysis complete.")