PyTorch has become the dominant framework for deep learning research and production, used by Google, Meta, Tesla, and nearly every top AI lab. Its dynamic computation graph, Pythonic API, and strong ecosystem make it the framework of choice for data scientists who want full control. This guide takes you from tensors to training production-ready neural networks.
Tensors – The Foundation
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
# Create tensors
x = torch.tensor([1.0, 2.0, 3.0])
y = torch.zeros(3, 4) # 3x4 zeros
z = torch.randn(2, 3) # random normal
I = torch.eye(4) # identity matrix
# GPU acceleration
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'Using: {device}')
x_gpu = x.to(device) # move to GPU
x_cpu = x_gpu.cpu().numpy() # back to numpy
# Autograd — automatic differentiation
x = torch.tensor(3.0, requires_grad=True)
y = x ** 2 + 2 * x + 1
y.backward()
print(f'dy/dx at x=3: {x.grad}') # 2*3 + 2 = 8
Building Neural Networks with nn.Module
class MLP(nn.Module):
def __init__(self, input_size, hidden_size, output_size, dropout=0.3):
super().__init__()
self.network = nn.Sequential(
nn.Linear(input_size, hidden_size),
nn.BatchNorm1d(hidden_size),
nn.ReLU(),
nn.Dropout(dropout),
nn.Linear(hidden_size, hidden_size // 2),
nn.BatchNorm1d(hidden_size // 2),
nn.ReLU(),
nn.Dropout(dropout),
nn.Linear(hidden_size // 2, output_size)
)
def forward(self, x):
return self.network(x)
model = MLP(input_size=20, hidden_size=128, output_size=2).to(device)
print(model)
print(f'Parameters: {sum(p.numel() for p in model.parameters()):,}')
Data Loading
from torch.utils.data import Dataset, DataLoader
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
class TabularDataset(Dataset):
def __init__(self, X, y):
self.X = torch.FloatTensor(X)
self.y = torch.LongTensor(y)
def __len__(self):
return len(self.X)
def __getitem__(self, idx):
return self.X[idx], self.y[idx]
X, y = make_classification(n_samples=5000, n_features=20, random_state=42)
X_tr, X_te, y_tr, y_te = train_test_split(X, y, test_size=0.2, random_state=42)
train_loader = DataLoader(TabularDataset(X_tr, y_tr),
batch_size=64, shuffle=True, num_workers=2)
test_loader = DataLoader(TabularDataset(X_te, y_te),
batch_size=64, shuffle=False)
Training Loop
criterion = nn.CrossEntropyLoss()
optimizer = optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=50)
def train_epoch(model, loader, criterion, optimizer, device):
model.train()
total_loss, correct = 0, 0
for X_batch, y_batch in loader:
X_batch, y_batch = X_batch.to(device), y_batch.to(device)
optimizer.zero_grad()
outputs = model(X_batch)
loss = criterion(outputs, y_batch)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
total_loss += loss.item()
correct += (outputs.argmax(1) == y_batch).sum().item()
return total_loss / len(loader), correct / len(loader.dataset)
def evaluate(model, loader, criterion, device):
model.eval()
total_loss, correct = 0, 0
with torch.no_grad():
for X_batch, y_batch in loader:
X_batch, y_batch = X_batch.to(device), y_batch.to(device)
outputs = model(X_batch)
total_loss += criterion(outputs, y_batch).item()
correct += (outputs.argmax(1) == y_batch).sum().item()
return total_loss / len(loader), correct / len(loader.dataset)
best_val_acc, patience, wait = 0, 10, 0
for epoch in range(100):
train_loss, train_acc = train_epoch(model, train_loader,
criterion, optimizer, device)
val_loss, val_acc = evaluate(model, test_loader, criterion, device)
scheduler.step()
if val_acc > best_val_acc:
best_val_acc = val_acc
torch.save(model.state_dict(), 'best_model.pt')
wait = 0
else:
wait += 1
if wait >= patience:
print(f'Early stopping at epoch {epoch}')
break
if epoch % 10 == 0:
print(f'Epoch {epoch:3d} | Train: {train_acc:.4f} | Val: {val_acc:.4f}')
print(f'Best val accuracy: {best_val_acc:.4f}')
Convolutional Neural Network (Image Classification)
class CNN(nn.Module):
def __init__(self, num_classes=10):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 32, kernel_size=3, padding=1),
nn.BatchNorm2d(32), nn.ReLU(),
nn.MaxPool2d(2, 2),
nn.Conv2d(32, 64, kernel_size=3, padding=1),
nn.BatchNorm2d(64), nn.ReLU(),
nn.MaxPool2d(2, 2),
nn.Conv2d(64, 128, kernel_size=3, padding=1),
nn.BatchNorm2d(128), nn.ReLU(),
nn.AdaptiveAvgPool2d((4, 4))
)
self.classifier = nn.Sequential(
nn.Flatten(),
nn.Linear(128 * 4 * 4, 256),
nn.ReLU(), nn.Dropout(0.5),
nn.Linear(256, num_classes)
)
def forward(self, x):
return self.classifier(self.features(x))
Transfer Learning
import torchvision.models as models
# Load pretrained ResNet-50
backbone = models.resnet50(weights='IMAGENET1K_V2')
# Freeze all layers except the final classifier
for param in backbone.parameters():
param.requires_grad = False
# Replace the head for your number of classes
backbone.fc = nn.Sequential(
nn.Linear(backbone.fc.in_features, 256),
nn.ReLU(), nn.Dropout(0.3),
nn.Linear(256, NUM_CLASSES)
)
# Only the new head will be trained
optimizer = optim.Adam(backbone.fc.parameters(), lr=1e-3)
Saving and Loading Models
# Save
torch.save(model.state_dict(), 'model.pt')
# Load
model = MLP(20, 128, 2).to(device)
model.load_state_dict(torch.load('model.pt', map_location=device))
model.eval()
# Export to TorchScript for production (no Python needed)
scripted = torch.jit.script(model)
scripted.save('model_scripted.pt')
Conclusion
PyTorch gives you full transparency into every computation, making debugging intuitive and customisation unlimited. Master the training loop pattern shown here — it works for any architecture from simple MLP to Transformer. Always use mixed precision training (`torch.cuda.amp`) for GPU jobs, use DataLoader with `num_workers > 0` for fast data loading, and save the best checkpoint based on validation accuracy, not final epoch. These three habits alone will save you hours on every deep learning project.
