PyTorch Foundations

Introduction

While the theoretical math of neural networks is beautiful, implementing forward passes and backpropagation from scratch in NumPy for a multi-billion parameter model is practically impossible. We need a framework that can handle the heavy lifting, automate calculus, and run on ultra-fast hardware (GPUs).

Enter PyTorch, the undisputed king of deep learning research and increasingly the standard for production systems. Developed by Meta (Facebook) AI, PyTorch provides a pythonic, flexible, and dynamic way to build neural networks. Unlike older frameworks that required compiling static computational graphs before running anything, PyTorch evaluates code dynamically (eager execution). This makes it feel exactly like standard Python—you can use print statements, loops, and standard debugging tools directly on your neural networks.

At its core, PyTorch provides two massive upgrades over standard Python:

1. Tensors: Multi-dimensional arrays (like NumPy arrays) that can run lightning-fast on NVIDIA GPUs via CUDA.
2. Autograd: An automatic differentiation engine that calculates gradients for backpropagation automatically, so you never have to write chain rule calculus by hand again.

Core Concepts Explained Intuitively

1. Tensors

A Tensor is simply a multi-dimensional grid of numbers.

• A 0-dimensional tensor is a scalar (a single number: 5).
• A 1-dimensional tensor is a vector (a list of numbers: [1, 2, 3]).
• A 2-dimensional tensor is a matrix (a grid of numbers, like a grayscale image).
• A 3-dimensional tensor is a cube (like a color image with Red, Green, Blue channels).
• A 4-dimensional tensor is a batch of color images.

While they look exactly like NumPy arrays, PyTorch Tensors have a special superpower: they can be moved to a GPU (tensor.to('cuda')), which accelerates matrix multiplications by orders of magnitude.

2. Autograd (Automatic Differentiation)

When you create a tensor and set requires_grad=True, PyTorch starts tracking every operation performed on that tensor. It builds a computational graph in the background. When you finish your forward pass and calculate the final loss, you simply call loss.backward(). PyTorch will traverse the graph backward, instantly computing the exact gradient of the loss with respect to every single tensor that had requires_grad=True.

3. nn.Module

torch.nn.Module is the base class for all neural network architectures in PyTorch. To create your own network, you subclass nn.Module. Inside your class, you define two things:

1. __init__(): Where you define the layers (e.g., linear layers, convolutions) your network will use.
2. forward(x): Where you define exactly how the input x flows through the layers to produce an output.

4. Dataset and DataLoader

Machine learning datasets are often terabytes in size—way too big to fit into RAM all at once.

• torch.utils.data.Dataset: An abstract class representing your data. You write custom code to load a single sample (like an image) from disk.
• torch.utils.data.DataLoader: Wraps the Dataset and automatically handles batching, shuffling, and multi-processing to feed data into your GPU as fast as possible.

Code Examples

1. Beginner Code Example: Tensors and Autograd

Let's see the magic of automatic differentiation in action. We will define an equation $y = 3x^2 + 4x + 2$ and let PyTorch calculate the derivative for us.

import torch

# Create a tensor for x. We set requires_grad=True because we want 
# PyTorch to calculate dy/dx
x = torch.tensor(2.0, requires_grad=True)

# Define the forward equation
y = 3 * x**2 + 4 * x + 2

print(f"y value at x=2: {y.item()}")

# MAGIC HAPPENS HERE
# Calculate the gradients automatically
y.backward()

# The mathematical derivative of y = 3x^2 + 4x + 2 is dy/dx = 6x + 4
# At x = 2, dy/dx = 6(2) + 4 = 16
# Let's check what PyTorch calculated:
print(f"PyTorch calculated gradient (dy/dx): {x.grad.item()}")

2. Intermediate Code Example: Building a Neural Network with nn.Module

Here we define a classic Multi-Layer Perceptron using PyTorch's object-oriented approach.

import torch
import torch.nn as nn
import torch.nn.functional as F

class SimpleMLP(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super(SimpleMLP, self).__init__() # Call the parent class init
        
        # Define the layers
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, num_classes)
        
    def forward(self, x):
        # Define the forward pass logic
        x = self.fc1(x)
        x = F.relu(x) # Apply activation function
        x = self.fc2(x)
        return x

# Instantiate the model
# e.g., for MNIST images (28x28 = 784 inputs), 128 hidden neurons, 10 output classes
model = SimpleMLP(input_size=784, hidden_size=128, num_classes=10)
print(model)

# Pass a dummy batch of 5 images through the network
dummy_input = torch.randn(5, 784) 
output = model(dummy_input)
print(f"Output shape: {output.shape}") # Should be [5, 10]

3. Advanced Code Example: A Complete PyTorch Training Loop

This is the standard boilerplate required to train a model in PyTorch on a GPU.

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import TensorDataset, DataLoader

# 1. Setup Device (Use GPU if available, else CPU)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Training on device: {device}")

# 2. Create Dummy Data and DataLoader
X_train = torch.randn(1000, 20)
y_train = torch.randint(0, 2, (1000,)) # Binary labels: 0 or 1
dataset = TensorDataset(X_train, y_train)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

# 3. Define Model, Loss, and Optimizer
model = nn.Sequential(
    nn.Linear(20, 64),
    nn.ReLU(),
    nn.Linear(64, 2)
).to(device) # Move model to GPU

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 4. The Training Loop
epochs = 5
for epoch in range(epochs):
    epoch_loss = 0.0
    for batch_X, batch_y in dataloader:
        # Move data to GPU
        batch_X, batch_y = batch_X.to(device), batch_y.to(device)
        
        # Step A: Zero the gradients (crucial!)
        optimizer.zero_grad()
        
        # Step B: Forward pass
        predictions = model(batch_X)
        
        # Step C: Calculate Loss
        loss = criterion(predictions, batch_y)
        
        # Step D: Backward pass (calculate gradients)
        loss.backward()
        
        # Step E: Update weights
        optimizer.step()
        
        epoch_loss += loss.item()
        
    print(f"Epoch {epoch+1}/{epochs}, Loss: {epoch_loss/len(dataloader):.4f}")

Common Mistakes

1. Forgetting optimizer.zero_grad():

   • Mistake: In PyTorch, gradients accumulate by default. If you don't zero them out at the start of a training loop iteration, the gradients from batch 2 will be added to the gradients from batch 1, pointing the optimizer in completely the wrong direction.
   • Fix: Always call optimizer.zero_grad() before loss.backward().

2. Device Mismatches (CPU vs GPU):

   • Mistake: You moved your model to the GPU (model.to('cuda')), but your data tensor is still on the CPU. PyTorch will crash with a RuntimeError because it cannot multiply a CPU tensor with a GPU tensor.
   • Fix: Ensure both model and data are explicitly moved to the same device using .to(device).

3. Mixing up .view() and .reshape():

   • Mistake: Getting confused about how to flatten a tensor for a Linear layer.
   • Fix: Use .reshape() generally. .view() is strictly for memory-contiguous tensors and will sometimes throw errors. Both achieve the same logical outcome (changing the shape of the tensor without altering data).

Interview Q&A

Q1: What is the difference between PyTorch and TensorFlow? A1: Historically, TensorFlow (1.x) used static computational graphs (define-and-run), making it harder to debug but easier to deploy. PyTorch uses dynamic computational graphs (define-by-run), making it highly intuitive, pythonic, and standard for research. Today (with TensorFlow 2.x), they are very similar, but PyTorch dominates the NLP and research ecosystems (e.g., HuggingFace).

Q2: How does PyTorch handle computational graphs? A2: PyTorch builds a Directed Acyclic Graph (DAG) dynamically on the fly during the forward pass. Tensors are nodes, and the mathematical operations are the edges. Once loss.backward() is called, PyTorch traverses this graph to compute gradients, and then the graph is immediately destroyed to free memory. A new graph is built on the next forward pass.

Q3: Why do we need torch.no_grad() during inference? A3: During inference (testing/evaluating), we do not need to update weights, so we don't need gradients. Wrapping your prediction code in with torch.no_grad(): tells PyTorch to stop tracking operations and building the computational graph. This massively reduces memory consumption and speeds up inference.

Q4: What is model.train() and model.eval()? A4: These methods toggle the mode of the model. Layers like Dropout and BatchNorm behave differently during training vs testing. Dropout drops neurons randomly during training to prevent overfitting, but during testing, it must be turned off (via model.eval()) so the network uses all its neurons for the best prediction.

Q5: Explain what loss.item() does and why it's used. A5: loss is a PyTorch Tensor containing a single scalar value, but it is still attached to the computational graph. Calling .item() extracts the raw Python float value from the tensor and detaches it from the graph. If you just append loss to a list to track history, you will keep the entire computational graph in memory, leading to a memory leak and OOM crash.

Practice Problems

1. Problem: Create a 2D tensor of shape (3, 3) initialized with random numbers. Use PyTorch operations to find the mean, the sum, and the max value of this tensor.
2. Problem: Write a custom torch.utils.data.Dataset class that generates 100 random image tensors of shape (3, 32, 32) and labels them 0 or 1. Wrap it in a DataLoader and print the shape of the first batch.
3. Problem: Modify the Advanced Training Loop example to include an evaluation phase. After every epoch, calculate the accuracy on a validation dataset. (Hint: Remember to use model.eval() and torch.no_grad()).

Summary

• Tensors are n-dimensional arrays that can run on GPUs for hardware acceleration.
• Autograd tracks tensor operations to automatically compute gradients via backpropagation.
• nn.Module is the standard way to define neural network architectures using object-oriented Python.
• DataLoaders efficiently manage massive datasets, handling batching and multi-processing.
• A standard PyTorch training loop consists of: zeroing gradients, forward pass, calculating loss, backward pass, and optimizer step.