PyTorch or Keras tour

Why this matters

You have written a forward pass, a backward pass and an update loop by hand. At work nobody does that: you use a deep learning framework that computes gradients automatically, runs on GPUs and ships tested layers. The two you will meet most are PyTorch (dominant in research and increasingly in industry) and Keras (a high-level API that runs on TensorFlow, JAX or PyTorch). Their vocabulary is the same: tensors, models, losses, optimisers and a training loop. Learn it once and you can read either.

Where to run this

Neither framework runs in this site's browser sandbox, so this lesson is reading and quiz only. To try the code, use a free notebook: Google Colab or Kaggle Notebooks both have PyTorch and Keras preinstalled and offer a free GPU with usage limits. Locally, pip install torch (see pytorch.org for the right command for your machine) or pip install keras tensorflow inside a virtual environment.

Tensors

A tensor is an n-dimensional array, like a NumPy array, with two extras: it can live on a GPU, and it can record the operations applied to it so gradients can be computed.

import torch

x = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
print(x.shape, x.dtype)        # torch.Size([2, 2]) torch.float32
w = torch.randn(2, 1, requires_grad=True)
y = (x @ w).sum()
y.backward()                   # autograd fills w.grad
print(w.grad)                  # dy/dw: the column sums of x, [[4.], [6.]]

device = "cuda" if torch.cuda.is_available() else "cpu"
x = x.to(device)               # move data to the GPU when there is one

requires_grad=True tells PyTorch to track operations on w. Calling .backward() on a scalar runs backpropagation for you, which is exactly what you wrote by hand last lesson.

Models

In PyTorch a model is a class that inherits from nn.Module. You declare layers in __init__ and describe the computation in forward:

from torch import nn

class TwoLayerNet(nn.Module):
    def __init__(self, n_in, n_hidden):
        super().__init__()
        self.hidden = nn.Linear(n_in, n_hidden)   # holds W1 and b1
        self.out = nn.Linear(n_hidden, 1)         # holds W2 and b2

    def forward(self, x):
        h = torch.tanh(self.hidden(x))
        return self.out(h).squeeze(-1)            # raw scores (logits)

model = TwoLayerNet(2, 8)

The same network in Keras is declarative:

import keras

model = keras.Sequential([
    keras.Input(shape=(2,)),
    keras.layers.Dense(8, activation="tanh"),
    keras.layers.Dense(1, activation="sigmoid"),
])

Loss and optimiser

The loss function measures error; the optimiser updates the weights from the gradients. SGD is plain gradient descent; Adam adapts a step size per weight and is the usual default.

loss_fn = nn.BCEWithLogitsLoss()     # sigmoid + cross-entropy, numerically stable
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)

Note that the PyTorch model returns logits (raw scores) and the loss applies the sigmoid. Combining them avoids overflow, so do not add your own sigmoid before BCEWithLogitsLoss.

The training loop

PyTorch makes you write the loop, which is a little more code and a lot more control:

from torch.utils.data import DataLoader, TensorDataset

loader = DataLoader(TensorDataset(X_train, y_train), batch_size=32, shuffle=True)

for epoch in range(20):
    model.train()                     # enables dropout and similar layers
    for xb, yb in loader:
        optimizer.zero_grad()         # clear old gradients
        loss = loss_fn(model(xb), yb) # forward pass
        loss.backward()               # backward pass
        optimizer.step()              # update weights

    model.eval()                      # disables dropout
    with torch.no_grad():             # no gradient tracking during evaluation
        val_loss = loss_fn(model(X_val), y_val).item()
    print(epoch, round(val_loss, 4))

The five lines inside the batch loop are the whole of training: zero the gradients, forward, compute the loss, backward, step. Keras wraps all of it in two calls:

model.compile(optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"])
history = model.fit(X_train, y_train, validation_data=(X_val, y_val),
                    epochs=20, batch_size=32)

history.history holds the per-epoch training and validation loss, ready for the over- and underfitting diagnosis from the previous lesson.

Saving and loading

torch.save(model.state_dict(), "model.pt")     # PyTorch: save the weights
model.load_state_dict(torch.load("model.pt"))

model.save("model.keras")                       # Keras: architecture + weights

Choosing between them

• PyTorch: explicit loops, easy debugging with ordinary Python, the largest ecosystem of research code and pretrained models (Hugging Face is PyTorch-first).
• Keras: less code for standard models, quick to prototype, built-in callbacks such as EarlyStopping.
• Both export models for serving (TorchScript or ONNX for PyTorch; SavedModel or ONNX for TensorFlow/Keras). Pick the one your team uses.

Common mistakes

• Forgetting optimizer.zero_grad(). PyTorch adds new gradients to old ones, so steps get larger and wrong.
• Forgetting model.eval() and torch.no_grad() when validating, which leaves dropout on and wastes memory.
• Device mismatch. The model and every batch must be on the same device: RuntimeError: Expected all tensors to be on the same device.
• Wrong dtype or shape for labels. BCEWithLogitsLoss wants float labels with the same shape as the logits.
• Double sigmoid. A sigmoid output layer plus BCEWithLogitsLoss applies it twice.

Quick recap

• Tensors are GPU-capable arrays that can track gradients; .backward() runs autograd.
• A model holds layers and defines forward; the optimiser (usually Adam) updates its parameters.
• One training step: zero gradients, forward, loss, backward, step. Validate with eval() and no_grad().
• Keras compile + fit wraps the same loop.