train_all_motif_finders.py

import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import GCNConv, global_mean_pool
from torch_geometric.data import Data, DataLoader
from torch_geometric.data import Batch
import numpy as np
from typing import Tuple
import argparse
import time
from tqdm import tqdm

from subgraph_dataset import SubgraphViewDataset
from centrality_utils import compute_centrality_and_cse
from feature_enhancer import FeatureEnhancer
from build_nano_db import load_node_dataset
from utils import process
from preprompt import pca_compression
import scipy.sparse as sp


class SubgraphEncoder(nn.Module):
    
    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int):
        super(SubgraphEncoder, self).__init__()
        
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.output_dim = output_dim
        
        self.conv1 = GCNConv(input_dim, hidden_dim)
        self.conv2 = GCNConv(hidden_dim, output_dim)
        
    def forward(self, data: Data) -> torch.Tensor:

        x, edge_index, batch = data.x, data.edge_index, data.batch
        
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        
        x = self.conv2(x, edge_index)
        
        graph_embedding = global_mean_pool(x, batch)
        
        return graph_embedding


class MotifContrastiveModel(nn.Module):
    """

    """
    
    def __init__(self, struct_input_dim: int, semantic_input_dim: int, 
                 hidden_dim: int, output_dim: int):
        super(MotifContrastiveModel, self).__init__()
        
        self.struct_input_dim = struct_input_dim
        self.semantic_input_dim = semantic_input_dim
        self.hidden_dim = hidden_dim
        self.output_dim = output_dim
        
        self.struct_encoder = SubgraphEncoder(
            input_dim=struct_input_dim,
            hidden_dim=hidden_dim,
            output_dim=output_dim
        )
        
        self.semantic_encoder = SubgraphEncoder(
            input_dim=semantic_input_dim,
            hidden_dim=hidden_dim,
            output_dim=output_dim
        )
        
    def forward(self, struct_view_batch: Data, semantic_view_batch: Data) -> torch.Tensor:
        """

        """
        h_struct = self.struct_encoder(struct_view_batch)
        h_semantic = self.semantic_encoder(semantic_view_batch)
        
        loss = self.calculate_infonce_loss(h_struct, h_semantic)
        
        return loss
    
    def calculate_infonce_loss(self, z1: torch.Tensor, z2: torch.Tensor, 
                              temperature: float = 0.1) -> torch.Tensor:
        """

        """
        batch_size = z1.size(0)
        
        z1 = F.normalize(z1, dim=1)
        z2 = F.normalize(z2, dim=1)
        
        sim_matrix = torch.matmul(z1, z2.T) / temperature
        
        labels = torch.arange(batch_size, device=z1.device)
        
        loss = F.cross_entropy(sim_matrix, labels)
        
        return loss


def prepare_dataset_data(dataset_name: str, args: argparse.Namespace, 
                        enhancer: FeatureEnhancer) -> Tuple[Data, torch.Tensor, torch.Tensor]:

    dataset_wrapper = load_node_dataset(root='data', name=dataset_name, args=args)
    data = dataset_wrapper.data
    
    batched_data = Batch.from_data_list([data])
    raw_features_np, adj = process.process_tu(batched_data, batched_data.x.shape[1])
    
    unified_raw_features_np = pca_compression(raw_features_np, k=50)
    
    enhanced_features_tensor = enhancer.enhance(unified_raw_features_np, data.raw_texts)
    
    data.x = enhanced_features_tensor

    K_STEPS = list(range(1, 9))  
    TOP_K_NODES = min(200, data.num_nodes // 10)  
    
    try:
        results = compute_centrality_and_cse(
            edge_index=data.edge_index,
            num_nodes=data.num_nodes,
            ksteps=K_STEPS,
            k=TOP_K_NODES
        )
        
        cse_encodings = results['cse_encodings']
        top_k_indices = results['top_k_indices']
        
    except Exception as e:
        print(f"    ✗ CSE calculation failed: {e}")

        cse_encodings = torch.randn(data.num_nodes, len(K_STEPS))
        top_k_indices = torch.randperm(data.num_nodes)[:TOP_K_NODES]
        
    
    return data, cse_encodings, top_k_indices


def train_motif_encoder(dataset_name: str, data: Data, cse_features: torch.Tensor, 
                       top_k_indices: torch.Tensor, args: argparse.Namespace) -> str:
    
    dataset = SubgraphViewDataset(
        full_graph_data=data,
        cse_features=cse_features,
        top_k_node_indices=top_k_indices
    )
    
    dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True)
   
    struct_input_dim = cse_features.shape[1]  
    semantic_input_dim = 100  
    
    model = MotifContrastiveModel(
        struct_input_dim=struct_input_dim,
        semantic_input_dim=semantic_input_dim,
        hidden_dim=args.hidden_dim,
        output_dim=args.output_dim
    )
    
    optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate)
    
    patience = getattr(args, 'patience', 20)  
    min_delta = getattr(args, 'min_delta', 1e-4)  
    best_loss = float('inf')
    patience_counter = 0
    best_epoch = 0
    
    model.train()
    
    for epoch in range(args.epochs):
        total_loss = 0.0
        num_batches = 0
        
        for batch_idx, (struct_view_batch, semantic_view_batch) in enumerate(dataloader):
            loss = model(struct_view_batch, semantic_view_batch)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()
            num_batches += 1

            if batch_idx % 10 == 0:
                print(f"    Epoch {epoch+1}/{args.epochs}, Batch {batch_idx+1}/{len(dataloader)}, "
                      f"Loss: {loss.item():.4f}")
        
        avg_loss = total_loss / num_batches
        print(f"    ✓ Epoch {epoch+1} completed, average loss: {avg_loss:.4f}")
        
        if avg_loss < best_loss - min_delta:
            best_loss = avg_loss
            best_epoch = epoch + 1
            patience_counter = 0

            output_dir = os.path.join(args.motif_lib_path, dataset_name)
            os.makedirs(output_dir, exist_ok=True)

            struct_encoder_state = model.struct_encoder.state_dict()

            encoder_path = os.path.join(output_dir, 'encoder.pth')
            torch.save(struct_encoder_state, encoder_path)

            config = {
                'struct_input_dim': struct_input_dim,
                'hidden_dim': args.hidden_dim,
                'output_dim': args.output_dim,
                'dataset_name': dataset_name,
                'num_epochs': epoch + 1,  
                'best_epoch': best_epoch,
                'best_loss': best_loss,
                'learning_rate': args.learning_rate,
                'early_stopped': False
            }
            
            config_path = os.path.join(output_dir, 'config.pth')
            torch.save(config, config_path)
            
        else:
            patience_counter += 1

            if patience_counter >= patience:
                config_path = os.path.join(args.motif_lib_path, dataset_name, 'config.pth')
                if os.path.exists(config_path):
                    config = torch.load(config_path)
                    config['early_stopped'] = True
                    config['stopped_epoch'] = epoch + 1
                    torch.save(config, config_path)
                
                break
    
    if patience_counter < patience:
        print(f"    ✅ training completed! best epoch: {best_epoch}, best loss: {best_loss:.4f}")
    
    encoder_path = os.path.join(args.motif_lib_path, dataset_name, 'encoder.pth')
    if not os.path.exists(encoder_path):
        output_dir = os.path.join(args.motif_lib_path, dataset_name)
        os.makedirs(output_dir, exist_ok=True)
        
        struct_encoder_state = model.struct_encoder.state_dict()
        
        torch.save(struct_encoder_state, encoder_path)
        
        config = {
            'struct_input_dim': struct_input_dim,
            'hidden_dim': args.hidden_dim,
            'output_dim': args.output_dim,
            'dataset_name': dataset_name,
            'num_epochs': epoch + 1,
            'best_epoch': best_epoch,
            'best_loss': best_loss,
            'learning_rate': args.learning_rate,
            'early_stopped': patience_counter >= patience
        }
        
        config_path = os.path.join(output_dir, 'config.pth')
        torch.save(config, config_path)
    
    return encoder_path


def main():
    parser = argparse.ArgumentParser(description='train all motif encoders')
    parser.add_argument('--pretrain_datasets', nargs='+', 
                       default=['Cora','Citeseer','Pubmed','home','wikics'],
                       help='list of pretrain datasets')
    parser.add_argument('--motif_lib_path', type=str, default='./motif_lib/',
                       help='root directory of motif lib')
    parser.add_argument('--epochs', type=int, default=10000,
                       help='number of epochs')
    parser.add_argument('--batch_size', type=int, default=64,
                       help='batch size')
    parser.add_argument('--learning_rate', type=float, default=0.0001,
                       help='learning rate')
    parser.add_argument('--hidden_dim', type=int, default=64,
                       help='hidden dimension')
    parser.add_argument('--output_dim', type=int, default=32,
                       help='output dimension')
    parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu',
                       help='compute device')
    parser.add_argument('--seed', type=int, default=42,
                       help='random seed')
    
    args = parser.parse_args()
    
    torch.manual_seed(args.seed)
    np.random.seed(args.seed)
    
    enhancer = FeatureEnhancer(
        text_encoder='bert',
        text_output_dim=50  
    )
    
    os.makedirs(args.motif_lib_path, exist_ok=True)
    
    for i, dataset_name in enumerate(args.pretrain_datasets, 1):
        print(f"start to train motif encoder for dataset '{dataset_name}' ({i}/{len(args.pretrain_datasets)})")
        data, cse_features, top_k_indices = prepare_dataset_data(dataset_name, args, enhancer)
        encoder_path = train_motif_encoder(dataset_name, data, cse_features, top_k_indices, args)
        print(f"✅ motif encoder for dataset '{dataset_name}' trained successfully!")

if __name__ == "__main__":
    main()