Building a Gin Config Controlled PyTorch Pipeline with Configurable MLP Variants, Cosine Scheduling, and Runtime Parameter Overrides

building-a-gin-config-controlled-pytorch-pipeline-with-configurable-mlp-variants,-cosine-scheduling,-and-runtime-parameter-overrides
Building a Gin Config Controlled PyTorch Pipeline with Configurable MLP Variants, Cosine Scheduling, and Runtime Parameter Overrides

In this tutorial, we implement a Gin Config–controlled PyTorch experiment pipeline in which the executable training code remains stable. At the same time, the experimental degrees of freedom are moved into declarative configuration files. We construct a nonlinear spiral binary classification task, define a configurable MLP with scoped architectural variants, and expose parameters for the optimizer, scheduler, loss, batching, seeding, and training loop via @gin.configurable bindings. We use Gin’s scoped references to instantiate separate model configurations, runtime bindings to override selected parameters without editing source code, and operative config export to capture the exact resolved configuration that produces each training run.

Installing Gin Config and Building the Spiral Dataset

!pip -q install gin-config import os import json import math import random import textwrap from pathlib import Path import gin import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import TensorDataset, DataLoader import matplotlib.pyplot as plt ROOT = Path("https://www.marktechpost.com/content/gin_config_sharp_tutorial") CONFIG_DIR = ROOT / "configs" RUN_DIR = ROOT / "runs" CONFIG_DIR.mkdir(parents=True, exist_ok=True) RUN_DIR.mkdir(parents=True, exist_ok=True) gin.clear_config() @gin.configurable def seed_everything(seed=42):    random.seed(seed)    np.random.seed(seed)    torch.manual_seed(seed)    torch.cuda.manual_seed_all(seed)    return seed @gin.configurable def make_spiral_dataset(    n_per_class=gin.REQUIRED,    noise=0.18,    rotations=1.75,    train_fraction=0.8,    seed=0, ):    rng = np.random.default_rng(seed)    radius_0 = np.linspace(0.05, 1.0, n_per_class)    theta_0 = rotations * 2 * np.pi * radius_0    theta_0 += rng.normal(0.0, noise, size=n_per_class)    x0 = np.stack(        [            radius_0 * np.cos(theta_0),            radius_0 * np.sin(theta_0),        ],        axis=1,    )    radius_1 = np.linspace(0.05, 1.0, n_per_class)    theta_1 = rotations * 2 * np.pi * radius_1 + np.pi    theta_1 += rng.normal(0.0, noise, size=n_per_class)    x1 = np.stack(        [            radius_1 * np.cos(theta_1),            radius_1 * np.sin(theta_1),        ],        axis=1,    )    x = np.concatenate([x0, x1], axis=0).astype(np.float32)    y = np.concatenate(        [            np.zeros((n_per_class, 1)),            np.ones((n_per_class, 1)),        ],        axis=0,    ).astype(np.float32)    order = rng.permutation(len(x))    x = x[order]    y = y[order]    split = int(train_fraction * len(x))    x_train, y_train = x[:split], y[:split]    x_val, y_val = x[split:], y[split:]    mean = x_train.mean(axis=0, keepdims=True)    std = x_train.std(axis=0, keepdims=True) + 1e-8    x_train = (x_train - mean) / std    x_val = (x_val - mean) / std    return {        "train": (            torch.tensor(x_train),            torch.tensor(y_train),        ),        "val": (            torch.tensor(x_val),            torch.tensor(y_val),        ),        "metadata": {            "n_train": int(len(x_train)),            "n_val": int(len(x_val)),            "n_features": int(x_train.shape[1]),            "noise": float(noise),            "rotations": float(rotations),            "seed": int(seed),        },    } @gin.configurable(denylist=["x", "y"]) def make_loader(    x,    y,    batch_size=128,    shuffle=True,    seed=0, ):    generator = torch.Generator()    generator.manual_seed(seed)    dataset = TensorDataset(x, y)    return DataLoader(        dataset,        batch_size=batch_size,        shuffle=shuffle,        generator=generator,        drop_last=False,    ) 

We start by installing Gin Config and importing the core Python libraries, PyTorch, NumPy, and the plotting libraries required for the experiment. We create a clean project directory structure and reset Gin’s global configuration state so the notebook runs reproducibly. We then define the seed function, generate a nonlinear spiral dataset, and build a configurable DataLoader that Gin can control through external bindings.

Defining a Gin-Configurable MLP, Optimizer, and Scheduler

def activation_layer(name):    name = name.lower()    if name == "relu":        return nn.ReLU()    if name == "gelu":        return nn.GELU()    if name == "tanh":        return nn.Tanh()    if name == "silu":        return nn.SiLU()    raise ValueError(f"Unknown activation: {name}") @gin.configurable class MLP(nn.Module):    def __init__(        self,        input_dim=gin.REQUIRED,        hidden_dims=(64, 64),        output_dim=1,        activation="gelu",        dropout=0.0,        use_layernorm=False,    ):        super().__init__()        layers = []        current_dim = input_dim        for hidden_dim in hidden_dims:            layers.append(nn.Linear(current_dim, hidden_dim))            if use_layernorm:                layers.append(nn.LayerNorm(hidden_dim))            layers.append(activation_layer(activation))            if dropout > 0:                layers.append(nn.Dropout(dropout))            current_dim = hidden_dim        layers.append(nn.Linear(current_dim, output_dim))        self.network = nn.Sequential(*layers)    def forward(self, x):        return self.network(x) @gin.configurable(denylist=["params"]) def make_optimizer(    params,    name="adamw",    lr=3e-3,    weight_decay=1e-3,    momentum=0.9, ):    name = name.lower()    if name == "adamw":        return torch.optim.AdamW(            params,            lr=lr,            weight_decay=weight_decay,        )    if name == "sgd":        return torch.optim.SGD(            params,            lr=lr,            momentum=momentum,            weight_decay=weight_decay,        )    raise ValueError(f"Unknown optimizer: {name}") @gin.configurable(denylist=["optimizer"]) def make_cosine_scheduler(    optimizer,    total_epochs=60,    warmup_epochs=5,    min_lr_factor=0.05, ):    def lr_lambda(epoch):        if epoch < warmup_epochs:            return float(epoch + 1) / float(max(1, warmup_epochs))        progress = (epoch - warmup_epochs) / float(            max(1, total_epochs - warmup_epochs)        )        cosine = 0.5 * (1.0 + math.cos(math.pi * progress))        return min_lr_factor + (1.0 - min_lr_factor) * cosine    return torch.optim.lr_scheduler.LambdaLR(        optimizer,        lr_lambda=lr_lambda,    ) @gin.configurable def bce_with_logits_loss(    logits,    targets,    label_smoothing=0.0, ):    if label_smoothing > 0:        targets = targets * (1.0 - label_smoothing) + 0.5 * label_smoothing    return F.binary_cross_entropy_with_logits(logits, targets) @torch.no_grad() def evaluate(model, loader, loss_fn, device):    model.eval()    total_loss = 0.0    total_correct = 0    total_count = 0    for x, y in loader:        x = x.to(device)        y = y.to(device)        logits = model(x)        loss = loss_fn(logits, y)        probs = torch.sigmoid(logits)        preds = (probs >= 0.5).float()        total_loss += loss.item() * len(x)        total_correct += (preds == y).sum().item()        total_count += len(x)    return {        "loss": total_loss / total_count,        "accuracy": total_correct / total_count,    } 

We define the neural network building blocks that form the configurable model and the training utilities. We create an MLP class whose architecture, activation function, dropout, and layer normalization behavior are controlled through Gin rather than hardcoded values. We also implement configurable optimizer, scheduler, loss, and evaluation functions so the training pipeline remains modular and experiment-ready.

Implementing the Training Loop and Experiment Runner

@gin.configurable(    denylist=[        "model",        "optimizer",        "scheduler",        "train_loader",        "val_loader",        "device",    ] ) def fit(    model,    optimizer,    scheduler,    train_loader,    val_loader,    device,    epochs=60,    grad_clip_norm=1.0,    log_every=10,    loss_fn=bce_with_logits_loss, ):    history = []    for epoch in range(1, epochs + 1):        model.train()        for x, y in train_loader:            x = x.to(device)            y = y.to(device)            optimizer.zero_grad(set_to_none=True)            logits = model(x)            loss = loss_fn(logits, y)            loss.backward()            if grad_clip_norm is not None:                nn.utils.clip_grad_norm_(                    model.parameters(),                    grad_clip_norm,                )            optimizer.step()        if scheduler is not None:            scheduler.step()        train_metrics = evaluate(            model,            train_loader,            loss_fn,            device,        )        val_metrics = evaluate(            model,            val_loader,            loss_fn,            device,        )        lr = optimizer.param_groups[0]["lr"]        row = {            "epoch": epoch,            "lr": lr,            "train_loss": train_metrics["loss"],            "train_accuracy": train_metrics["accuracy"],            "val_loss": val_metrics["loss"],            "val_accuracy": val_metrics["accuracy"],        }        history.append(row)        if epoch == 1 or epoch % log_every == 0 or epoch == epochs:            print(                f"epoch={epoch:03d} | "                f"lr={lr:.6f} | "                f"train_loss={row['train_loss']:.4f} | "                f"train_acc={row['train_accuracy']:.3f} | "                f"val_loss={row['val_loss']:.4f} | "                f"val_acc={row['val_accuracy']:.3f}"            )    return history @gin.configurable def run_experiment(    tag=gin.REQUIRED,    model=gin.REQUIRED,    dataset_fn=make_spiral_dataset,    optimizer_factory=make_optimizer,    scheduler_factory=make_cosine_scheduler,    prefer_gpu=True, ):    seed_everything()    device = "cuda" if prefer_gpu and torch.cuda.is_available() else "cpu"    data = dataset_fn()    x_train, y_train = data["train"]    x_val, y_val = data["val"]    train_loader = make_loader(        x_train,        y_train,        shuffle=True,    )    val_loader = make_loader(        x_val,        y_val,        shuffle=False,    )    model = model.to(device)    optimizer = optimizer_factory(model.parameters())    scheduler = None    if scheduler_factory is not None:        scheduler = scheduler_factory(optimizer)    print("n" + "=" * 80)    print(f"Experiment: {tag}")    print("=" * 80)    print(f"Device: {device}")    print(f"Dataset: {data['metadata']}")    print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")    history = fit(        model=model,        optimizer=optimizer,        scheduler=scheduler,        train_loader=train_loader,        val_loader=val_loader,        device=device,    )    result = {        "tag": tag,        "device": device,        "metadata": data["metadata"],        "parameters": sum(p.numel() for p in model.parameters()),        "final": history[-1],        "history": history,    }    return result 

We implement the main training loop, in which the model performs forward passes, computes binary cross-entropy loss, backpropagates gradients, applies gradient clipping, and updates parameters. We evaluate the model after each epoch on both the training and validation sets, while storing loss, accuracy, and learning rate history. We then define the top-level experiment runner that connects the dataset, model, optimizer, scheduler, and training loop through Gin-managed dependencies.

Writing Gin Config Files with Scoped Bindings and Runtime Overrides

BASE_CONFIG = CONFIG_DIR / "base.gin" COMPACT_CONFIG = CONFIG_DIR / "compact_adamw.gin" WIDE_CONFIG = CONFIG_DIR / "wide_sgd.gin" BASE_CONFIG.write_text(    textwrap.dedent(        """        SEED = 123        N_PER_CLASS = 900        EPOCHS = 50        BATCH = 128        seed_everything.seed = %SEED        make_spiral_dataset.n_per_class = %N_PER_CLASS        make_spiral_dataset.noise = 0.20        make_spiral_dataset.rotations = 1.85        make_spiral_dataset.train_fraction = 0.80        make_spiral_dataset.seed = %SEED        make_loader.batch_size = %BATCH        make_loader.seed = %SEED        MLP.input_dim = 2        MLP.output_dim = 1        MLP.activation = 'gelu'        MLP.dropout = 0.05        MLP.use_layernorm = True        make_optimizer.name = 'adamw'        make_optimizer.lr = 0.003        make_optimizer.weight_decay = 0.001        make_optimizer.momentum = 0.9        make_cosine_scheduler.total_epochs = %EPOCHS        make_cosine_scheduler.warmup_epochs = 5        make_cosine_scheduler.min_lr_factor = 0.05        bce_with_logits_loss.label_smoothing = 0.02        fit.epochs = %EPOCHS        fit.grad_clip_norm = 1.0        fit.log_every = 10        fit.loss_fn = @bce_with_logits_loss        run_experiment.dataset_fn = @make_spiral_dataset        run_experiment.optimizer_factory = @make_optimizer        run_experiment.scheduler_factory = @make_cosine_scheduler        run_experiment.prefer_gpu = True        """    ).strip() ) COMPACT_CONFIG.write_text(    textwrap.dedent(        f"""        include '{BASE_CONFIG.as_posix()}'        run_experiment.tag = 'compact_gelu_adamw'        run_experiment.model = @compact/MLP()        compact/MLP.hidden_dims = (64, 64, 64)        compact/MLP.dropout = 0.05        compact/MLP.use_layernorm = True        make_optimizer.name = 'adamw'        make_optimizer.lr = 0.003        make_optimizer.weight_decay = 0.001        """    ).strip() ) WIDE_CONFIG.write_text(    textwrap.dedent(        f"""        include '{BASE_CONFIG.as_posix()}'        run_experiment.tag = 'wide_relu_sgd'        run_experiment.model = @wide/MLP()        wide/MLP.hidden_dims = (128, 128, 128, 64)        wide/MLP.activation = 'relu'        wide/MLP.dropout = 0.02        wide/MLP.use_layernorm = True        make_optimizer.name = 'sgd'        make_optimizer.lr = 0.035        make_optimizer.momentum = 0.92        make_optimizer.weight_decay = 0.0005        bce_with_logits_loss.label_smoothing = 0.0        """    ).strip() ) def run_from_gin_file(config_path, runtime_bindings=None):    runtime_bindings = runtime_bindings or []    gin.clear_config()    gin.parse_config_files_and_bindings(        config_files=[str(config_path)],        bindings=runtime_bindings,        skip_unknown=False,        finalize_config=True,    )    print("nLoaded config file:")    print(config_path)    print("nSelected queried parameters:")    print("fit.epochs =", gin.query_parameter("fit.epochs"))    print("make_loader.batch_size =", gin.query_parameter("make_loader.batch_size"))    print("make_spiral_dataset.noise =", gin.query_parameter("make_spiral_dataset.noise"))    try:        gin.bind_parameter("fit.epochs", 999)    except RuntimeError as error:        print("nConfig lock check:")        print(str(error).splitlines()[0])    result = run_experiment()    tag = result["tag"]    out_dir = RUN_DIR / tag    out_dir.mkdir(parents=True, exist_ok=True)    result_path = out_dir / "result.json"    operative_path = out_dir / "operative_config.gin"    result_path.write_text(json.dumps(result, indent=2))    operative_path.write_text(gin.operative_config_str())    print("nSaved:")    print(result_path)    print(operative_path)    return result, operative_path compact_result, compact_operative = run_from_gin_file(    COMPACT_CONFIG,    runtime_bindings=[        "fit.epochs = 45",        "make_spiral_dataset.noise = 0.18",        "run_experiment.tag = 'compact_gelu_adamw_runtime_override'",    ], ) wide_result, wide_operative = run_from_gin_file(    WIDE_CONFIG,    runtime_bindings=[        "fit.epochs = 45",        "make_spiral_dataset.noise = 0.18",        "run_experiment.tag = 'wide_relu_sgd_runtime_override'",    ], ) 

We create the actual Gin configuration files that control the experiment without modifying the Python source code. We define a shared base configuration and then compose two scoped experiments: a compact GELU-based AdamW model and a wider ReLU-based SGD model. We also demonstrate runtime overrides, parameter queries, config locking, result serialization, and operative config export for reproducible experiment tracking.

Comparing Results and Exporting the Operative Config

def plot_metric(results, metric, title):    plt.figure(figsize=(9, 4))    for result in results:        epochs = [row["epoch"] for row in result["history"]]        values = [row[metric] for row in result["history"]]        plt.plot(epochs, values, label=result["tag"])    plt.xlabel("Epoch")    plt.ylabel(metric)    plt.title(title)    plt.grid(True, alpha=0.3)    plt.legend()    plt.tight_layout()    plt.show() plot_metric(    [compact_result, wide_result],    "val_loss",    "Validation Loss Controlled by Gin Config", ) plot_metric(    [compact_result, wide_result],    "val_accuracy",    "Validation Accuracy Controlled by Gin Config", ) summary = [    {        "tag": compact_result["tag"],        "params": compact_result["parameters"],        "val_loss": compact_result["final"]["val_loss"],        "val_accuracy": compact_result["final"]["val_accuracy"],    },    {        "tag": wide_result["tag"],        "params": wide_result["parameters"],        "val_loss": wide_result["final"]["val_loss"],        "val_accuracy": wide_result["final"]["val_accuracy"],    }, ] print("n" + "=" * 80) print("Final comparison") print("=" * 80) for row in summary:    print(        f"{row['tag']} | "        f"params={row['params']:,} | "        f"val_loss={row['val_loss']:.4f} | "        f"val_acc={row['val_accuracy']:.3f}"    ) print("n" + "=" * 80) print("Compact experiment operative config preview") print("=" * 80) print(compact_operative.read_text()[:2500]) print("n" + "=" * 80) print("Generated files") print("=" * 80) for path in sorted(ROOT.rglob("*")):    if path.is_file():        print(path) 

We visualize the validation loss and validation accuracy curves for both Gin-controlled experiments. We summarize the final parameter counts, validation losses, and validation accuracies to clearly compare the two configurations. We also print the operative configuration and the generated files, which provide a complete record of the exact settings used during execution.

Conclusion

In conclusion, we have a reproducible experiment-management workflow that demonstrates how Gin Config improves control, traceability, and modularity in PyTorch projects. We ran multiple scoped experiments from composed .gin files, compared AdamW and SGD training behavior under controlled dataset and epoch settings, verified configuration locking after parsing, and saved both metrics and operative configs for later inspection. It gives us a pattern for scaling Colab experiments into research-grade pipelines, in which model architecture, optimization strategy, data generation, and training schedules must be systematically adjusted without breaking the core implementation.


Check out the Full Codes with Notebook hereAlso, feel free to follow us on Twitter and don’t forget to join our 150k+ML SubReddit and Subscribe to our Newsletter. Wait! are you on telegram? now you can join us on telegram as well.

Need to partner with us for promoting your GitHub Repo OR Hugging Face Page OR Product Release OR Webinar etc.? Connect with us

Sana Hassan, a consulting intern at Marktechpost and dual-degree student at IIT Madras, is passionate about applying technology and AI to address real-world challenges. With a keen interest in solving practical problems, he brings a fresh perspective to the intersection of AI and real-life solutions.

Leave a Reply

Your email address will not be published. Required fields are marked *