gtn-base-v4

GTN-Base v4 (Topology-Disjoint Surrogate Model)

This repository contains the GTN-Base v4 surrogate model trained for the Shared Energy project.

Model Description

The model is a Graph Transformer Network (GTN) that acts as a fast surrogate evaluator. It maps a power converter topology graph $T$ and a target duty cycle $d$ to the predicted efficiency ($\eta$) and normalized output voltage ($V_{\text{out}}$).

Training Details

• Dataset: Cleaned version of 5-component topology dataset (gtn_dataset_5comp_v2_corrected_cleaned.jsonl)
• Split Mode: topology_random (Topology-disjoint split. Train, Val, and Test topologies are completely mutually exclusive.)
• Epochs: 681 (with early stopping, best epoch at 601)
• Batch Size: 32
• Warmup Epochs: 450
• Loss: Relative Squared Error (RSE)

Final Evaluation Metrics (Test Split)

• Efficiency RSE: 0.016024 (1.60%)
• Vout RSE: 0.048169 (4.82%)

Files Included

• gtn_surrogate_model.pt: Model weights (model_state_dict), config, and training metadata.
• manifest.json: Training and evaluation manifest summary.
• training_history.jsonl: Metrics per epoch.
• test_predictions.jsonl: Model predictions on the test set.
• validation_predictions.jsonl: Model predictions on the validation set.

Usage

import torch
from shared_energy.gtn_surrogate import GTNSurrogateModel

checkpoint_path = "gtn_surrogate_model.pt"
payload = torch.load(checkpoint_path, map_location="cpu")
model = GTNSurrogateModel(**payload["model_config"])
model.load_state_dict(payload["model_state_dict"])
model.eval()