A Google tutorial demonstrates how to build a PyTorch training pipeline where the executable code stays fixed while experimental variables move into external configuration files. The guide sets up a nonlinear spiral binary classification task, defines a multi-layer perceptron with scoped architectural variants, and exposes parameters for the optimizer, scheduler, loss, batching, seeding, and training loop via @gin.configurable bindings. It uses 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
The process begins by installing Gin Config and importing the core Python libraries, PyTorch, NumPy, and the plotting libraries required for the experiment. A clean project directory structure is created and Gin’s global configuration state is reset so the notebook runs reproducibly. The seed function is defined, a nonlinear spiral dataset is generated, and a configurable DataLoader is built that Gin can control through external bindings.
!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("/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,
)
Defining a Gin-Configurable MLP, Optimizer, and Scheduler
The tutorial defines the neural network building blocks that form the configurable model and the training utilities. An MLP class is created whose architecture, activation function, dropout, and layer normalization behavior are controlled through Gin rather than hardcoded values. Configurable functions for the optimizer, scheduler, loss, and evaluation are also implemented so the training pipeline remains modular and experiment-ready.
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,
}
