magnus919_agent-skills/data-scientist/references/pytorch-integration.md

PyTorch Integration Reference

Source validated against: PyTorch 2.12 documentation (pytorch.org/docs/stable)
Last reviewed: 2026-05-23
When to load: The campaign protocol (Phase 3-7) or any task requires implementing, training, or deploying a PyTorch model.

---

Device Management

Canonical Device Pattern

Always parameterize the device. Never hardcode "cuda".

import torch

device = torch.device(
    "cuda" if torch.cuda.is_available()
    else "mps" if torch.backends.mps.is_available()
    else "cpu"
)
# Usage: model.to(device); tensor.to(device)

Checking Device Properties

if torch.cuda.is_available():
    print(f"Device: {torch.cuda.get_device_name(0)}")
    print(f"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
    print(f"CUDA Capability: {torch.cuda.get_device_capability(0)}")

MPS (Apple Silicon): Available on macOS 12.3+. Not all operations are supported — check torch.backends.mps.is_available() and torch.backends.mps.is_built(). Some operations fall back to CPU automatically.

---

Training Loop Patterns

Basic Supervised Training Loop

model.train()
for epoch in range(num_epochs):
    running_loss = 0.0
    for batch_x, batch_y in dataloader:
        batch_x, batch_y = batch_x.to(device), batch_y.to(device)

        optimizer.zero_grad()
        outputs = model(batch_x)
        loss = criterion(outputs, batch_y)
        loss.backward()

        # Gradient clipping (essential for stability)
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

        optimizer.step()

        running_loss += loss.item()

    avg_loss = running_loss / len(dataloader)
    print(f"Epoch {epoch}: loss={avg_loss:.4f}")

Gradient Accumulation (for large models on limited VRAM)

accumulation_steps = 4  # Effective batch_size = physical_batch * accumulation

for i, (batch_x, batch_y) in enumerate(dataloader):
    outputs = model(batch_x)
    loss = criterion(outputs, batch_y) / accumulation_steps
    loss.backward()

    if (i + 1) % accumulation_steps == 0:
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()
        optimizer.zero_grad()

Validation Loop

model.eval()
val_loss = 0.0
correct = 0
total = 0

with torch.no_grad():
    for batch_x, batch_y in val_loader:
        batch_x, batch_y = batch_x.to(device), batch_y.to(device)
        outputs = model(batch_x)
        val_loss += criterion(outputs, batch_y).item()
        _, predicted = torch.max(outputs, 1)
        total += batch_y.size(0)
        correct += (predicted == batch_y).sum().item()

print(f"Val Loss: {val_loss/len(val_loader):.4f}, Acc: {100*correct/total:.2f}%")

---

Dataset & DataLoader

Custom Dataset Class

from torch.utils.data import Dataset, DataLoader

class CustomDataset(Dataset):
    def __init__(self, features, labels, transform=None):
        self.features = torch.tensor(features, dtype=torch.float32)
        self.labels = torch.tensor(labels, dtype=torch.long)
        self.transform = transform

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        x = self.features[idx]
        y = self.labels[idx]
        if self.transform:
            x = self.transform(x)
        return x, y

DataLoader Configuration

dataloader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4,          # Set to 0 on Windows if multiprocessing issues
    pin_memory=True,        # Speeds up GPU transfer (only with CUDA)
    persistent_workers=True if num_workers > 0 else False,  # PyTorch 2.0+
    collate_fn=None,        # Custom collation for variable-length data
)

Note on num_workers: On Linux, 4-8 workers is typical. On macOS, keep at 0-2. On Windows, 0 is safest. The optimal value depends on the data loading speed vs GPU speed.

Collate Function for Variable-Length Data

def collate_fn(batch):
    """Pad sequences in a batch to the same length."""
    inputs, labels = zip(*batch)
    # Pad to max length in this batch
    inputs_padded = torch.nn.utils.rnn.pad_sequence(inputs, batch_first=True)
    return inputs_padded, torch.tensor(labels)

---

Model Saving & Loading

Save/Load State Dict (Recommended)

# Save
torch.save({
    "epoch": epoch,
    "model_state_dict": model.state_dict(),
    "optimizer_state_dict": optimizer.state_dict(),
    "loss": loss,
    "config": {"n_layers": 4, "hidden_dim": 256},  # metadata
}, "checkpoint.pt")

# Load
checkpoint = torch.load("checkpoint.pt", map_location=device)
model.load_state_dict(checkpoint["model_state_dict"])
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
start_epoch = checkpoint["epoch"] + 1

Save for Inference (weights only)

torch.save(model.state_dict(), "model_weights.pt")

# Load for inference
model = YourModel()
model.load_state_dict(torch.load("model_weights.pt", map_location="cpu"))
model.eval()

TorchScript / torch.export (for production deployment)

# torch.export (PyTorch 2.x preferred, producesExportedProgram)
exported_program = torch.export.export(model, (example_input,))
torch.export.save(exported_program, "model.pt2")

# TorchScript (legacy, still widely supported)
scripted_model = torch.jit.script(model)
scripted_model.save("model_scripted.pt")

---

Mixed Precision (AMP)

Automatic Mixed Precision trains with float16 (or bfloat16) where safe and float32 where needed. Typically 1.5-2x faster with minimal accuracy loss.

from torch.amp import autocast, GradScaler

scaler = GradScaler("cuda")  # "cuda" or "cpu"

for batch_x, batch_y in dataloader:
    batch_x, batch_y = batch_x.to(device), batch_y.to(device)

    optimizer.zero_grad()

    # Autocast context manager
    with autocast(device_type="cuda"):  # or "cpu"
        outputs = model(batch_x)
        loss = criterion(outputs, batch_y)

    # Scale loss, backward, unscale, step
    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    scaler.step(optimizer)
    scaler.update()

Key points:

• autocast wraps forward pass + loss computation only
• GradScaler prevents underflow in small gradients
• Use bfloat16 on Ampere+ GPUs (less numerical issues than float16)
• AMP is most beneficial for large models (CNNs, Transformers) — small models may not see speedup

---

torch.compile (PyTorch 2.x JIT Compilation)

torch.compile compiles model graphs for faster execution with minimal code changes.

# Basic usage (reduces model — returns a compiled wrapper)
compiled_model = torch.compile(model)

# With options
compiled_model = torch.compile(
    model,
    mode="reduce-overhead",  # "default", "reduce-overhead", "max-autotune"
    fullgraph=True,          # Fail if graph breaks
    dynamic=True,            # Support dynamic tensor shapes
)

# Training is identical:
output = compiled_model(batch_x)  # First call compiles, subsequent calls are fast

When to use: Models with large tensor operations (CNNs, Transformers). Not beneficial for very small models or heavy data loading bottlenecks.

When not to use: Dynamic control flow, custom CUDA extensions, models compiled for TorchScript export.

Mode selection:

• "default" — conservative, works for most models
• "reduce-overhead" — good for inference, reduces Python overhead
• "max-autotune" — benchmarks and picks the best backend (slow first run)

---

Loss Functions

Problem Type	Loss Function	Import
Binary classification	BCEWithLogitsLoss	torch.nn.BCEWithLogitsLoss
Multi-class classification	CrossEntropyLoss	torch.nn.CrossEntropyLoss
Multi-label classification	BCEWithLogitsLoss	Combines sigmoid + BCELoss
Regression (MSE)	MSELoss	torch.nn.MSELoss
Regression (MAE)	L1Loss	torch.nn.L1Loss
Regression (Huber)	HuberLoss	torch.nn.HuberLoss (delta parameter)
Contrastive / Siamese	TripletMarginLoss / ContrastiveLoss	torch.nn.TripletMarginLoss
Imbalanced classes	Weighted CrossEntropyLoss	Pass weight tensor to constructor
Sequence (CTC)	CTCLoss	torch.nn.CTCLoss

# Weighted loss for imbalanced data
class_weights = torch.tensor([0.2, 0.8]).to(device)  # inverse frequency
criterion = torch.nn.CrossEntropyLoss(weight=class_weights)

---

Optimizers & Schedulers

Optimizer	When	Learning Rate
AdamW	Default for most models (Transformers, CNNs)	1e-4 to 3e-4
Adam	Legacy; prefer AdamW (proper weight decay)	1e-4 to 3e-4
SGD + momentum	When Adam overfits; vision models	1e-2 to 1e-1
AdamW (LoRA)	Fine-tuning with LoRA	2e-4 to 5e-4

import torch.optim as optim

# AdamW — the default
optimizer = optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0.01)

# SGD with momentum (for vision fine-tuning)
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=1e-4)

Learning Rate Schedulers

from torch.optim.lr_scheduler import (
    ReduceLROnPlateau,  # Reduce when metric plateaus
    CosineAnnealingLR,  # Cosine decay
    OneCycleLR,         # Warmup + cosine ("super-convergence")
)

# Reduce on plateau (works well for most tasks)
scheduler = ReduceLROnPlateau(optimizer, mode="min", patience=5, factor=0.5)

# One-cycle (fast convergence, needs max_lr)
scheduler = OneCycleLR(
    optimizer,
    max_lr=1e-3,
    steps_per_epoch=len(train_loader),
    epochs=num_epochs,
)

# Cosine annealing (good for Transformers)
scheduler = CosineAnnealingLR(optimizer, T_max=num_epochs)

# Warmup + cosine (standard for LLM fine-tuning)
# Implement manually:
def get_cosine_schedule_with_warmup(optimizer, warmup_steps, total_steps):
    def lr_lambda(current_step):
        if current_step < warmup_steps:
            return float(current_step) / float(max(1, warmup_steps))
        progress = float(current_step - warmup_steps) / float(max(1, total_steps - warmup_steps))
        return 0.5 * (1.0 + math.cos(math.pi * progress))
    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)

---

Transfer Learning

Feature Extraction (Freeze Backbone)

import torchvision.models as models

model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)

# Freeze all layers
for param in model.parameters():
    param.requires_grad = False

# Replace classifier
num_features = model.fc.in_features
model.fc = torch.nn.Linear(num_features, num_classes)

# Only the new layer's params are trainable
optimizer = optim.AdamW(model.fc.parameters(), lr=1e-3)

Full Fine-Tuning

# Unfreeze all layers
for param in model.parameters():
    param.requires_grad = True

# Use lower learning rate
optimizer = optim.AdamW(model.parameters(), lr=2e-5)

# Often helps to freeze early layers, fine-tune later layers:
for name, param in model.named_parameters():
    if "layer1" in name or "layer2" in name:
        param.requires_grad = False

LoRA for Transformer Models

LoRA (Low-Rank Adaptation) trains small rank-decomposition matrices while keeping the base model frozen. Requires the peft library.

from peft import LoraConfig, get_peft_model

lora_config = LoraConfig(
    r=8,                    # Rank — higher = more expressiveness, more params
    lora_alpha=32,          # Scaling factor
    target_modules=["q_proj", "v_proj"],  # Which modules to apply LoRA to
    lora_dropout=0.1,
    bias="none",            # Don't train bias terms
    task_type="SEQ_CLS",    # Task type (SEQ_CLS, CAUSAL_LM, TOKEN_CLS, etc.)
)

model = get_peft_model(base_model, lora_config)

# Train with slightly higher LR
optimizer = optim.AdamW(model.parameters(), lr=3e-4)

# LoRA adds very few params (< 1% of base model)
model.print_trainable_parameters()  # e.g., "trainable params: 294,912 || all params: 110,080,512 || %: 0.2679"

---

Knowledge Distillation

import torch.nn.functional as F

def distillation_loss(student_logits, teacher_logits, labels, temperature=4.0, alpha=0.7):
    """
    Combined distillation + supervised loss.

    Args:
        temperature: Higher = softer probability distribution (more information from teacher)
        alpha: Weight for distillation loss (vs standard cross-entropy)
    """
    # Soft target loss (distillation)
    soft_student = F.log_softmax(student_logits / temperature, dim=-1)
    soft_teacher = F.softmax(teacher_logits.detach() / temperature, dim=-1)
    distill_loss = F.kl_div(soft_student, soft_teacher, reduction="batchmean")
    distill_loss *= temperature ** 2  # Scale to keep gradients in right range

    # Hard target loss
    hard_loss = F.cross_entropy(student_logits, labels)

    return alpha * distill_loss + (1 - alpha) * hard_loss


# Training loop with distillation
teacher_model.eval()
for batch_x, batch_y in dataloader:
    with torch.no_grad():
        teacher_logits = teacher_model(batch_x)

    student_logits = student_model(batch_x)
    loss = distillation_loss(student_logits, teacher_logits, batch_y)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

Temperature tuning: Start with T=4.0. Higher values produce softer targets (more small-class information). Lower values (T=1.0) collapse to standard cross-entropy.

Alpha tuning: α=0.7 (weight on distillation) is a common starting point. Increase α when the teacher is much better than the student.

---

Model Pruning

import torch.nn.utils.prune as prune

# Apply pruning to specific layers
prune.l1_unstructured(module=model.fc, name="weight", amount=0.3)  # Remove 30% of weights

# Make pruning permanent (removes the pruning mask)
prune.remove(module=model.fc, name="weight")

# Structured pruning (removes entire neurons/channels)
prune.ln_structured(module=model.conv1, name="weight", amount=0.2, n=2, dim=0)

# Global pruning (prune all layers together by importance)
parameters_to_prune = [
    (model.layer1, "weight"),
    (model.layer2, "weight"),
    (model.fc, "weight"),
]
prune.global_unstructured(
    parameters_to_prune,
    pruning_method=prune.L1Unstructured,
    amount=0.2,  # Remove 20% of weights globally
)

After pruning: Fine-tune the pruned model. Pruning then fine-tuning almost always recovers accuracy. Pruning without fine-tuning degrades performance significantly.

---

Distributed Data Parallel (DDP)

For multi-GPU training. DDP wraps the model and handles gradient synchronization.

import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP

def setup(rank, world_size):
    """Initialize the distributed process group."""
    dist.init_process_group(
        backend="nccl",              # "nccl" for NVIDIA, "gloo" for CPU
        init_method="env://",         # Use env vars MASTER_ADDR and MASTER_PORT
        rank=rank,
        world_size=world_size,
    )

def cleanup():
    dist.destroy_process_group()

def train(rank, world_size):
    setup(rank, world_size)

    # Model must be on the correct device BEFORE wrapping
    model = YourModel().to(rank)
    ddp_model = DDP(model, device_ids=[rank])

    # DataLoader must use DistributedSampler
    sampler = torch.utils.data.distributed.DistributedSampler(
        dataset, num_replicas=world_size, rank=rank
    )
    dataloader = DataLoader(dataset, batch_size=32, sampler=sampler)

    # Training loop (same structure, use ddp_model)
    for epoch in range(num_epochs):
        sampler.set_epoch(epoch)  # Shuffle each epoch
        for batch_x, batch_y in dataloader:
            batch_x, batch_y = batch_x.to(rank), batch_y.to(rank)
            outputs = ddp_model(batch_x)
            loss = criterion(outputs, batch_y)
            loss.backward()
            optimizer.step()

    cleanup()

# Launch
if __name__ == "__main__":
    world_size = torch.cuda.device_count()
    mp.spawn(train, args=(world_size,), nprocs=world_size)

When DDP is worth it: Models that take > 1 hour to train on a single GPU. For short experiments, single-GPU + AMP is often faster due to communication overhead.

---

Debugging

Common Failure Patterns

Symptom	Likely Cause	Fix
loss = nan	Exploding gradients, bad learning rate	Lower LR, add gradient clipping, check for NaN in input data
loss = nan after AMP	Gradient underflow	Increase GradScaler init_scale, or use bfloat16
Loss doesn't decrease	Wrong LR, wrong loss function	Check LR range, verify loss function matches task
CUDA OOM	Batch size too large	Reduce batch size, enable gradient checkpointing, use AMP
Expected all tensors to be on...	Device mismatch	Always .to(device) tensors before model forward
CUDA error: device-side assert	Wrong label class (out of range)	Check label values are in [0, num_classes)
Model doesn't overfit 1 batch	Bug in model architecture	Try overfitting on a single batch (batch of 2-4 samples for 100 steps)

Overfit on One Batch (Diagnostic)

# If model can't overfit a single batch, something is fundamentally wrong
single_batch = next(iter(dataloader))

for step in range(100):
    outputs = model(single_batch[0].to(device))
    loss = criterion(outputs, single_batch[1].to(device))
    loss.backward()
    optimizer.step()

print(f"Loss after 100 steps on 1 batch: {loss.item():.6f}")
# If loss is not near 0, model has a bug or LR is way off.

Gradient Checking

# Check gradient norms during training
total_norm = 0.0
for p in model.parameters():
    if p.grad is not None:
        param_norm = p.grad.data.norm(2)
        total_norm += param_norm.item() ** 2
total_norm = total_norm ** 0.5

if total_norm > 10.0:
    print(f"WARNING: Large gradient norm: {total_norm:.4f}")

---

Reproducibility

import random
import numpy as np
import torch

def set_seed(seed: int = 42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False  # Trades off speed for determinism

set_seed(42)

Note on cudnn.deterministic: Setting this to True may reduce performance (5-15%). For experiments where exact reproducibility isn't critical, leave benchmark=True and accept minor stochastic variation.

---

See Also

• references/experimental-campaign-protocol.md — where this reference fits in the campaign workflow
• references/sklearn-integration.md — for non-deep-learning methods
• scripts/detect-compute.py — check your hardware before choosing model size
• pytorch.org/docs/stable — official API documentation