PyTorch basics for developers: a practical AI & ML guide

Introduction

If you've been building applications in production and want to add machine learning capabilities without reinventing the wheel, PyTorch basics for developers: a practical AI & ML guide will get you up to speed fast. PyTorch has become the de facto standard in the ML community because it feels natural to Python developers—no mystical abstractions, just computational graphs that behave like code. By the end of this guide, you'll understand tensors, build a working model, and run a full training loop on GPU.

Why PyTorch for Production Teams

Before we dive into tensors and training loops, let's acknowledge why we're here. You've shipped REST APIs, managed databases, and debugged distributed systems. PyTorch respects that experience. Unlike some frameworks that hide the machinery, PyTorch keeps you close to the math—you write the forward pass, you call backward, you step the optimizer. It's Pythonic. It's imperative. It feels like the code you already write.

We've used PyTorch in production for everything from real-time recommendation engines in e-commerce to risk scoring in fintech. The framework scales from a laptop prototype to multi-GPU training on thousands of samples per second. That alone makes it worth learning.

Tensors: The Foundation

A tensor is PyTorch's core data structure. Think of it as a NumPy array's sophisticated cousin—it can live on CPU or GPU, and it tracks gradients for automatic differentiation. That last point is crucial: tensors remember how they were created so they can compute derivatives.

Here's what tensors look like in practice:

import torch

# Create a tensor on CPU
x = torch.randn(3, 4)  # 3x4 matrix of random values
print(x.shape)  # torch.Size([3, 4])
print(x.device)  # cpu

# Move to GPU if available
if torch.cuda.is_available():
    x = x.to('cuda')
    print(x.device)  # cuda:0

# Requires gradient tracking for parameters
w = torch.randn(4, 2, requires_grad=True)
print(w.requires_grad)  # True

# Basic operations
y = torch.matmul(x, w)  # Matrix multiplication
print(y.shape)  # torch.Size([3, 2])
print(y.requires_grad)  # True (inherited from w)

Notice that requires_grad=True. When a tensor has this flag, PyTorch builds a computation graph as you perform operations. Later, when you call backward(), it traces back through that graph to compute gradients. This is autograd, and it's the engine that powers training.

For most developers, you won't create tensors manually with requires_grad=True. Your model's parameters will already have that flag set. But understanding what it means will save you hours debugging silent failures where gradients vanish.

Define a Model by Subclassing nn.Module

PyTorch models inherit from nn.Module. This class gives you the scaffolding: parameter tracking, device management, train/eval modes, and more. Every model you build in production follows this pattern.

import torch
import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(10, 32),
            nn.ReLU(),
            nn.Linear(32, 1)
        )
    
    def forward(self, x):
        return self.layers(x)

# Instantiate the model
model = Net()

# Inspect parameters
for name, param in model.named_parameters():
    print(f"{name}: {param.shape}, requires_grad={param.requires_grad}")

This is a simple feedforward network: input layer with 10 features, hidden layer with 32 units and ReLU activation, output layer with 1 unit. In production, you might replace nn.Sequential with explicit layer definitions if you need branches or skip connections, but the pattern stays the same.

Key points:

• __init__ declares all layers as attributes. PyTorch's nn.Module metaclass finds them and registers them.
• forward defines the computation. Never call it directly; use model(x) instead, which triggers hooks for training/eval state.
• Parameters are automatically set to requires_grad=True so gradients flow during backpropagation.

We've seen teams build impressive models this way—recommendation systems with embeddings, sequence models for anomaly detection—all following this simple blueprint.

The Training Loop: Forward, Loss, Backward, Step

Now for the core rhythm of training. Every iteration, you:

1. Forward pass: Feed data through the model, get predictions.
2. Compute loss: Compare predictions to ground truth.
3. Backward pass: Compute gradients.
4. Optimizer step: Update parameters using those gradients.

import torch
import torch.nn as nn
import torch.optim as optim

# Setup
model = Net()
criterion = nn.MSELoss()  # Loss function for regression
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Dummy data
X = torch.randn(100, 10)  # 100 samples, 10 features
y = torch.randn(100, 1)   # 100 targets

# Training loop
num_epochs = 5
batch_size = 16

for epoch in range(num_epochs):
    epoch_loss = 0.0
    
    # Iterate over batches
    for i in range(0, len(X), batch_size):
        batch_X = X[i:i+batch_size]
        batch_y = y[i:i+batch_size]
        
        # Forward pass
        predictions = model(batch_X)
        
        # Compute loss
        loss = criterion(predictions, batch_y)
        
        # Backward pass
        optimizer.zero_grad()  # Clear old gradients
        loss.backward()         # Compute new gradients
        
        # Optimizer step
        optimizer.step()        # Update parameters
        
        epoch_loss += loss.item()
    
    print(f"Epoch {epoch+1}, Loss: {epoch_loss / (len(X) // batch_size):.4f}")

A few details matter here:

• optimizer.zero_grad() is essential. PyTorch accumulates gradients by default; without this, old gradients persist.
• loss.backward() traverses the computation graph and populates the .grad attribute of every parameter.
• optimizer.step() uses those gradients to update parameters. The update rule depends on your optimizer (Adam, SGD, etc.).
• loss.item() extracts the scalar value from a tensor, which is useful for logging.

In production, you'd wrap this in validation logic, checkpoint the model, and handle edge cases. But this loop is the heartbeat. Everything else is scaffolding.

Move to GPU When Available

Once your training loop works on CPU, moving to GPU is straightforward—but easy to get wrong if you're not careful. Both the model and the data must be on the same device.

device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"Using device: {device}")

# Move model to device
model = model.to(device)

# Move data to device in the training loop
for epoch in range(num_epochs):
    for i in range(0, len(X), batch_size):
        batch_X = X[i:i+batch_size].to(device)
        batch_y = y[i:i+batch_size].to(device)
        
        predictions = model(batch_X)
        loss = criterion(predictions, batch_y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

GPU training is typically 10–100× faster than CPU, depending on your model size and hardware. We've seen a model that took 45 minutes per epoch on a production CPU server run in 3 minutes on a single GPU. That difference compounds over thousands of training steps, turning a 2-week project into a 2-day project.

Pro tip: Use torch.cuda.is_available() to gracefully fall back to CPU in your development environment or CI/CD pipelines where GPUs may not be present.

Common Mistakes

Forgetting optimizer.zero_grad(): Gradients accumulate. Your first iteration trains normally; the second iteration has double the gradient; by iteration ten, gradients explode and training diverges. Always zero before backward.

Mixing devices: Calling model(cpu_tensor) when the model is on GPU raises "Tensors on different devices" errors. Move all data to the same device before passing it to the model.

Training on evaluation data: If you evaluate loss on training data every batch, you're not measuring generalization. Use separate train and test sets, and toggle model.eval() during evaluation (this disables dropout and batch norm updates).

Not detaching or no_grad for inference: If you run inference in a training loop—say, to validate on a test set—wrap it in torch.no_grad(). Otherwise, PyTorch builds a computation graph for operations you won't backpropagate through, wasting memory.

model.eval()
with torch.no_grad():
    test_predictions = model(test_X)
    test_loss = criterion(test_predictions, test_y)
print(f"Test Loss: {test_loss.item():.4f}")

Putting It Together: A Complete Example

Here's a minimal but complete training script:

import torch
import torch.nn as nn
import torch.optim as optim

# Define model
class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(10, 32),
            nn.ReLU(),
            nn.Linear(32, 1)
        )
    
    def forward(self, x):
        return self.layers(x)

# Setup
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = SimpleNet().to(device)
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.MSELoss()

# Dummy data
X_train = torch.randn(100, 10).to(device)
y_train = torch.randn(100, 1).to(device)

# Training
for epoch in range(10):
    model.train()
    loss = criterion(model(X_train), y_train)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")

This works. It trains. It runs on GPU if available. From here, you'd add validation, checkpointing, hyperparameter sweeps—but you have the foundation.

Key Takeaways

• Tensors are like NumPy arrays but support GPU and automatic differentiation. Use requires_grad=True to track gradients.
• Define models by subclassing nn.Module, implementing __init__ and forward. PyTorch automatically manages parameters.
• The training loop is straightforward: forward, compute loss, zero gradients, backward, optimizer step. This pattern scales to complex models.
• Move models and data to GPU with .to(device). Check torch.cuda.is_available() for graceful fallback.
• Always call optimizer.zero_grad() before backward() to avoid gradient accumulation.
• Use torch.no_grad() and model.eval() for inference to save memory and disable training-specific behavior.
• Start simple, validate on a test set, and iterate. PyTorch rewards clear, explicit code over magic.