PyTorch - Deep Learning & Tensors
PyTorch is a Python-based scientific computing package that uses the power of Graphics Processing Units (GPUs) and provides maximum flexibility and speed through its dynamic computational graph system.
When to Use
- Building and training Deep Neural Networks (CNN, RNN, Transformers).
- Researching new AI architectures with dynamic graph needs.
- Accelerating tensor math on NVIDIA (CUDA) or Mac (MPS) hardware.
- Solving Physics-Informed Neural Networks (PINNs).
- Implementing Generative models (GANs, Diffusion).
- Large-scale optimization using Autograd (automatic differentiation).
- Production-grade AI deployment (via TorchScript/ONNX).
Reference Documentation
Official docs: https://pytorch.org/docs/ Tutorials: https://pytorch.org/tutorials/ Search patterns: torch.nn, torch.optim, torch.utils.data, Autograd, Tensor.to(device)
Core Principles
The Tensor
The central data structure, similar to NumPy's ndarray, but with two key additions: it can live on a GPU and it supports automatic differentiation.
Dynamic Computational Graph (Autograd)
PyTorch builds the graph "on the fly" as code executes. This allows for standard Python control flow (if/for) inside your models.
Modules and Parameters
nn.Module is the base class for all neural network components. It automatically tracks nn.Parameter objects (weights/biases) for optimization.
Quick Reference
Installation
# CPU
pip install torch torchvision
# GPU (Check pytorch.org for specific CUDA versions)
pip install torch --index-url https://download.pytorch.org/whl/cu121Standard Imports
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import DataLoader, DatasetBasic Pattern - Simple Linear Regression (The "PyTorch Way")
import torch
# 1. Data (Tensors)
X = torch.tensor([[1.0], [2.0], [3.0]], requires_grad=True)
y = torch.tensor([[2.0], [4.0], [6.0]])
# 2. Simple Model
model = torch.nn.Linear(1, 1) # y = w*x + b
# 3. Loss and Optimizer
criterion = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
# 4. Training Loop
for epoch in range(100):
prediction = model(X)
loss = criterion(prediction, y)
optimizer.zero_grad() # Clear previous gradients
loss.backward() # Compute gradients (Autograd)
optimizer.step() # Update weightsCritical Rules
✅ DO
- Always zero_grad() - PyTorch accumulates gradients by default. Forget this, and your model will fail to converge.
- Use.to(device) - Explicitly move your model AND your data to the same device (CPU or CUDA).
- Use DataLoader - Never feed data manually in a loop; DataLoader handles batching, shuffling, and multi-process loading.
- Set model.train() / model.eval() - This is vital for layers like Dropout and BatchNorm that behave differently during inference.
- Use torch.no_grad() for inference - This saves significant memory and compute by not building a graph.
- Specify dtypes - Be conscious of float32 (standard) vs float64 (scientific precision) vs float16 (speed/GPU).
❌ DON'T
- Mix CPU and GPU Tensors -
RuntimeError: Expected all tensors to be on the same deviceis the most common error. - Use standard Python loops for math - Use vectorized tensor operations for performance.
- Forget.item() - When getting a scalar value from a tensor for logging, use
loss.item()to detach it from the graph. - Overuse float64 on GPU - Many consumer GPUs have poor double-precision performance; use float32 if possible.
Anti-Patterns (NEVER)
import torch
# ❌ BAD: Mixing Python lists/arrays with Tensors in a loop
# for x in data:
# res = model(torch.tensor(x)) # Extremely slow re-allocation!
# ✅ GOOD: Batching
# data_tensor = torch.stack([torch.tensor(x) for x in data])
# res = model(data_tensor)
# ❌ BAD: Calculating loss without zeroing gradients
loss.backward()
optimizer.step()
# Next iteration... gradients will be double what they should be!
# ✅ GOOD:
optimizer.zero_grad()
loss.backward()
optimizer.step()
# ❌ BAD: Standard NumPy for inference
# with torch.no_grad():
# pred = model(X).numpy() # Can be slow on GPU if not handled
# ✅ GOOD: Explicit move to CPU
# pred = model(X).detach().cpu().numpy()Tensors and Device Management
Moving between CPU and GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
# Create tensor on device
x = torch.randn(3, 3, device=device)
# Move model to device
model = MyModel().to(device)
# Move data to device during loop
for inputs, labels in dataloader:
inputs, labels = inputs.to(device), labels.to(device)
# ...Building Models (nn.Module)
Flexible Architectures
class ScientificNet(nn.Module):
def __init__(self, input_dim, hidden_dim):
super().__init__()
self.layer1 = nn.Linear(input_dim, hidden_dim)
self.layer2 = nn.Linear(hidden_dim, 1)
self.dropout = nn.Dropout(0.2)
def forward(self, x):
x = F.relu(self.layer1(x))
x = self.dropout(x)
x = torch.sigmoid(self.layer2(x))
return x
model = ScientificNet(10, 50)Custom Datasets (torch.utils.data)
Handling Scientific Files (e.g., HDF5 or CSV)
class MyScientificDataset(Dataset):
def __init__(self, file_path):
self.data = pd.read_csv(file_path)
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
# Convert row to tensor
sample = torch.tensor(self.data.iloc[idx, :-1].values, dtype=torch.float32)
label = torch.tensor(self.data.iloc[idx, -1], dtype=torch.float32)
return sample, label
dataset = MyScientificDataset("experiment_results.csv")
loader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4)Advanced Autograd
Gradients for Physics (Jacobians/Hessians)
x = torch.linspace(-5, 5, 100, requires_grad=True)
y = x**3
# First derivative dy/dx
# create_graph=True allows for higher-order derivatives
dy_dx = torch.autograd.grad(y.sum(), x, create_graph=True)[0]
# Second derivative (Hessian) d2y/dx2
d2y_dx2 = torch.autograd.grad(dy_dx.sum(), x)[0]Practical Workflows
1. Physics-Informed Neural Network (PINN) Fragment
def pde_loss(model, x):
"""Simple ODE: u'(x) = u(x)."""
x.requires_grad = True
u = model(x)
u_x = torch.autograd.grad(u.sum(), x, create_graph=True)[0]
return F.mse_loss(u_x, u)
# Training loop combines data_loss + pde_loss2. Early Stopping for Scientific Training
best_loss = float('inf')
patience = 10
counter = 0
for epoch in range(1000):
train_loss = train_one_epoch()
val_loss = validate()
if val_loss < best_loss:
best_loss = val_loss
torch.save(model.state_dict(), 'best_model.pth')
counter = 0
else:
counter += 1
if counter >= patience:
print("Early stopping triggered")
break3. Feature Extraction for Chemistry
def extract_embeddings(model, loader):
model.eval()
embeddings = []
with torch.no_grad():
for batch in loader:
# Assume model has a .get_features() method
features = model.get_features(batch.to(device))
embeddings.append(features.cpu())
return torch.cat(embeddings)Performance Optimization
Using torch.compile (PyTorch 2.0+)
Significant speedups for modern models with one line:
model = MyModel()
compiled_model = torch.compile(model)Mixed Precision (torch.cuda.amp)
Saves memory and speeds up training on modern GPUs (Tensor Cores).
scaler = torch.cuda.amp.GradScaler()
for inputs, labels in loader:
with torch.cuda.amp.autocast():
output = model(inputs)
loss = criterion(output, labels)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()Common Pitfalls and Solutions
Vanishing/Exploding Gradients
# ❌ Problem: Loss becomes 'nan' or weights don't update
# ✅ Solution: Use Gradient Clipping
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)Dead ReLU
If a neuron's output is always <= 0, it stops learning.
# ✅ Solution: Use LeakyReLU or ELU
self.act = nn.LeakyReLU(0.01)Memory Leak (Tensors staying in Graph)
Logging loss directly keeps the whole graph in memory.
# ❌ BAD: total_loss += loss
# ✅ GOOD: total_loss += loss.item()PyTorch is the engine of the AI revolution. For scientists, it provides the bridge from classical data analysis to the world of differentiable models, allowing for the discovery of patterns that were previously invisible.