PyTorch Fundamentals: Tensors, Autograd & Building a Neural Network
In the previous lesson, we derived Backpropagation and Gradient Descent mathematically and built an MLP from scratch using pure NumPy. In real-world deep learning, managing millions of matrix operations and manual analytical gradients across deep architectures is computationally intractable. In this lesson, we master PyTorch—the premier industrial and research framework—exploring n-dimensional tensors, hardware acceleration on CUDA GPUs, reverse-mode automatic differentiation (Autograd), modular network abstraction (nn.Module), data pipelines (DataLoader), and an end-to-end production training pipeline.
What You Will Master
- PyTorch Architecture: Imperative "Define-by-Run" philosophy vs static computational graphs.
- Tensors & Dtypes: Rank 0D to 5D representations, memory striding, FP32, FP16, and BF16.
- Hardware Acceleration: Device-agnostic code paradigms across CPU, CUDA, and Apple MPS.
-
Autograd Engine: Dynamic DAG construction, gradient accumulation,
backward(), andno_grad(). -
nn.Module & Layers: Fully connected
nn.Linear, activation functions, and modular composition. -
Losses & Optimizers:
BCEWithLogitsLossnumerical stability, Adam, and AdamW. - Dataset & DataLoader: Mini-batch streaming, worker multiprocessing, and memory pinning.
- End-to-End Pipeline: Complete training, validation, testing, saving, and loading on non-linear datasets.
1. What is PyTorch and Why Do We Need It?
PyTorch is an open-source deep learning framework created by the Facebook AI Research (FAIR) team in 2016, based on the original Lua-based Torch library. Unlike earlier first-generation deep learning frameworks (such as TensorFlow 1.x and Theano) that operated on rigid, static "compile-then-run" computational graphs, PyTorch pioneered the imperative, "Define-by-Run" paradigm. In PyTorch, every line of neural network code executes immediately in Python, just like standard NumPy code.
In Lesson 24, we implemented forward propagation and backpropagation using pure NumPy. While mathematically educational, NumPy lacks:
- No built-in Automatic Differentiation (gradients must be derived by hand).
- No native GPU/TPU acceleration (computations are locked to single-threaded or multi-core CPU).
- No high-level abstractions for layers, loss functions, or optimizers.
PyTorch resolves every fundamental bottleneck of scientific computing:
- Autograd: Dynamically traces tape executions to calculate exact reverse-mode derivatives.
- Massive Parallelism: Seamlessly transfers tensor buffers to NVIDIA CUDA and Apple MPS hardware.
- Industrial Ecosystem:
torch.nn,torch.optim, TorchScript, and ONNX deployment.
1.1 Deep Architectural Comparison: NumPy vs PyTorch
| Feature Dimension | NumPy (numpy) |
PyTorch (torch) |
|---|---|---|
| Core Purpose | General Scientific & Numerical Computing | Deep Learning & Accelerated Tensor Mathematics |
| Primary Data Structure | np.ndarray |
torch.Tensor |
| Automatic Differentiation | None (Manual symbolic or finite difference) | Autograd Engine (Dynamic Directed Acyclic Graphs) |
| Hardware Acceleration | CPU only (OpenBLAS, MKL) | NVIDIA CUDA, AMD ROCm, Apple Metal Performance Shaders (MPS) |
| Neural Layer Primitives | None (Write weights and matrix math manually) | torch.nn (Linear, Conv2d, Transformer, MultiheadAttention) |
| Optimization Algorithms | None (Write parameter updates manually) | torch.optim (SGD, Adam, AdamW, RMSprop, LBFGS) |
| Zero-Copy Memory Bridge |
Seamless shared-memory conversion: torch.from_numpy(arr) and tensor.numpy() point to identical CPU memory addresses without copying!
|
|
2. Installation, Environment Setup & Device Verification
PyTorch can be installed in an isolated Python virtual environment. Because deep learning models heavily benefit from GPU acceleration, selecting the correct wheel matching your system's hardware driver is essential.
Always verify your installation and device capabilities with the standard environment diagnostic snippet:
# Check library version
print(f"PyTorch Version: {torch.__version__}")
# Verify CUDA and MPS availability
cuda_avail = torch.cuda.is_available()
mps_avail = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
print(f"CUDA Available: {cuda_avail}")
if cuda_avail:
print(f"GPU Device Name: {torch.cuda.get_device_name(0)}")
print(f"Device Compute Capability: {torch.cuda.get_device_capability(0)}")
elif mps_avail:
print("Apple Metal Performance Shaders (MPS) Acceleration Enabled")
else:
print("Running on standard Host CPU")
torch.manual_seed(42); torch.cuda.manual_seed_all(42); np.random.seed(42)
3. What is a Tensor? Dimensionality & Real-World Hierarchy
In mathematics and computer science, a Tensor is a geometric object that maps linear relationships between vectors, scalars, and other tensors. In PyTorch, a torch.Tensor is a multi-dimensional matrix containing elements of a single data type, with a unified strided memory layout.
A single numerical value with zero dimensions.
A 1-dimensional array of numbers.
A 2-dimensional table of rows and columns.
A 3-dimensional cube of stacked matrices.
A batch of multi-channel 2D spatial grids.
A temporal sequence of multi-channel spatial volumes.
4. Tensor Creation, Factory Methods & Dtypes
PyTorch provides dozens of dedicated factory functions to allocate and initialize tensors with specific distributions, geometries, and precision.
# 1. Direct creation from Python data structures
t_data = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32)
# 2. Structural allocation factories
zeros_t = torch.zeros((2, 3)) # All zeros
ones_t = torch.ones((2, 3)) # All ones
empty_t = torch.empty((2, 3)) # Uninitialized raw memory
full_t = torch.full((2, 3), 3.1415) # Filled with constant
# 3. Random distribution factories
rand_uniform = torch.rand((2, 3)) # Uniform distribution in [0.0, 1.0)
rand_normal = torch.randn((2, 3)) # Standard Gaussian N(0, 1)
rand_int = torch.randint(low=0, high=10, size=(2, 3))
# 4. Range and progression factories
step_range = torch.arange(start=0, end=10, step=2) # tensor([0, 2, 4, 6, 8])
linspace_t = torch.linspace(start=0, end=1, steps=5) # 5 equally spaced values
# 5. Like-factories (matching shape and dtype of existing tensor)
zeros_like_t = torch.zeros_like(rand_normal)
4.1 Data Types (Dtypes) & Numerical Precision Engineering
Selecting the appropriate data type directly affects VRAM memory consumption and hardware throughput. In modern LLM training and vision inference, precision tuning is paramount.
| Data Type | Bit Width & Format | Primary Deep Learning Use Case |
|---|---|---|
| torch.float32 | 32-bit (1 sign, 8 exponent, 23 mantissa) | Standard default precision for weights, activations, and gradient accumulators. |
| torch.bfloat16 | 16-bit Brain Float (1 sign, 8 exponent, 7 mantissa) | Modern LLM pre-training & Transformers; preserves FP32 dynamic range without underflow. |
| torch.float16 | 16-bit Half Precision (1 sign, 5 exponent, 10 mantissa) | FP16 Automatic Mixed Precision (AMP) training on NVIDIA Tensor Cores. |
| torch.float64 | 64-bit Double Precision | High-precision scientific computing and gradient checking validation. |
| torch.int64 | 64-bit Signed Integer (Long) | Categorical class labels in CrossEntropyLoss and vocabulary token indices. |
| torch.bool | 1-byte Boolean | Attention padding masks and boolean tensor masking. |
5. Tensor Operations, Broadcasting & Linear Algebra
PyTorch supports over 200 mathematical operators. Crucially, operations follow NumPy broadcasting semantics, expanding singleton dimensions automatically during arithmetic without allocating redundant memory.
x.add_(y), x.zero_()) modify underlying memory in place. Caution: Modifying a leaf variable required for backward gradient computation causes an immediate Autograd crash.
5.1 Reshaping, Slicing & Memory Layout: view() vs reshape()
In neural network pipelines, altering tensor geometry without modifying data is continuous. PyTorch differentiates between zero-copy viewing and memory reallocation:
# 1. view(): Requires contiguous memory, creates zero-copy shared view
y_view = x.view(3, 4)
# 2. reshape(): Returns view if memory is contiguous; otherwise copies memory
y_reshape = x.reshape(3, 4)
# 3. unsqueeze(): Adds a singleton dimension of size 1 at specified axis
x_1d = torch.tensor([1.0, 2.0, 3.0]) # shape: torch.Size([3])
x_row = x_1d.unsqueeze(0) # shape: torch.Size([1, 3]) (Batch of 1)
x_col = x_1d.unsqueeze(1) # shape: torch.Size([3, 1]) (Column vector)
# 4. squeeze(): Strips all singleton dimensions (or specified dimension)
clean_1d = x_row.squeeze() # shape: torch.Size([3])
# 5. Permute: Arbitrary dimension reordering (e.g., NCHW -> NHWC for plotting)
img_nchw = torch.randn(16, 3, 64, 64)
img_nhwc = img_nchw.permute(0, 2, 3, 1) # shape: torch.Size([16, 64, 64, 3])
# 6. Flattening: Collapses spatial dimensions for linear dense layers
flattened = torch.flatten(img_nchw, start_dim=1) # shape: torch.Size([16, 12288])
6. CPU & GPU Hardware Acceleration: Device-Agnostic Programming
Deep neural network training relies fundamentally on dense matrix multiplications that run up to $50\times$ faster on massively parallel GPU architectures than on standard CPUs. PyTorch allows seamless movement of tensors and model parameters across memory boundaries using the .to(device) abstraction.
"cuda" if torch.cuda.is_available()
else ("mps" if hasattr(torch.backends, "mps") and torch.backends.mps.is_available() else "cpu")
)
print(f"Active Compute Target: {device}")
In PyTorch, a tensor residing in CPU RAM cannot perform arithmetic or matrix multiplication with a tensor located in GPU VRAM. Both the model parameters and the input features and targets must reside on the exact same physical device. Failing to ensure this results in the fatal runtime exception:
Remedy: Always migrate your models and mini-batch tensors explicitly inside your training loop:
model = model.to(device); features = features.to(device); labels = labels.to(device)
7. Autograd: Automatic Differentiation Engine
PyTorch's defining capability is Autograd. When you perform operations on tensors flagged with requires_grad=True, Autograd dynamically constructs a Directed Acyclic Graph (DAG) in real time. Each intermediate tensor stores a grad_fn pointer referencing the backward mathematical operator required to compute its local derivative.
Consider the scalar polynomial expression:
By elementary calculus, the exact analytical derivative with respect to $x$ is:
Evaluating this derivative at $x = 2.0$:
Now let us observe how PyTorch computes this identical quantity automatically without any manual calculus formulas:
# 1. Instantiate input leaf tensor with gradient tracking
x = torch.tensor(2.0, requires_grad=True)
# 2. Forward pass: evaluate algebraic computation graph
y = x ** 2 + 3 * x + 1
print(f"Forward Output y: {y.item()}") # 2^2 + 3(2) + 1 = 11.0
print(f"Gradient Function: {y.grad_fn}") # <AddBackward0 object>
# 3. Backward pass: initiate reverse-mode Vector-Jacobian Product
y.backward()
# 4. Inspect accumulated gradient
print(f"Computed Gradient dy/dx: {x.grad}") # tensor(7.)
7.1 The Gradient Accumulation Trait & zero_grad()
A paramount operational characteristic of PyTorch's Autograd is that gradients accumulate by default. Whenever .backward() is called, newly computed partial derivatives are added to the existing values in .grad via x.grad += dL/dx.
If you call loss.backward() across multiple training iterations without clearing gradients, your parameters will update using the sum of gradients from all prior batches! This results in exploding gradients and severe training divergence.
Always clear gradients before each backward pass: optimizer.zero_grad() (or in modern PyTorch: optimizer.zero_grad(set_to_none=True) to save memory bandwidth).
7.2 Disabling Gradient Tracking: torch.no_grad() & Inference Mode
During model evaluation, validation, or live inference, tracking computational history consumes unnecessary memory and compute. Wrapping evaluation code inside with torch.no_grad(): disables graph construction, speeding up execution and freeing VRAM:
model.eval()
with torch.no_grad():
predictions = model(test_features)
probabilities = torch.sigmoid(predictions)
# In modern PyTorch 1.9+, torch.inference_mode() is even faster:
with torch.inference_mode():
prod_preds = model(live_stream_features)
8. torch.nn.Module: The Neural Network Architecture Blueprint
In PyTorch, all neural network architectures inherit from the base class torch.nn.Module. Modules can contain other modules (layers, sub-networks), automatically tracking parameters, buffers, and hooks.
Every custom neural network module requires two core methods:
-
__init__(self, ...): Initializes superclass viasuper().__init__(), defines layers (nn.Linear,nn.Dropout, etc.), and initializes learnable weights. -
forward(self, x): Defines the forward computation flow from input tensor $x$ to output tensor $\hat{y}$. Never invokemodel.forward(x)directly; always callmodel(x)so PyTorch can execute registered hooks.
class MultiLayerPerceptron(nn.Module):
def __init__(self, in_features=2, hidden_dim=32, out_features=1):
super().__init__()
self.net = nn.Sequential(
nn.Linear(in_features, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim // 2),
nn.ReLU(),
nn.Linear(hidden_dim // 2, out_features)
)
def forward(self, x):
return self.net(x)
model = MultiLayerPerceptron()
print(model)
8.1 Inside nn.Linear: Weights, Biases & Geometry
The fundamental building block of dense neural networks is nn.Linear(in_features, out_features). It executes an affine transformation:
Notice that PyTorch stores the weight matrix in transposed orientation: shape is (out_features, in_features), while the bias vector has shape (out_features,). Both parameters are automatically initialized using Kaiming Uniform initialization to prevent symmetry collapse.
9. Activation Functions in PyTorch
Without non-linear activation functions, stacking multiple linear layers collapses mathematically into a single linear projection: $W_2(W_1 x + b_1) + b_2 = W_{\text{eff}} x + b_{\text{eff}}$. Non-linearities allow the network to warp input feature space and learn arbitrary non-linear decision boundaries.
| Activation Function | Mathematical Definition | PyTorch Module / Functional | Primary Application |
|---|---|---|---|
| ReLU | $\max(0, z)$ | nn.ReLU() / F.relu(x) | Default standard for hidden layers in MLPs and CNNs. |
| Leaky ReLU | $\max(\alpha z, z), \alpha=0.01$ | nn.LeakyReLU(0.01) | Prevents the "Dying ReLU" neuron extinction problem. |
| GELU | $z \cdot \Phi(z)$ | nn.GELU() / F.gelu(x) | Industry standard for modern Transformers (BERT, GPT, LLaMA). |
| Sigmoid | $\frac{1}{1 + e^{-z}}$ | nn.Sigmoid() | Binary classification output probability mapping $[0, 1]$. |
| Tanh | $\frac{e^z - e^{-z}}{e^z + e^{-z}}$ | nn.Tanh() | Zero-centered hidden layers, RNNs, and LSTMs $[-1, 1]$. |
| Softmax | $\frac{e^{z_i}}{\sum_j e^{z_j}}$ | nn.Softmax(dim=-1) | Multiclass probability distribution summing to $1.0$. |
10. Loss Functions & Numerical Stability
Loss functions measure the empirical error between model predictions and ground truth targets. Selecting the mathematically harmonious loss function is crucial for training stability.
Computes Mean Squared Error: $\frac{1}{n} \sum (y_i - \hat{y}_i)^2$. Penalizes large outliers heavily. Alternative: nn.SmoothL1Loss() (Huber) for outlier robustness.
Combines Sigmoid activation and Binary Cross Entropy into a single numerically stable Log-Sum-Exp operator. Never pass Sigmoid outputs here; pass unnormalized raw logits!
Combines LogSoftmax and NLLLoss. Accepts raw logits of shape $(B, C)$ and ground truth integer class targets of shape $(B,)$ with dtype torch.int64. No one-hot encoding required!
11. Optimizers & The Sacred 3-Step Parameter Update Sequence
While Autograd computes parameter gradients ($\nabla_{\mathbf{w}} L$), the Optimizer determines how those gradients are applied to adjust model weights. PyTorch houses standard optimizers in torch.optim.
Every neural network training loop in PyTorch executes this exact sequence in every mini-batch:
optimizer.zero_grad()
Purges accumulated gradients from the previous mini-batch before backward propagation starts.
loss.backward()
Autograd traces reverse graph from scalar loss, populating .grad for all trainable parameters.
optimizer.step()
Applies optimizer update rule (Adam, SGD) using current gradients: $w \leftarrow w - \eta \cdot \Delta w$.
12. Data Ingestion Architecture: Dataset & DataLoader
In enterprise machine learning, training sets exceed available VRAM and system memory. PyTorch cleanly decouples dataset storage from batch sampling mechanics via two fundamental classes:
Represents the data collection. Custom datasets implement two methods:
__len__(self): Returns total dataset size $N$.__getitem__(self, idx): Loads and returns a single sample $(x_i, y_i)$.
For tensors already loaded in memory: TensorDataset(X, y) provides instantaneous indexing.
An asynchronous iterator wrapping the Dataset to manage mini-batching, shuffling, multi-process loading, and memory pinning:
batch_size=32: Group 32 samples into a single tensor.shuffle=True: Shuffles indices before every epoch.pin_memory=True: Pinned RAM for fast host-to-GPU transfer.
13. End-to-End Production Pipeline: Non-Linear Binary Classification (make_moons)
Let us synthesize everything into a complete, leak-free PyTorch deep learning pipeline. We tackle the canonical non-linear interlocking moons dataset (make_moons), performing rigorous train/validation/test splitting, feature normalization, mini-batch training, and test set evaluation.
import torch
import torch.nn as nn
from sklearn.datasets import make_moons
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import accuracy_score, f1_score
from torch.utils.data import TensorDataset, DataLoader
# 1. Device configuration
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 2. Synthetic dataset generation & stratified splitting (70% Train, 15% Val, 15% Test)
X_raw, y_raw = make_moons(n_samples=2000, noise=0.20, random_state=42)
X_train, X_temp, y_train, y_temp = train_test_split(X_raw, y_raw, test_size=0.30, random_state=42, stratify=y_raw)
X_val, X_test, y_val, y_test = train_test_split(X_temp, y_temp, test_size=0.50, random_state=42, stratify=y_temp)
# 3. Leak-free feature scaling (Fit ONLY on training set!)
scaler = StandardScaler()
X_train = scaler.fit_transform(X_train)
X_val = scaler.transform(X_val)
X_test = scaler.transform(X_test)
# 4. Conversion to PyTorch Tensors & DataLoaders
train_loader = DataLoader(TensorDataset(torch.tensor(X_train, dtype=torch.float32), torch.tensor(y_train[:, None], dtype=torch.float32)), batch_size=32, shuffle=True)
val_X_t = torch.tensor(X_val, dtype=torch.float32).to(device)
val_y_t = torch.tensor(y_val[:, None], dtype=torch.float32).to(device)
test_X_t = torch.tensor(X_test, dtype=torch.float32).to(device)
test_y_t = torch.tensor(y_test[:, None], dtype=torch.float32).to(device)
# 5. Architecture Definition: Multi-Layer Perceptron
class BinaryClassifier(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Linear(2, 32),
nn.ReLU(),
nn.Linear(32, 16),
nn.ReLU(),
nn.Linear(16, 1) # Raw logit output
)
def forward(self, x):
return self.net(x)
model = BinaryClassifier().to(device)
criterion = nn.BCEWithLogitsLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.005)
# 6. Rigorous Training Loop
epochs = 80
for epoch in range(epochs):
model.train()
running_loss = 0.0
for batch_x, batch_y in train_loader:
batch_x, batch_y = batch_x.to(device), batch_y.to(device)
logits = model(batch_x)
loss = criterion(logits, batch_y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
running_loss += loss.item() * batch_x.size(0)
train_loss = running_loss / len(train_loader.dataset)
# Validation step
model.eval()
with torch.no_grad():
val_loss = criterion(model(val_X_t), val_y_t).item()
if (epoch + 1) % 20 == 0:
print(f"Epoch {epoch+1:03d} | Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f}")
# 7. Final Test Evaluation
model.eval()
with torch.no_grad():
test_logits = model(test_X_t)
test_preds = (torch.sigmoid(test_logits) >= 0.5).int().cpu().numpy()
acc = accuracy_score(y_test, test_preds)
f1 = f1_score(y_test, test_preds)
print(f"Final Test Accuracy: {acc * 100:.2f}% | F1-Score: {f1:.4f}")
14. Model Persistence: Saving & Loading Checkpoints
In production systems, trained models must be serialized to disk for containerized deployment, microservice hosting, or resumption of interrupted training runs. PyTorch recommends saving the state_dict (an ordered Python dictionary mapping each layer name to its parameter tensor).
Serializes weights safely using Python's pickle protocol:
For complete resume checkpoints, save a dictionary containing epoch, model_state_dict, optimizer_state_dict, and best_loss.
Instantiate architecture first, load state dict, and set to eval:
15. Top 10 PyTorch Pitfalls & Anti-Patterns Matrix
Deep learning bugs in PyTorch rarely throw explicit compile errors; instead, they silently corrupt mathematical convergence. Study this field-tested anti-pattern diagnostic matrix:
| # | Common Pitfall | Silent Failure Symptom | Production Correction |
|---|---|---|---|
| 1 | Omitting zero_grad() |
Gradients sum endlessly across batches, exploding training loss. | optimizer.zero_grad() before loss.backward() |
| 2 | Double Sigmoid with BCE | Applying Sigmoid in model AND using BCEWithLogitsLoss squashes probabilities twice. |
Output raw unbounded logits from model; let loss compute Sigmoid internally. |
| 3 | Softmax before CrossEntropy | CrossEntropyLoss applies LogSoftmax internally; double Softmax corrupts gradient magnitudes. |
Pass raw logits directly to nn.CrossEntropyLoss(). |
| 4 | Device Mismatch | Runtime crash: Expected all tensors to be on same device. | Ensure model.to(device) and x.to(device), y.to(device) match. |
| 5 | Missing model.eval() |
Dropout continues dropping units and BatchNorm normalizes by mini-batch during inference. | Always invoke model.eval() before validation and inference. |
| 6 | Missing torch.no_grad() |
Validation passes construct unneeded autograd graphs, causing GPU Out-Of-Memory (OOM). | Wrap all validation and testing in with torch.no_grad():. |
| 7 | Accumulating Tensor History | Writing total_loss += loss accumulates entire graph history, leaking VRAM. |
total_loss += loss.item() * batch_size |
| 8 | Target Dtype Mismatch | Passing float targets to CrossEntropyLoss raises runtime type error. |
Cast target labels to torch.int64 (.long()). |
| 9 | Data Leakage in Normalization | Fitting Scaler on whole dataset leaks test distribution into training. | Call scaler.fit_transform(X_train), then scaler.transform(X_test). |
| 10 | Mutating Leaf Variables In-Place | Directly editing param.data += ... breaks autograd version counters. |
Wrap manual parameter updates inside with torch.no_grad():. |
Practical Assignment: Production PyTorch Classifier
Construct an end-to-end PyTorch script meeting these production specifications:
make_moons(n_samples=3000, noise=0.25, random_state=42). Split into 70% Train, 15% Validation, and 15% Test. Standardize with leak-free scaling.
nn.Module with 3 dense layers: 2 → 64 (ReLU) → 32 (ReLU) → 1 (Linear Logit).
TensorDataset with batch_size=64 and shuffle=True.
BCEWithLogitsLoss and AdamW(lr=0.003, weight_decay=1e-4). Log validation loss every 10 epochs.
state_dict() to disk, reload into a clean instance, and confirm test accuracy $\ge 95\%$.
PyTorch Interactive Exploration Workbench
Experiment hands-on with tensor geometries, dynamic autograd graphs, non-linear decision boundaries, and mini-batch iterators.
Tensor Rank, Memory & Geometry Inspector
Select rank and dimensions to visualize multi-dimensional layout and memory footprints.