pytorch

What this skill does

# 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 tor