Building a TGraphX Experiment from a YAML Config
Research experiments that live entirely in Python scripts tend to accumulate hardcoded hyperparameters, dataset paths, and architecture choices scattered through the code. Changing the number of hidden layers requires editing the model class. Changing the learning rate requires finding where it is hardcoded in the optimizer call. After a few weeks of iteration, the script becomes a tangle of commented-out configurations and is difficult to share or reproduce.
Config-driven experiment design solves this by externalizing all choices that vary between experiments into a YAML file. The Python code becomes a generic runner that builds the model, dataset, and training loop from whatever the config specifies. This article shows how to design a YAML config structure for GNN experiments and wire it up with TGraphX.
What This Builds On
This tutorial assumes familiarity with TGraphX's layers and basic training loops. The GNN research reproducibility article explains why reproducibility infrastructure matters and how TGraphX's reproducibility utilities fit into it. The sklearn-style pipeline tutorial is a related approach to structuring experiments.
You will need pyyaml installed: pip install pyyaml.
Why YAML for GNN Experiments?
YAML is a human-readable serialization format that maps cleanly to Python dictionaries and lists. It is widely used for configuration in machine learning tools (Hydra, Weights & Biases sweeps, PyTorch Lightning) and requires no custom parser. YAML files can be:
- Version-controlled alongside code
- Diff'd to see exactly what changed between two experiments
- Used to generate hyperparameter sweeps by programmatically generating config variants
- Shared directly as the experiment specification in a paper's supplementary material
The alternative — argparse with many arguments — becomes unwieldy past a dozen hyperparameters and does not naturally support hierarchical configuration (e.g., separate sections for model, optimizer, data).
Designing the YAML Config Structure
A good config separates concerns: what data to use, what model to build, how to train it, and metadata about the experiment.
# experiment.yaml
experiment:
name: "citation_sage_baseline"
seed: 42
output_dir: "./runs/citation_sage_baseline"
data:
num_nodes: 500
num_features: 64
num_classes: 7
num_edges: 2000
homophily: 0.8
model:
type: "sage" # sage | gat | gin
hidden_channels: 128
num_layers: 2
dropout: 0.5
optimizer:
type: "adam"
lr: 0.005
weight_decay: 5.0e-4
training:
max_epochs: 300
early_stopping_patience: 50
early_stopping_min_delta: 1.0e-4
train_ratio: 0.14
val_ratio: 0.10
test_ratio: 0.20
This structure is self-documenting. Anyone reading the YAML file understands the full experiment specification without reading the Python code.
Loading the Config in Python
import yaml
from pathlib import Path
def load_config(config_path: str) -> dict:
"""Load a YAML config file and return as a nested dict."""
with open(config_path, 'r') as f:
cfg = yaml.safe_load(f)
return cfg
cfg = load_config("experiment.yaml")
# Access nested values
print(cfg['model']['hidden_channels']) # 128
print(cfg['optimizer']['lr']) # 0.005
print(cfg['experiment']['seed']) # 42
For convenience, convert the nested dict to a namespace-style object so you can use dot notation:
from types import SimpleNamespace
def dict_to_namespace(d):
"""Recursively convert dict to SimpleNamespace for dot-access."""
if isinstance(d, dict):
return SimpleNamespace(**{k: dict_to_namespace(v) for k, v in d.items()})
return d
config = dict_to_namespace(cfg)
print(config.model.hidden_channels) # 128
print(config.optimizer.lr) # 0.005
Building the Model from Config
The config's model.type field selects which TGraphX layer to use. A factory function maps this string to the correct class:
import torch.nn as nn
import torch.nn.functional as F
from tgraphx.layers.sage import TensorGraphSAGELayer
from tgraphx.layers.gat import TensorGATLayer
from tgraphx.layers.gin import TensorGINLayer
def build_model(cfg) -> nn.Module:
"""Build a GNN model from config."""
m = cfg.model
layer_class = {
'sage': TensorGraphSAGELayer,
'gat': TensorGATLayer,
'gin': TensorGINLayer,
}[m.type]
layers = nn.ModuleList()
in_dim = cfg.data.num_features
for i in range(m.num_layers):
out_dim = cfg.data.num_classes if i == m.num_layers - 1 else m.hidden_channels
layers.append(layer_class(in_dim, out_dim))
in_dim = out_dim
class ConfigGNN(nn.Module):
def __init__(self, layers, dropout):
super().__init__()
self.layers = layers
self.dropout = dropout
def forward(self, x, edge_index):
for i, layer in enumerate(self.layers):
x = layer(x, edge_index)
if i < len(self.layers) - 1:
x = F.relu(x)
x = F.dropout(x, p=self.dropout, training=self.training)
return x
return ConfigGNN(layers, m.dropout)
Building the Optimizer and Training Loop from Config
import torch
def build_optimizer(model: nn.Module, cfg) -> torch.optim.Optimizer:
"""Build optimizer from config."""
o = cfg.optimizer
if o.type == 'adam':
return torch.optim.Adam(
model.parameters(),
lr=o.lr,
weight_decay=o.weight_decay
)
elif o.type == 'sgd':
return torch.optim.SGD(
model.parameters(),
lr=o.lr,
momentum=getattr(o, 'momentum', 0.9)
)
else:
raise ValueError(f"Unknown optimizer type: {o.type}")
The complete runner function assembles all pieces:
import torch
from tgraphx import Graph
from tgraphx.reproducibility import set_reproducibility
from tgraphx.estimators.early_stopping import EarlyStopping
from tgraphx.estimators.splits import random_node_split
def run_experiment(config_path: str):
cfg = dict_to_namespace(load_config(config_path))
with set_reproducibility(seed=cfg.experiment.seed):
# Build dataset (synthetic here; replace with real data loader)
d = cfg.data
x = torch.randn(d.num_nodes, d.num_features)
y = torch.randint(0, d.num_classes, (d.num_nodes,))
edge_index = torch.randint(0, d.num_nodes, (2, d.num_edges), dtype=torch.long)
g = Graph(node_features=x, edge_index=edge_index)
# Splits
t = cfg.training
train_mask, val_mask, test_mask = random_node_split(
d.num_nodes, t.train_ratio, t.val_ratio, seed=cfg.experiment.seed
)
# Model and optimizer
model = build_model(cfg)
optimizer = build_optimizer(model, cfg)
stopper = EarlyStopping(
patience=t.early_stopping_patience,
min_delta=t.early_stopping_min_delta
)
best_val_loss = float('inf')
best_state = None
for epoch in range(t.max_epochs):
model.train()
logits = model(g.node_features, g.edge_index)
loss = F.cross_entropy(logits[train_mask], y[train_mask])
optimizer.zero_grad()
loss.backward()
optimizer.step()
model.eval()
with torch.no_grad():
val_logits = model(g.node_features, g.edge_index)
val_loss = F.cross_entropy(val_logits[val_mask], y[val_mask]).item()
val_acc = (val_logits[val_mask].argmax(1) == y[val_mask]).float().mean().item()
if val_loss < best_val_loss:
best_val_loss = val_loss
best_state = {k: v.clone() for k, v in model.state_dict().items()}
if stopper.step(val_loss):
print(f"[{cfg.experiment.name}] Early stop at epoch {epoch}")
break
model.load_state_dict(best_state)
model.eval()
with torch.no_grad():
test_logits = model(g.node_features, g.edge_index)
test_acc = (test_logits[test_mask].argmax(1) == y[test_mask]).float().mean().item()
print(f"[{cfg.experiment.name}] Test accuracy: {test_acc:.4f}")
return test_acc
# Run the experiment
# result = run_experiment("experiment.yaml")
Running Hyperparameter Sweeps
Once the runner takes a config path, running a sweep is just generating config variants and calling the runner:
import copy
base_cfg = load_config("experiment.yaml")
for lr in [1e-3, 5e-3, 1e-2]:
for hidden in [64, 128, 256]:
cfg_variant = copy.deepcopy(base_cfg)
cfg_variant['optimizer']['lr'] = lr
cfg_variant['model']['hidden_channels'] = hidden
cfg_variant['experiment']['name'] = f"sweep_lr{lr}_h{hidden}"
# Write variant to a temp file and run
with open('/tmp/sweep_config.yaml', 'w') as f:
yaml.dump(cfg_variant, f)
result = run_experiment('/tmp/sweep_config.yaml')
print(f"lr={lr}, hidden={hidden}: {result:.4f}")
Saving and Archiving Experiment Results
For full reproducibility, save the config alongside the model checkpoint:
import os
import json
def save_experiment(cfg, model, metrics, output_dir):
os.makedirs(output_dir, exist_ok=True)
# Save config
with open(os.path.join(output_dir, 'config.yaml'), 'w') as f:
yaml.dump(cfg if isinstance(cfg, dict) else vars(cfg), f)
# Save model weights
torch.save(model.state_dict(), os.path.join(output_dir, 'model.pt'))
# Save metrics
with open(os.path.join(output_dir, 'metrics.json'), 'w') as f:
json.dump(metrics, f, indent=2)
print(f"Saved experiment to {output_dir}")
This ensures that every archived run has its configuration and weights together, making it possible to reproduce or extend any previous experiment.
Versioning Configs with Git
One of the strongest arguments for YAML configs is that they are version-controlled as first-class citizens alongside code. A practical workflow:
- Keep a
configs/directory at the project root. - Name each config file descriptively:
sage_cora_baseline.yaml,gin_sweep_lr_v2.yaml. - Commit configs in the same pull request as the code changes that required them.
- When a paper is submitted, tag the commit and note which config files correspond to the reported results.
This workflow makes it possible to reproduce any experiment from the repository history:
# Reproduce experiment from 3 weeks ago
git checkout v1.2.3 # tag or commit hash
python run_experiment.py --config configs/sage_cora_baseline.yaml
Combining this with TGraphX's reproducibility context manager (which sets all seeds) ensures that the numerical results can be reproduced exactly, not just approximately.
Environment Config: Separating Code Config from Run Config
In practice, it is useful to separate config into two layers:
Experiment config (checked into git): Model architecture, hyperparameters, dataset settings — everything that defines the experiment.
Environment config (local, not checked in): Paths to data files, GPU settings, number of workers — everything that varies by machine.
# env.yaml (in .gitignore)
data_root: "/data/graphs"
num_workers: 8
cuda_device: 0
import os
def load_merged_config(experiment_config, env_config=None):
"""Load and merge experiment + environment configs."""
cfg = load_config(experiment_config)
if env_config and os.path.exists(env_config):
env = load_config(env_config)
cfg.update(env) # environment settings override experiment settings
return cfg
This separation keeps experiment configs portable and shareable while allowing each machine to define its own data paths and resource settings.
Limitations and Honest Notes
YAML has several footguns worth knowing. Unquoted strings that look like booleans (yes, no, true, false) are parsed as Python booleans, not strings. Scientific notation in floats (5e-4) is correctly parsed as a float, but 5e4 (without decimal) may parse as an integer in some YAML parsers. Always use yaml.safe_load rather than yaml.load to avoid arbitrary code execution.
The SimpleNamespace approach for dot-access is convenient but lacks validation. If a key is missing from the config, you get an AttributeError at runtime rather than at config load time. For production use, a config validation library like pydantic or cerberus provides better error messages and type checking.
This tutorial does not include integration with experiment tracking systems like MLflow or Weights & Biases. TGraphX's tgraphx.tracking module provides lightweight local tracking. For team-scale experiment management, external tools are more appropriate.
Config files version-control well but can accumulate stale or unused keys as the codebase evolves. Periodically audit config files against the code to remove keys that are no longer used.
Integrating with TGraphX Doctor and Validation
Before a training run, it is good practice to validate that the config produces a valid model and that the graph data is well-formed. TGraphX's tgraphx.doctor module provides validation utilities:
from tgraphx.doctor import GraphDoctor
def validate_before_training(g, config):
"""Run validation checks before starting an expensive training run."""
doctor = GraphDoctor(g)
report = doctor.run()
if report.has_isolated_nodes:
print(f"Warning: {report.num_isolated_nodes} isolated nodes detected.")
if report.has_self_loops and config.model.type == 'gin':
print("Warning: self-loops present — GIN sum aggregation may count self-edges.")
if not report.is_connected:
print(f"Warning: Graph has {report.num_components} components.")
# Validate model can handle the feature shape
model = build_model(config)
x_sample = torch.randn(10, config.data.num_features)
edge_sample = torch.randint(0, 10, (2, 20), dtype=torch.long)
try:
with torch.no_grad():
_ = model(x_sample, edge_sample)
print("Model forward pass: OK")
except Exception as e:
print(f"Model forward pass failed: {e}")
raise
Running these checks before a long experiment catches configuration mistakes (wrong feature dimension, incompatible layer types) before they waste compute time.
Frequently Asked Questions
Can I override individual config values from the command line?
Not directly with this approach. A common pattern is to accept --override key.path=value arguments and apply them to the loaded config dict before building the model. Alternatively, use a purpose-built config management library like Hydra, which provides command-line overrides natively.
How do I handle dataset-specific configs vs model-specific configs?
Use separate YAML files and merge them in the runner: one base config for dataset settings, one for model settings. Python's dict update semantics make shallow merging easy; for deep merging (nested dicts), use a utility like deepmerge.
Should I commit experiment configs to git?
Yes. Config files are small, human-readable, and provide a record of what was tried. Committing them alongside code ensures that the exact experiment setup is reproducible by anyone with access to the repository.