Introduction to PyTorch: Tensors, Autograd, and Custom Models
Learn the basics of PyTorch. Master PyTorch Tensors, dynamic computation graphs, Autograd, and building custom neural networks in Python.
Introduction
In the early days of Deep Learning, researchers had to write hundreds of lines of C++ code to implement a single neural network layer and manually calculate massive derivative matrices on paper.
In 2016, Meta AI released PyTorch. Today, it is the most popular deep learning framework in research and industry, powering platforms like Tesla Autopilot, Hugging Face, and OpenAI. PyTorch provides a GPU-accelerated library for multi-dimensional arrays (Tensors) and an automatic differentiation engine (Autograd) that calculates all backpropagation gradients for you automatically.
What You Will Learn
- What a PyTorch Tensor is and how it compares to NumPy.
- How Autograd performs automatic differentiation.
- The structure of PyTorch models (
nn.Module). - How to write a complete training loop in PyTorch.
Why This Topic Matters
PyTorch is the native language of modern deep learning. If you want to customize transformer layers, run inference on pre-trained models from Hugging Face, or deploy computer vision pipelines, you must understand how PyTorch manages computation graphs and updates weights.
Prerequisites
Detailed Explanation
PyTorch has three core pillars: Tensors, Autograd, and the Neural Network Module (nn.Module).
1. PyTorch Tensors
A Tensor is a multi-dimensional matrix similar to NumPy's ndarray. However, unlike NumPy arrays, PyTorch Tensors can run on GPUs to accelerate mathematical calculations by 50x or more.
import torch
# Create a tensor from a list
x = torch.tensor([1.0, 2.0, 3.0])
# Send tensor to GPU if available
device = "cuda" if torch.cuda.is_available() else "cpu"
x = x.to(device)
2. Autograd (Automatic Differentiation)
Autograd is PyTorch's automatic differentiation engine. If you set requires_grad=True on a Tensor, PyTorch starts tracking all mathematical operations applied to it.
It builds a dynamic Directed Acyclic Graph (DAG) of operations. When you call .backward() on the final output (typically the loss), PyTorch calculates all the gradients and stores them in the .grad attribute of each tensor.
graph LR
X[Tensor x requires_grad=True] -->|Multiply by 2| Y[Tensor y = 2*x]
Y -->|Square| Z[Loss L = y^2]
Z -->|Call L.backward| X
X -.->|Updates| Xgrad[x.grad = dL/dx]
3. Building Models with torch.nn
In PyTorch, custom neural networks inherit from torch.nn.Module. You define your layers inside __init__(), and specify the forward pass logic in forward():
import torch.nn as nn
class SimpleClassifier(nn.Module):
def __init__(self, input_dim, output_dim):
super(SimpleClassifier, self).__init__()
# Define layers
self.linear1 = nn.Linear(input_dim, 64)
self.relu = nn.ReLU()
self.linear2 = nn.Linear(64, output_dim)
def forward(self, x):
# Define forward flow
out = self.linear1(x)
out = self.relu(out)
out = self.linear2(out)
return out
Python Code Examples
Here is a complete, production-ready PyTorch training loop that trains a simple linear model to fit a line ($y = 2x + 1$).
import torch
import torch.nn as nn
import torch.optim as optim
# 1. Prepare synthetic dataset
# 100 samples, 1 feature
X_train = torch.randn(100, 1)
y_train = 2 * X_train + 1 + torch.randn(100, 1) * 0.1 # y = 2x + 1 + noise
# 2. Define Model, Loss, and Optimizer
model = nn.Linear(1, 1) # Single linear layer (weights & bias)
criterion = nn.MSELoss() # Mean Squared Error loss
optimizer = optim.SGD(model.parameters(), lr=0.01) # Stochastic Gradient Descent
# 3. Training Loop
epochs = 100
for epoch in range(epochs):
# Forward Pass
predictions = model(X_train)
loss = criterion(predictions, y_train)
# Backward Pass (Zero gradients, backpropagate, update)
optimizer.zero_grad() # Clears old gradients from previous step
loss.backward() # Computes dL/dw and dL/db
optimizer.step() # Updates weights: w = w - lr * grad
if (epoch + 1) % 20 == 0:
print(f"Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.4f}")
# 4. View learned parameters
print("\nLearned Weights:", model.weight.item())
print("Learned Bias:", model.bias.item())
Industry Use Cases
- Fine-Tuning LLMs: Loading open-source models using Hugging Face's
transformerslibrary, which runs on PyTorch, to adapt models to specific company datasets. - Computer Vision Pipelines: Using
torchvisionmodels (like ResNet or YOLO) to run real-time object detection on live video feeds.
Advantages & Limitations
Advantages
- Pythonic & Intuitive: Employs standard Python control flow, making debugging with
pdbeasy. - Dynamic Computation Graphs: Builds the graph on-the-fly, allowing you to change layer shapes or behaviors during execution.
- Strong Ecosystem: Massive support from Meta, AWS, Google, and the open-source community.
Limitations
- Deployment Overhead: Traditional PyTorch models can be heavy for low-latency web production. Typically, models must be compiled to TorchScript or exported to ONNX formats first.
- Steeper Learning Curve: Requires writing custom training loops (unlike Keras, which provides a high-level
model.fit()API, though PyTorch Lightning solves this).
FAQs
Q: Why do we call optimizer.zero_grad()?
A: By default, PyTorch accumulates (adds) gradients on subsequent .backward() calls instead of overwriting them. If you do not call zero_grad(), gradients from the previous epoch will add up to the current gradients, causing updates to explode.
Q: What is the difference between model.train() and model.eval()?
A: model.train() configures layers like Dropout and Batch Normalization to active learning modes. model.eval() disables them, ensuring consistent, deterministic outputs during model evaluation or inference.
Summary
PyTorch is a flexible, developer-friendly deep learning library. By using Tensors (which run on GPUs) and Autograd (which computes gradients automatically), PyTorch simplifies neural network development. Custom models are built by extending nn.Module, wrapping training states, and updating variables using built-in optimizers like SGD and Adam.
Next Topic
Now that we have mastered the foundations of Deep Learning, let's explore how neural networks process human language. Move on to: Natural Language Processing (NLP) Fundamentals.