PyTorch Basics — Taking the First Step to Implementing Principles with Code

After Completing This Topic

You will be able to implement the principles of neural networks (dial, gradient descent, backpropagation) learned in Episodes 2-4 into actual PyTorch code. This marks the beginning of the Phase 3 Tool section. This episode serves as the foundation for the practical exercises in Episodes 13 (Hugging Face/API) and 14 (Claude Code/Cursor).

The goal is not to go through short code snippets, but to construct a complete training loop from start to finish. We will use the cell type classifier (the scenario in Episode 3) as an example.

What is PyTorch?

PyTorch is a deep learning framework developed by Meta AI. It provides Python APIs for the necessary components for training and inferencing neural networks.

It has four core components:

• Tensor: A multi-dimensional array. Similar to NumPy, but with GPU acceleration and automatic differentiation support.
• Autograd: Automates backpropagation from Episode 4. It creates a computational graph in the background, and a single line of code backward() calculates all the partial derivatives.
• nn.Module: A class interface for defining neural network layers and models.
• Optimizer: Implements optimizers such as Adam and AdamW from Episode 3.

By assembling these components, we can create a training loop.

Installation: It is already installed in Colab. For local installations, use pip install torch. When using a GPU, specify the wheel that matches the CUDA version.

Alternatives: TensorFlow, JAX, MLX (Apple). PyTorch is the current standard for research and practical applications. Most open-source models (LLaMA, Mistral, Qwen, etc.) are distributed in PyTorch.

Tensor — The Container for the Dials and Pixels in Episode 2

A Tensor is a multi-dimensional array of real numbers. The layer activations, weights, and biases discussed in Episode 2 are all tensors.

import torch

# Scalar (0-dimensional)
x = torch.tensor(3.14)

# Vector (1-dimensional)
v = torch.tensor([1.0, 2.0, 3.0])

# Matrix (2-dimensional)
m = torch.tensor([[1.0, 2.0], [3.0, 4.0]])

# 3-dimensional (e.g., batch size 32, sequence length 100, embedding size 512)
t = torch.randn(32, 100, 512)
print(t.shape)  # torch.Size([32, 100, 512])

dtype and device: A tensor has a dtype (data type) and a device (CPU/GPU).

# fp32 CPU (default)
x = torch.randn(1000, 1000)

# fp16 GPU
y = torch.randn(1000, 1000, dtype=torch.float16, device="cuda")

# Moving the device
x_gpu = x.to("cuda")
x_cpu = y.to("cpu")

GPU computations are much faster. Most practical training is done on GPUs.

Tensor operations: Almost the same API as NumPy.

a = torch.randn(3, 4)
b = torch.randn(4, 5)

c = a @ b        # Matrix multiplication, shape (3, 5)
d = a + 1        # Broadcasting
e = a.sum(dim=1) # Sum along rows, shape (3,)
f = a.mean()     # Scalar
g = a.relu()     # Element-wise ReLU

The activation functions (ReLU, GELU, Sigmoid, etc.) discussed in Episodes 2 and 5 are all provided as tensor methods.

Autograd — Backpropagation in One Line from Episode 4

If you attach requires_grad=True to a tensor, PyTorch will track that tensor for partial derivative calculations. The computational graph from Episode 4 is automatically created.

# Parameter (target for partial differentiation)
w = torch.randn(3, requires_grad=True)
b = torch.randn(1, requires_grad=True)

# Data
x = torch.tensor([1.0, 2.0, 3.0])
y = torch.tensor(10.0)

# Forward pass
y_pred = w @ x + b
loss = (y_pred - y) ** 2

# Backpropagation - all partial derivatives calculated with this single line
loss.backward()

# Calculated gradients
print(w.grad)   # ∂loss/∂w
print(b.grad)   # ∂loss/∂b

All four steps of Episode 4 are included here.

• Forward pass (y_pred = w @ x + b, loss = ...): Automatic construction of the computational graph.
• Loss calculation ((y_pred - y) ** 2).
• Backpropagation (loss.backward()): Traverses the graph in reverse, calculating all partial derivatives.
• Parameter update: To be covered next.

Important convention: After calling .backward(), the gradient of the parameter is accumulated. You must initialize the gradient to 0 before the next step (using optimizer.zero_grad() later).

nn.Module — Defining Models from Episode 5

Large models are defined by inheriting from the nn.Module class. Let's actually implement the cell type classifier from Episode 3.

import torch.nn as nn

class CellTypeClassifier(nn.Module):
    def __init__(self, n_genes, n_types, hidden=512):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(n_genes, hidden),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(hidden, hidden // 2),
            nn.ReLU(),
            nn.Linear(hidden // 2, n_types)
        )

    def forward(self, x):
        return self.encoder(x)

model = CellTypeClassifier(n_genes=20000, n_types=25)
print(sum(p.numel() for p in model.parameters()))
# Prints the number of parameters

Explanation:

• __init__: Defines the layers. nn.Linear(in, out) is y = W·x + b. The fully connected layer from Episode 2.
• nn.Sequential: Connects the layers sequentially.
• nn.ReLU(), nn.Dropout(0.3): Activation function and regularization.
• forward: Forward pass logic. PyTorch automatically calls this function.

Automatic parameter tracking: The weights and biases within nn.Linear are automatically set to requires_grad=True and included in model.parameters(). There is no need to manage them manually using torch.tensor(..., requires_grad=True).

Calling: The model is called like a function.

x = torch.randn(32, 20000)  # Batch size 32, 20000 genes
logits = model(x)           # Automatically calls forward
print(logits.shape)         # torch.Size([32, 25])

Loss and Optimizer — Gradient Descent from Episode 3

The loss function and optimizer learned in Episode 3.

Loss:

# Classification problem - CE from Episode 3 A.1
loss_fn = nn.CrossEntropyLoss()

# Regression problem - MSE from Episode 3 A.1
loss_fn = nn.MSELoss()

Optimizer:

# SGD - from Episode 3 A.3
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

# Adam - from Episode 3 A.7
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# AdamW - from Episode 3 A.8 (standard for LLM training)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)

Learning rate scheduler:

scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)

Call scheduler.step() at the end of each epoch to automatically decrease the learning rate.

Training Loop — Putting It All Together

Now, let's combine all the pieces into a complete training loop.

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

# Prepare data (dummy example)
n_samples = 10000
n_genes = 20000
n_types = 25

X = torch.randn(n_samples, n_genes)
y = torch.randint(0, n_types, (n_samples,))

dataset = TensorDataset(X, y)
train_loader = DataLoader(dataset, batch_size=128, shuffle=True)

# Model
model = CellTypeClassifier(n_genes, n_types)
model = model.to("cuda")

# Loss and optimizer
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)

# Training loop
n_epochs = 10
for epoch in range(n_epochs):
    model.train()
    total_loss = 0.0
    for batch_x, batch_y in train_loader:
        batch_x = batch_x.to("cuda")
        batch_y = batch_y.to("cuda")
        
        # Forward pass (Part #2)
        logits = model(batch_x)
        loss = loss_fn(logits, batch_y)
        
        # Backward pass (Part #4)
        optimizer.zero_grad()  # Initialize gradients from the previous step
        loss.backward()        # Autograd
        
        # Parameter update (Part #3)
        optimizer.step()
        
        total_loss += loss.item()
    
    print(f"Epoch {epoch}: loss = {total_loss / len(train_loader):.4f}")

This loop corresponds exactly to the diagram in Part #3.

Initialize parameters randomly ← Automatically done when creating the model
Repeat (tens of thousands to millions of times):
    Randomly select a mini-batch ← DataLoader
    Perform a forward pass on the mini-batch ← logits = model(batch_x)
    Calculate the cost based on predictions and ground truth ← loss = loss_fn(logits, batch_y)
    Perform a backward pass → Calculate gradients ← loss.backward()
    Update parameters using the optimizer ← optimizer.step()

PyTorch exposes the principles of Parts #2 to #4 as APIs.

Dataset and DataLoader — Mini-Batches from Part #3

In Part #3, we discussed why mini-batch SGD is necessary. PyTorch supports this with Dataset + DataLoader.

Dataset: A class that defines __len__ and __getitem__.

from torch.utils.data import Dataset

class ScRNADataset(Dataset):
    def __init__(self, expression_matrix, labels):
        self.expr = expression_matrix
        self.labels = labels
    
    def __len__(self):
        return len(self.labels)
    
    def __getitem__(self, idx):
        return self.expr[idx], self.labels[idx]

DataLoader: Handles batching, shuffling, and parallel loading.

train_loader = DataLoader(
    dataset,
    batch_size=128,
    shuffle=True,        # Shuffle the order of data in each epoch
    num_workers=4,       # Parallel loading (utilizes CPU cores)
    pin_memory=True,     # Accelerate GPU transfer
    drop_last=True       # Drop the last incomplete batch
)

Bio Practice: Implement custom Datasets for scRNA-seq, images, sequence data, etc., specific to each domain. Scanpy, AnnData, and PyTorch Geometric provide domain-specific wrappers.

GPU Utilization — Fast Training

GPUs are tens to hundreds of times faster than CPUs. Most practical training is done on GPUs.

# Check for available GPUs
print(torch.cuda.is_available())
print(torch.cuda.device_count())

# Move tensors and models
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
batch_x = batch_x.to(device)

Multi-GPU: Use multiple GPUs for large models and datasets.

• DataParallel (DP): Multiple GPUs on a single node. Simple but can be a bottleneck.
• DistributedDataParallel (DDP): Standard. Each GPU is a separate process, with gradient all-reduce.
• Fully Sharded Data Parallel (FSDP): Parameters and optimizer states are also distributed. Essential for training 70B+ models.

Apple Silicon: Supports the mps device. device = "mps". Allows small-scale experiments on personal laptops.

Mixed Precision — Reduced Training Time and Memory Usage

Calculate using FP16 or BF16 instead of FP32 to reduce training time and memory usage.

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for batch_x, batch_y in train_loader:
    optimizer.zero_grad()
    
    with autocast(dtype=torch.bfloat16):
        logits = model(batch_x)
        loss = loss_fn(logits, batch_y)
    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

BF16 vs FP16:

• FP16: Half-precision. Requires GradScaler to avoid overflow.
• BF16 (bfloat16): Has the same exponent range as FP32. Lower risk of overflow. Supported by A100 and H100 series GPUs. Moving towards becoming the standard in practice.

Effect: Training time is 2-4 times faster, and memory usage is halved. Minimal loss of accuracy.

Saving and Loading — Checkpoints

Save the intermediate and final states of training.

# Save
torch.save({
    "epoch": epoch,
    "model_state_dict": model.state_dict(),
    "optimizer_state_dict": optimizer.state_dict(),
    "loss": total_loss,
}, "checkpoint.pt")

# Load
checkpoint = torch.load("checkpoint.pt")
model.load_state_dict(checkpoint["model_state_dict"])
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
start_epoch = checkpoint["epoch"] + 1

If training takes days, be sure to save checkpoints regularly as a precaution against crashes or power outages.

Inference — Using the Trained Model

Make predictions with the trained model.

model.eval()  # Set dropout and batchnorm to inference mode
with torch.no_grad():  # Disable gradient calculation (saves memory and speeds up computation)
    for batch_x, _ in test_loader:
        batch_x = batch_x.to(device)
        logits = model(batch_x)
        preds = logits.argmax(dim=-1)
        # ... use preds

no_grad prevents autograd from creating a computation graph, saving memory and increasing speed.

Deployment: There are several ways to deploy trained models to production.

• TorchScript: JIT-compiles PyTorch models for deployment to C++ and mobile.
• ONNX: Framework-independent format. Optimize with TensorRT and OpenVINO.
• HuggingFace Hub: Share and serve models.
• vLLM, TGI, SGLang: Frameworks specifically designed for serving LLMs.

Bio Practical Tips

Data Loading

scRNA-seq: scanpy + AnnData → PyTorch Dataset wrapper. Protein sequences: Parse with biopython and create a Dataset. Images (pathology/microscopy): Utilize torchvision.datasets.ImageFolder and apply augmentation.

Training Optimization

• Gradient accumulation: Simulate a large batch when GPU memory is insufficient.
• Gradient checkpointing (Part #4 A.7): Save memory by recomputing activations.
• DeepSpeed/FSDP: Train large models.

Experiment Management

• Weights & Biases (wandb): Experiment logging and visualization.
• MLflow: Open-source experiment tracking.
• Lightning: Reduce boilerplate code in the training loop.

PyTorch Lightning. A wrapper that hides the repetitive code in the training loop.

import pytorch_lightning as pl

class CellClassifierLightning(pl.LightningModule):
    def __init__(self, model, lr=1e-3):
        super().__init__()
        self.model = model
        self.lr = lr
        self.loss_fn = nn.CrossEntropyLoss()
    
    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self.model(x)
        loss = self.loss_fn(logits, y)
        self.log("train_loss", loss)
        return loss
    
    def configure_optimizers(self):
        return torch.optim.AdamW(self.parameters(), lr=self.lr)

trainer = pl.Trainer(max_epochs=10, accelerator="gpu", devices=1)
trainer.fit(model_lightning, train_loader)

Automatically handles multi-GPU, mixed precision, and checkpoints. Convenient for production training.

Bio Application Scenarios

Scenario 1 — Cell Type Classifier in Practice

Fully implement Scenario 1 from Part #3. AnnData → PyTorch Dataset → Training. Refer to the Scanpy tutorial.

Scenario 2 — Protein Contact Prediction (mini AlphaFold)

Protein sequence → Attention-based neural network → Probability of amino acid pair contact. Practice the concepts from Parts #5 and #6. Can be trained on a small dataset (e.g., CATH).

Scenario 3 — Tissue Pathology Image Classification

Implement the pathology slide scenario from Parts #2 and #6. ResNet-50 fine-tuning. Utilize Torchvision.datasets.

Key Summary

• PyTorch 4 components: Tensor, Autograd, nn.Module, Optimizer. The principles from Parts #2 to #4 are directly applied in the API.
• Training loop 5 steps: forward, loss, backward, zero_grad, step. Exactly as shown in the diagram in Part #3.
• Use DataLoader for batch processing and to(device) to utilize the GPU.
• Accelerate training by 2-4 times with Mixed precision (BF16).
• Regularly save Checkpoints.
• Train large models with Multi-GPU (DDP/FSDP).
• Use tools like Lightning and wandb for practical convenience.

📐 Appendix — Practical Tips and Formulas for Experts

Difficulty: Very Hard Target Audience: Readers seeking advanced optimization and system-level understanding of PyTorch training.

A.1 Internal Mechanics of Autograd

When a tensor with requires_grad=True is used in an operation, each operation is recorded as a graph node.

Function Class: Each operation (add, mul, matmul, etc.) is a subclass of torch.autograd.Function, defining the static methods forward and backward.

Custom Autograd:

class MyReLU(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        return x.clamp(min=0)

    @staticmethod
    def backward(ctx, grad_output):
        x, = ctx.saved_tensors
        grad_input = grad_output.clone()
        grad_input[x < 0] = 0
        return grad_input

ctx.save_for_backward stores the tensors needed for the backward pass. The activation value saving in section 4 utilizes this mechanism.

A.2 Gradient Checkpointing

See section 4, A.7. torch.utils.checkpoint.

from torch.utils.checkpoint import checkpoint

class MyBlock(nn.Module):
    def forward(self, x):
        # The activations of this function are not stored and will be recomputed during the backward pass.
        return checkpoint(self._forward, x, use_reentrant=False)

    def _forward(self, x):
        # Actual computation
        return self.layers(x)

Memory usage is O(√L). The runtime increases by a factor of 1.5. Essential for large models.

A.3 Numerical Aspects of Mixed Precision

FP32: exponent 8, mantissa 23. Range: ~1e-38 to 1e+38, precision: ~7 digits.

FP16: exponent 5, mantissa 10. Range: ~6e-5 to 6.5e+4, precision: ~3 digits. Risk of overflow.

BF16: exponent 8, mantissa 7. Range: same as FP32, precision: ~2 digits. Overflow safe.

Loss Scaling (required for FP16):

scaled_loss = loss * scale_factor
scaled_loss.backward()  # Gradients are multiplied by the scale_factor.
# Unscale the gradients before the optimizer step.

Loss scaling is not necessary for BF16.

A.4 DDP Synchronization

All-Reduce Gradient Synchronization: Each GPU computes the gradients for its mini-batch, and then all GPUs compute the average of the gradients.

Time Complexity: O(P / bandwidth), where P is the number of parameters.

NCCL Backend: NVIDIA's optimized library that utilizes InfiniBand and NVLink for communication between GPUs.

Bucketing: Groups gradients into chunks to overlap communication and computation. torch.nn.parallel.DistributedDataParallel does this automatically.

A.5 FSDP (Fully Sharded Data Parallel)

Parameters and optimizer state are sharded across GPUs.

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

model = FSDP(model, sharding_strategy=ShardingStrategy.FULL_SHARD)

Memory Reduction: Each GPU only stores 1/N of the parameters. A 70B parameter model can fit on 8 × 40GB GPUs.

Communication Overhead: Parameters are gathered and released during the forward and backward passes. NCCL optimization is essential.

A.6 Custom CUDA Kernel

Some operations are not optimally implemented in PyTorch. Custom kernels can be written using Triton or CUDA C.

Flash Attention 2: The online softmax from section 6, A.8 is implemented in Triton and CUDA.

from flash_attn import flash_attn_func

output = flash_attn_func(q, k, v, causal=True)

2-4 times faster than nn.functional.scaled_dot_product_attention.

Recent Trends: torch.compile (PyTorch 2.0) automatically optimizes kernels. model = torch.compile(model). In most cases, it improves performance by 2x or more.

A.7 Profiling

PyTorch Profiler:

with torch.profiler.profile(
    activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA],
    schedule=torch.profiler.schedule(wait=1, warmup=1, active=3),
    on_trace_ready=torch.profiler.tensorboard_trace_handler("./log")
) as prof:
    for step, batch in enumerate(train_loader):
        # ... training step
        prof.step()

Visualize with TensorBoard. Identify bottlenecks (slow kernels, CPU-GPU transfer, data loading).

A.8 Reproducibility

Ensure that the same experiment yields the same results when repeated.

import torch
import random
import numpy as np

def set_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    random.seed(seed)
    np.random.seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

Perfect reproducibility is difficult (some CUDA operations are non-deterministic). In practice, fix the seed and record the hardware and PyTorch version.

A.9 Debugging Tips

NaN Detection:

torch.autograd.set_detect_anomaly(True)

If a NaN occurs during the backward pass, this will trace back to the operation where it originated.

Gradient Clipping (section 4, A.8):

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

Apply after loss.backward() and before optimizer.step().

Batch size finder (Lightning): Automatically searches for the maximum possible batch size.

A.10 Latest Trends

torch.compile (PyTorch 2.0+): Automatic optimization through JIT compilation.

torch.export: Extracts a model into a graph (ONNX, TensorRT, mobile deployment).

AOT autograd: Generates the backward graph at compile time.

Nested Tensor: Processes variable-length sequences without padding.

FunctorchModule → torch.func: JAX-style functional API.

These tools are becoming the standard for training and serving modern large models.

References

All content, scenarios, analogies, and formulas in this section are developed internally by BioPlayground. The following are external resources that can help with concept learning.

• PyTorch Official Documentation: pytorch.org/docs
• PyTorch 60-Minute Tutorial: pytorch.org/tutorials/beginner/deep_learning_60min_blitz.html
• PyTorch Lightning: lightning.ai
• HuggingFace Transformers: huggingface.co/docs/transformers
• Flash Attention: github.com/Dao-AILab/flash-attention
• DeepSpeed: deepspeed.ai
• PyTorch Profiler: pytorch.org/tutorials/recipes/recipes/profiler_recipe.html
• Scanpy: scanpy.readthedocs.io — standard library for scRNA-seq
• PyTorch Geometric: pyg.org — graph neural networks (molecules, proteins)
• TorchVision: pytorch.org/vision — image domain

Section 12 completes the review of PyTorch basics. Section 13 will cover using the HuggingFace model hub and OpenAI/Anthropic APIs.

Next Concepts

• Ep. #13 huggingface-and-openai — Leveraging Hugging Face Model Hub and Commercial APIs.
• Ep. #14 claude-code-and-cursor — AI Coding Agent in Action.
• Ep. #15 bio-ai-integration — Integrating It All into a Bio Pipeline.