dataset-profiling.py

import matplotlib.pyplot as plt
import numpy as np
import os
from tgb.linkproppred.dataset import LinkPropPredDataset

def profile_datasets():
    datasets = ["tgbl-wiki", "tgbl-subreddit", "tgbl-flight"]
    
    names = []
    struct_sizes = []
    feat_sizes = []

    print("==================================================")
    print(" Temporal Graph Memory Profiler")
    print("==================================================\n")

    for name in datasets:
        print(f"--- Profiling {name} ---")
        try:
            dataset = LinkPropPredDataset(name=name, root="datasets")
            data = dataset.full_data
            
            # Feature Footprint
            edge_feat = data['edge_feat']
            feat_mb = edge_feat.nbytes / (1024 * 1024) if edge_feat is not None else 0
            
            # Structural Footprint (Sources, Destinations, Timestamps)
            src_mb = data['sources'].nbytes / (1024 * 1024)
            dst_mb = data['destinations'].nbytes / (1024 * 1024)
            ts_mb = data['timestamps'].nbytes / (1024 * 1024)
            struct_mb = src_mb + dst_mb + ts_mb
            
            print(f"Feature RAM:   {feat_mb:.2f} MB")
            print(f"Structure RAM: {struct_mb:.2f} MB\n")
            
            names.append(name)
            struct_sizes.append(struct_mb)
            feat_sizes.append(feat_mb)
            
        except Exception as e:
            print(f"Error loading {name}: {e}\n")

    # Generate Stacked Bar Chart
    plt.rcParams.update({'font.size': 12, 'font.family': 'serif'})
    fig, ax = plt.subplots(figsize=(8, 6))

    x = np.arange(len(names))
    width = 0.5

    ax.bar(x, struct_sizes, width, label='Structural Data (IDs, Time)', color='#c44e52')
    ax.bar(x, feat_sizes, width, bottom=struct_sizes, label='Continuous Features', color='#4c72b0')

    ax.set_ylabel('Memory Footprint (MB)')
    ax.set_title('Resting Dataset Anatomy Prior to Quantization')
    ax.set_xticks(x)
    ax.set_xticklabels(names)
    ax.legend()

    # Use a logarithmic scale if Subreddit crushes Wiki too much to see
    ax.set_yscale('log')
    ax.set_ylabel('Memory Footprint (MB) - Log Scale')

    os.makedirs("results/plots", exist_ok=True)
    plot_path = "results/plots/dataset_anatomy.png"
    fig.tight_layout()
    fig.savefig(plot_path, dpi=300)
    print(f"Saved dataset anatomy plot to {plot_path}")

if __name__ == "__main__":
    profile_datasets()