pytorch

# 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 ```bash # CPU pip install torch torchvision # GPU (Check pytorch.org for specific CUDA versions) pip install torch --index-url https://download.pytorch.org/whl/cu121 ``` ### Standard Imports ```python import torch import torch.nn as nn import torch.optim as optim import torch.nn.functional as F from torch.utils.data import DataLoader, Dataset ``` ### Basic Pattern - Simple Linear Regression (The "PyTorch Way") ```python 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 weights ``` ## Critical 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 device` is 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) ```python 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 ```python 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 ```python 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) ```python 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) ```python 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 ```python 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_loss ``` ### 2. Early Stopping for Scientific Training ```python 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") break ``` ### 3. Feature Extraction for Chemistry ```python 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: ```python model = MyModel() compiled_model = torch.compile(model) ``` ### Mixed Precision (torch.cuda.amp) Saves memory and speeds up training on modern GPUs (Tensor Cores). ```python 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 ```python # ❌ 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. ```python # ✅ 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. ```python # ❌ 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.