ai_ml / ml frameworks / 01_what_is_pytorch.md

What is PyTorch?

4 interview angles 4 min read source

What is PyTorch?

PyTorch is an open-source deep learning framework originally from Meta AI (released 2016), now under the PyTorch Foundation (Linux Foundation). It became the de facto standard for ML research and is now dominant in production too — most modern LLMs, including Llama, Mistral, and most Hugging Face models, are PyTorch-native.

Core idea

PyTorch gives you:

  1. Tensors — n-dimensional arrays (like NumPy) with GPU acceleration.
  2. Autograd — automatic differentiation by recording operations on a dynamic computation graph.
  3. nn.Module — composable building blocks for neural networks.
  4. Optimizers and losses — SGD/Adam/AdamW, cross-entropy, MSE, etc.
  5. DataLoader — batched, shuffled, parallelized data loading.

The philosophy is “Pythonic and imperative” — you write Python, run it, and it works. No graph-compilation step you have to think about (unlike TensorFlow 1.x).

Tensors — the foundation

import torch

x = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
y = torch.randn(2, 2)                       # random from N(0,1)
z = x @ y + 1                               # matmul + broadcast
print(z.shape, z.dtype, z.device)           # torch.Size([2,2]) float32 cpu

# Move to GPU
x = x.to("cuda")                            # or torch.device("cuda:0")

Tensor API mirrors NumPy where reasonable but adds: device, requires_grad, autograd, deferred CUDA execution.

Autograd — the killer feature

Set requires_grad=True, do operations, call .backward(), gradients appear in .grad:

x = torch.tensor(2.0, requires_grad=True)
y = x ** 3 + 2 * x
y.backward()
print(x.grad)   # 3*x^2 + 2 = 14

PyTorch builds the computation graph dynamically as you run code — this is “define-by-run” / eager mode. Branching, loops, dynamic shapes all just work because the graph is rebuilt each forward pass.

A typical training loop

import torch
from torch import nn
from torch.utils.data import DataLoader

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        return self.fc2(x)

model = Net().to("cuda")
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()

for epoch in range(10):
    for xb, yb in train_loader:
        xb, yb = xb.to("cuda"), yb.to("cuda")
        logits = model(xb)
        loss = loss_fn(logits, yb)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

Notice how literal it is — no tf.GradientTape, no session.run, no compile(). This is the appeal.

Key concepts

Concept What it is
torch.Tensor n-d array with optional gradients, on CPU or GPU
nn.Module base class for layers / models; tracks parameters automatically
nn.Parameter a tensor registered as learnable when assigned to a Module
torch.optim optimizers (SGD, Adam, AdamW, RMSprop)
autograd.grad / .backward() computes gradients
DataLoader batches a Dataset, with shuffling, multi-worker loading
torch.no_grad() context manager that disables grad tracking (use during eval/inference)
model.train() / model.eval() toggle dropout / batchnorm modes
torch.compile() (2.0+) JIT-compile model to fused kernels for speed

The PyTorch ecosystem

  • torchvision — image models, datasets, transforms.
  • torchaudio — audio data / models.
  • torchtext — text utilities (legacy; HF tokenizers now dominate).
  • PyTorch Lightning — high-level training loop framework; reduces boilerplate.
  • Hugging Face transformers — pretrained LLMs/vision models, PyTorch-first.
  • accelerate, deepspeed, FSDP — distributed / sharded training.
  • torch.compile + Inductor — graph compiler for speed.
  • TorchScript / ONNX / ExecutorTorch — deployment paths.

Why PyTorch dominated research

  • Dynamic graphs — you can use Python control flow naturally. TF 1.x required tracing in advance.
  • Debuggability — set a breakpoint anywhere; print tensors mid-computation; standard Python stack traces.
  • NumPy-like API — small learning curve for scientific Python users.
  • Strong academic adoption — most new papers ship PyTorch code.

By 2022, ~80%+ of papers at top ML conferences used PyTorch. The Hugging Face ecosystem cemented its production lead.

Production paths

PyTorch was historically weaker than TF for production deploys, but that gap has closed:

  • TorchServe — model server.
  • ONNX export — interop with other runtimes (TensorRT, OpenVINO).
  • torch.compile — 2-3x speedups for many models.
  • Quantization — int8 / int4 via torch.ao.quantization.
  • ExecutorTorch — mobile / edge deployment.
  • vLLM, TensorRT-LLM, Triton Inference Server — serving for LLMs.

Common gotchas

  • Forgetting optimizer.zero_grad() — gradients accumulate by default; without zeroing, you train on the wrong gradient.
  • Forgetting model.eval() at inference — dropout stays on, batchnorm uses batch stats; metrics get noisy.
  • Forgetting torch.no_grad() at inference — wastes memory on the autograd graph.
  • .cuda() then optimizer = ... — always create optimizer after moving the model to GPU.
  • Modifying tensors in-placex += 1 on a tensor with requires_grad can break autograd; prefer x = x + 1.
  • DataLoader num_workers > 0 on Windows — needs if __name__ == "__main__": guard.

Where to go next

Interview angle

  • “What is PyTorch and why is it popular?” — Pythonic, define-by-run deep learning framework with dynamic graphs, autograd, and a NumPy-like tensor API. Popular because of debuggability, research-friendliness, and the HF ecosystem.
  • “What’s a nn.Module?” — base class for neural network components. Tracks parameters automatically (any nn.Parameter or sub-Module assigned as an attribute is registered). Defines forward() for the computation.
  • “How does autograd work?” — PyTorch records every operation on tensors with requires_grad=True into a computation graph. Calling .backward() traverses that graph in reverse, applying the chain rule to populate .grad on leaf tensors.
  • “What does optimizer.zero_grad() do and why is it needed?” — clears .grad on all parameters. PyTorch accumulates gradients across .backward() calls (useful for gradient accumulation), so you must explicitly zero them at the start of each step.