PyTorch and TensorFlow

Tensors, models, losses, data pipelines, training loops, optimizers, mixed precision, checkpoints, and debugging.

Tensor Basics

Python

import torch
x = torch.randn(32, 3, 224, 224)  # (N, C, H, W)
x.shape, x.dtype, x.device
x.view(32, -1)          # needs contiguous memory
x.reshape(32, -1)       # copies if needed
x.permute(0, 2, 3, 1)   # NCHW -> NHWC
x.unsqueeze(1); x.squeeze()

Python

import tensorflow as tf
t = tf.random.normal([32, 224, 224, 3])  # NHWC
tf.reshape(t, [32, -1]); tf.transpose(t, [0, 3, 1, 2])

Layout: PyTorch defaults to NCHW, TensorFlow to NHWC.
Broadcasting: Same right-aligned rules as NumPy.

Devices

Python

device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
xb, yb = xb.to(device), yb.to(device)

Same device: Model and batch must live on the same device.
Pinned memory: DataLoader(pin_memory=True) speeds host-to-GPU copies.
TensorFlow: Places ops on the GPU automatically; use tf.device("/CPU:0") to override.

Defining Models

Python

import torch.nn as nn

class MLP(nn.Module):
    def __init__(self, d_in, d_hidden, n_cls):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d_in, d_hidden), nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(d_hidden, n_cls))

    def forward(self, x):
        return self.net(x)   # logits

Python

model = tf.keras.Sequential([
    tf.keras.layers.Dense(256, activation="relu"),
    tf.keras.layers.Dropout(0.1),
    tf.keras.layers.Dense(n_cls),   # logits
])

Return logits: Let the loss apply softmax or sigmoid for numerical stability.
Parameters: sum(p.numel() for p in model.parameters()) in PyTorch; model.summary() in Keras.

Common Layers

Layer	PyTorch	Keras
Dense	nn.Linear	Dense
Convolution	nn.Conv2d	Conv2D
Normalization	nn.BatchNorm2d, nn.LayerNorm	BatchNormalization, LayerNormalization
Embedding	nn.Embedding	Embedding
Recurrent	nn.LSTM, nn.GRU	LSTM, GRU
Attention	nn.MultiheadAttention	MultiHeadAttention

H_out = ⌊(H_in + 2p − k) / s⌋ + 1

Conv output size: kernel k, padding p, stride s

Attention

Attention(Q, K, V) = softmax(Q K^T / √dₖ) · V

Python

import torch.nn.functional as F
out = F.scaled_dot_product_attention(
    q, k, v, is_causal=True)   # (B, heads, T, d)

Causal mask: Blocks attention to future positions in decoders.
Fused kernels: SDPA dispatches to FlashAttention when shapes and dtypes allow.

Loss Functions

Task	PyTorch	Keras
Multiclass	CrossEntropyLoss (logits)	SparseCategoricalCrossentropy(from_logits=True)
Binary / multilabel	BCEWithLogitsLoss	BinaryCrossentropy(from_logits=True)
Regression	MSELoss, HuberLoss	MeanSquaredError, Huber

Pass logits, not probabilities

CrossEntropyLoss applies log-softmax itself. Adding softmax in the model first makes training slow and unstable.

Data Pipelines

Python

from torch.utils.data import DataLoader, Dataset

class Rows(Dataset):
    def __init__(self, X, y): self.X, self.y = X, y
    def __len__(self): return len(self.X)
    def __getitem__(self, i): return self.X[i], self.y[i]

loader = DataLoader(Rows(X, y), batch_size=64,
                    shuffle=True, num_workers=4)

Python

ds = (tf.data.Dataset.from_tensor_slices((X, y))
      .shuffle(10_000).batch(64)
      .prefetch(tf.data.AUTOTUNE))

Shuffle training data only: Keep validation and test order fixed.
Workers / prefetch: Overlap data loading with GPU compute so the GPU never waits.
drop_last: Avoids a tiny final batch that can destabilize BatchNorm.

PyTorch Training Loop

Forward
Loss
zero_grad
Backward
Optimizer step

Clear old gradients before backward, then step

Train

Python

for epoch in range(epochs):
    model.train()
    for xb, yb in loader:
        xb, yb = xb.to(device), yb.to(device)
        logits = model(xb)
        loss = loss_fn(logits, yb)
        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        optimizer.step()
    scheduler.step()

Evaluate

Python

model.eval()
with torch.no_grad():
    preds = model(x_val.to(device)).argmax(dim=1)

Keras Training

Python

model.compile(
    optimizer=tf.keras.optimizers.AdamW(1e-3),
    loss=tf.keras.losses.SparseCategoricalCrossentropy(
        from_logits=True),
    metrics=["accuracy"])
model.fit(ds, validation_data=val_ds, epochs=10,
          callbacks=[tf.keras.callbacks.EarlyStopping(
              patience=3, restore_best_weights=True)])

compile: Binds optimizer, loss, and metrics to the model.
fit: Runs the training loop, validation, and callbacks for you.
Custom loops: Use tf.GradientTape when you need full control over each step.

Optimizers and Schedules

SGD + momentum: Strong for vision with a tuned schedule.
Adam: Adaptive per-parameter steps; fast convergence.
AdamW: Decoupled weight decay; the default for transformers.
Warmup + cosine decay: Common schedule for large models.
Gradient clipping: clip_grad_norm_(params, 1.0) prevents exploding updates.

θ ← θ − η (m̂ / (√v̂ + ε) + λθ)

AdamW: decay λθ applied separately from the adaptive step

Regularization and Initialization

Weight decay: Penalizes large weights; exclude biases and norm layers.
Dropout: Randomly zeroes activations in training only.
Data augmentation: Often the strongest regularizer for vision and audio.
Early stopping: Stop when the validation metric stops improving.
Init: He (Kaiming) for ReLU, Xavier (Glorot) for tanh; frameworks default sensibly.

Mixed Precision

Python

scaler = torch.amp.GradScaler("cuda")
with torch.autocast("cuda", dtype=torch.float16):
    loss = loss_fn(model(xb), yb)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

bfloat16: Same range as FP32; usually needs no loss scaling.
Keras: tf.keras.mixed_precision.set_global_policy("mixed_float16").

Checkpoints and Reproducibility

Python

torch.save({"model": model.state_dict(),
            "opt": optimizer.state_dict(),
            "epoch": epoch}, "ckpt.pt")
ckpt = torch.load("ckpt.pt", map_location=device)
model.load_state_dict(ckpt["model"])
torch.manual_seed(0)

Save state_dict: More portable than pickling the whole model.
Keras: model.save("model.keras") and tf.keras.utils.set_random_seed(0).
Determinism: Seeds alone do not make GPU kernels deterministic.

Shape Conventions

Data	Typical shape
Tabular batch	(N, features)
Images (PyTorch)	(N, C, H, W)
Sequences	(N, T, d) or (T, N, d)
Attention scores	(N, heads, T, T)
Logits (classification)	(N, classes)

Inference and Deployment

eval() + no_grad(): Disable dropout and gradient tracking for inference.
TorchScript / torch.compile: Optimize and package PyTorch models.
ONNX: Framework-neutral export for many runtimes.
SavedModel / TF Serving: Standard TensorFlow deployment path.
Quantization: INT8 or FP16 weights for faster, smaller models.

Debug Checklist

Overfit one small batch; if the loss won't drop, the bug is in the code.
Check input shapes, dtypes, and value ranges at the model boundary.
Confirm labels match the loss (class indices vs one-hot).
Verify train() and eval() modes are switched correctly.
Plot loss curves; NaN usually means the learning rate is too high.
Check that gradients are non-zero for every trainable layer.

Start with a tiny model and a small subset of data, then scale up once the full loop works end to end.