harsh040506/ml-engineering

ML Engineering

Production-grade guidance for training, optimizing, and deploying machine learning models.

---

The ML Development Lifecycle

Problem definition → Data collection → EDA → Feature engineering
→ Model selection → Training → Evaluation → Hyperparameter tuning
→ Error analysis → Deployment → Monitoring → Retraining

Never skip the problem definition or data stages. The most common ML failures are:

1. Solving the wrong problem
2. Bad data quality masquerading as a model quality problem
3. Data leakage causing inflated offline metrics
4. Ignoring production distribution shift

---

Data Best Practices

Train/Val/Test Splits

For time-series or temporal data: Split by time, not randomly.

• Training: oldest data
• Validation: middle period
• Test: most recent data (held out completely until final evaluation)

Random splits on time-series cause data leakage — future information leaks into training.

For i.i.d. data:

• 70/10/20 or 80/10/10 split
• Stratified split for classification to preserve class ratios
• Group-aware split when multiple samples share an entity (same user, document, patient)

Preventing Data Leakage

Leakage = test-set information contaminating the training set. It makes models look better than they are.

Common leakage sources:

• Applying scaling/normalization using statistics from the full dataset before splitting
• Target-encoding using the target variable from test samples
• Temporal overlap in time-series splits
• Duplicated rows that appear in both train and test

# WRONG — leaks test statistics
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)  # Fits on entire dataset
X_train, X_test = train_test_split(X_scaled)

# CORRECT
X_train, X_test = train_test_split(X)
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)   # Fit only on train
X_test_scaled = scaler.transform(X_test)          # Transform test with train statistics

---

Model Selection

Starting Point Heuristic

| Task | Start here | Then try | |------|-----------|---------| | Tabular classification | LightGBM / XGBoost | Neural network if still underperforming | | Tabular regression | LightGBM / XGBoost | Neural network | | Image classification | ResNet50 (fine-tune) | EfficientNet, ConvNeXt | | Text classification | DistilBERT / RoBERTa (fine-tune) | Larger model if needed | | Text generation | Fine-tune instruction-tuned LLM (Mistral 7B, Llama 3) | GPT-4 API if quality > cost | | Sequence-to-sequence | T5/mT5 fine-tune | LLM fine-tune | | Tabular anomaly detection | Isolation Forest | Autoencoder | | Embeddings | sentence-transformers/all-MiniLM-L6-v2 | Fine-tune on domain data |

Always establish a simple baseline first (logistic regression, rule-based heuristic, random). Beat the baseline before adding complexity.

---

PyTorch Training Loop (Production Template)

import torch
import torch.nn as nn
from torch.cuda.amp import GradScaler, autocast
import wandb
from pathlib import Path

def train(
    model: nn.Module,
    train_loader,
    val_loader,
    config: dict,
    output_dir: str = "./checkpoints",
):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)

    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=config["lr"],
        weight_decay=config.get("weight_decay", 0.01),
    )
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer, T_max=config["epochs"]
    )
    scaler = GradScaler()  # Mixed precision training
    criterion = nn.CrossEntropyLoss()

    wandb.init(project=config["project"], config=config)
    
    best_val_loss = float("inf")
    patience_counter = 0
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)

    for epoch in range(config["epochs"]):
        # ── Training ──────────────────────────────
        model.train()
        train_loss = 0.0
        for batch_idx, (inputs, targets) in enumerate(train_loader):
            inputs, targets = inputs.to(device), targets.to(device)
            optimizer.zero_grad()

            with autocast():  # Mixed precision — 2x speedup on modern GPUs
                outputs = model(inputs)
                loss = criterion(outputs, targets)

            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)  # Gradient clipping
            scaler.step(optimizer)
            scaler.update()
            train_loss += loss.item()

        # ── Validation ────────────────────────────
        model.eval()
        val_loss = 0.0
        correct = 0
        total = 0
        with torch.no_grad():
            for inputs, targets in val_loader:
                inputs, targets = inputs.to(device), targets.to(device)
                with autocast():
                    outputs = model(inputs)
                    loss = criterion(outputs, targets)
                val_loss += loss.item()
                _, predicted = outputs.max(1)
                correct += predicted.eq(targets).sum().item()
                total += targets.size(0)

        train_loss /= len(train_loader)
        val_loss /= len(val_loader)
        val_acc = correct / total

        wandb.log({
            "epoch": epoch,
            "train/loss": train_loss,
            "val/loss": val_loss,
            "val/accuracy": val_acc,
            "lr": scheduler.get_last_lr()[0],
        })

        print(f"Epoch {epoch+1}/{config['epochs']} | "
              f"Train Loss: {train_loss:.4f} | "
              f"Val Loss: {val_loss:.4f} | "
              f"Val Acc: {val_acc:.4f}")

        # ── Checkpointing ─────────────────────────
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            patience_counter = 0
            torch.save({
                "epoch": epoch,
                "model_state_dict": model.state_dict(),
                "optimizer_state_dict": optimizer.state_dict(),
                "val_loss": val_loss,
                "config": config,
            }, output_path / "best_model.pt")
            wandb.save(str(output_path / "best_model.pt"))
        else:
            patience_counter += 1
            if patience_counter >= config.get("patience", 10):
                print(f"Early stopping at epoch {epoch+1}")
                break

        scheduler.step()

    wandb.finish()
    return output_path / "best_model.pt"

---

Diagnosing Training Problems

Loss is NaN

1. Check for NaN in data — df.isnull().any(), np.isnan(X).any()
2. Learning rate too high — try 1/10th the current LR
3. Exploding gradients — add gradient clipping: clip_grad_norm_(model.parameters(), 1.0)
4. Log of zero — numerical instability in loss function; add epsilon: log(x + 1e-8)
5. Infinity in features — check for division by zero in preprocessing

Loss Not Decreasing

1. Learning rate too low — try 10× higher, or use learning rate finder
2. Learning rate too high — loss oscillates or diverges
3. Bug in data loading — confirm labels are correct with dataloader[0]
4. Weight initialization — use appropriate init (Xavier for tanh, Kaiming for ReLU)
5. Model capacity too low — try larger model

Overfitting

Gap between training and validation metrics is growing:

1. More data — always the best solution
2. Data augmentation — random crop, flip, color jitter for images; back-translation, synonym replacement for text
3. Regularization — increase dropout rate (0.1 → 0.3), L2 weight decay
4. Reduce model size — fewer layers or smaller hidden dim
5. Early stopping — stop when val loss stops improving

Training Too Slow

1. Enable mixed precision — GradScaler + autocast → 2–3× speedup
2. Increase batch size — more GPU utilization (check GPU utilization first: nvidia-smi)
3. DataLoader workers — set num_workers=4 (or num CPU cores)
4. Pin memory — DataLoader(pin_memory=True) for faster CPU→GPU transfer
5. Profile — use torch.profiler to identify the bottleneck

---

Hyperparameter Search

Use Optuna for efficient hyperparameter search. Prefer Tree-structured Parzen Estimators (TPE) over grid search — it finds good parameters 10× faster.

import optuna

def objective(trial):
    config = {
        "lr": trial.suggest_float("lr", 1e-5, 1e-2, log=True),
        "batch_size": trial.suggest_categorical("batch_size", [16, 32, 64, 128]),
        "dropout": trial.suggest_float("dropout", 0.1, 0.5),
        "hidden_dim": trial.suggest_categorical("hidden_dim", [128, 256, 512]),
        "epochs": 20,  # Short runs for search
        "patience": 5,
    }
    
    # Train with these params and return val metric
    val_loss = train_and_evaluate(config)
    return val_loss

study = optuna.create_study(direction="minimize")
study.optimize(objective, n_trials=50, timeout=3600)  # 50 trials or 1 hour

print(f"Best params: {study.best_params}")
print(f"Best val loss: {study.best_value}")

---

Model Deployment Patterns

ONNX Export (Framework-Agnostic Inference)

import torch.onnx

model.eval()
dummy_input = torch.randn(1, input_size)
torch.onnx.export(
    model, dummy_input,
    "model.onnx",
    export_params=True,
    opset_version=17,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}},
)

# Validate export
import onnxruntime as ort
session = ort.InferenceSession("model.onnx")
output = session.run(None, {"input": dummy_input.numpy()})

FastAPI Serving

from fastapi import FastAPI
from pydantic import BaseModel
import onnxruntime as ort
import numpy as np

app = FastAPI()
session = ort.InferenceSession("model.onnx", providers=["CUDAExecutionProvider", "CPUExecutionProvider"])

class PredictRequest(BaseModel):
    features: list[float]

class PredictResponse(BaseModel):
    prediction: int
    confidence: float

@app.post("/predict", response_model=PredictResponse)
def predict(request: PredictRequest):
    input_array = np.array([request.features], dtype=np.float32)
    logits = session.run(None, {"input": input_array})[0]
    probs = softmax(logits[0])
    return PredictResponse(
        prediction=int(probs.argmax()),
        confidence=float(probs.max()),
    )

Deeper Reference

For complete training recipes and hyperparameter tuning guides, see:

• references/training-recipes.md — end-to-end PyTorch and HuggingFace training scripts with mixed-precision, gradient checkpointing, and distributed training
• references/hyperparameter-guide.md — learning rate schedules, batch size scaling rules, regularization strategies, and Optuna/Ray Tune search configurations