PyTorch Cheat Sheet — Tensors, Models & Training — PyTorch R

import torch

# Create tensors
x = torch.tensor([[1.0, 2.0], [3.0, 4.0]])  # from data
y = torch.zeros(3, 4)  # zeros
z = torch.randn(2, 5)  # random normal

# Reshape and device
x_reshaped = x.reshape(4, 1)
x_gpu = x.to('cuda') if torch.cuda.is_available() else x.cpu()

# Operations
result = torch.matmul(x, x.T)  # matrix multiply
result = x + y  # broadcasting works
result = x.sum(dim=1)  # reduce
print(result.shape, result.dtype)

import torch

# Enable gradient tracking
x = torch.tensor([2.0, 3.0], requires_grad=True)
y = x ** 2
z = y.sum()

# Backprop
z.backward()
print(x.grad)  # [4.0, 6.0] — dz/dx = 2x

# Gradient accumulation (default behavior)
z.backward()  # grad is ADDED to existing x.grad
print(x.grad)  # [8.0, 12.0] — accumulated!

# Clear gradients before next backward
x.grad.zero_()
z.backward()
print(x.grad)  # [4.0, 6.0] — fresh

# Disable grad tracking (inference)
with torch.no_grad():
    pred = model(x)  # no gradient computation

import torch
import torch.nn as nn

class SimpleNN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super().__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, output_size)
        self.dropout = nn.Dropout(0.2)
    
    def forward(self, x):
        x = self.fc1(x)
        x = torch.relu(x)  # activation
        x = self.dropout(x)  # regularization
        x = self.fc2(x)
        return x

# Instantiate
model = SimpleNN(input_size=10, hidden_size=64, output_size=2)
print(sum(p.numel() for p in model.parameters()))  # total params

# Move to device
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = model.to(device)

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

# Dummy data
X_train = torch.randn(100, 10)
y_train = torch.randint(0, 2, (100,))
dataset = TensorDataset(X_train, y_train)
loader = DataLoader(dataset, batch_size=16, shuffle=True)

# Model, loss, optimizer
model = nn.Sequential(
    nn.Linear(10, 64),
    nn.ReLU(),
    nn.Linear(64, 2)
)
criterion = nn.CrossEntropyLoss()
optimizer = Adam(model.parameters(), lr=0.001)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model.to(device)

# Training
num_epochs = 3
for epoch in range(num_epochs):
    total_loss = 0
    for X_batch, y_batch in loader:
        X_batch, y_batch = X_batch.to(device), y_batch.to(device)
        
        # Forward
        logits = model(X_batch)
        loss = criterion(logits, y_batch)
        
        # Backward
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        total_loss += loss.item()
    
    print(f"Epoch {epoch+1}/{num_epochs}, Loss: {total_loss/len(loader):.4f}")

import torch

# Check GPU
print(f"GPU available: {torch.cuda.is_available()}")
print(f"GPU count: {torch.cuda.device_count()}")
print(f"Current GPU: {torch.cuda.current_device()}")
print(f"GPU name: {torch.cuda.get_device_name(0)}")

# Move tensors and models to GPU
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
x = torch.randn(1000, 1000).to(device)
model = model.to(device)

# Automatic Mixed Precision (AMP) for speed + memory
from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for epoch in range(num_epochs):
    for X_batch, y_batch in loader:
        X_batch, y_batch = X_batch.to(device), y_batch.to(device)
        
        optimizer.zero_grad()
        with autocast():
            logits = model(X_batch)
            loss = criterion(logits, y_batch)
        
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

# Synchronize (for timing)
torch.cuda.synchronize()

import torch
import torch.nn as nn
from pathlib import Path

model = nn.Linear(10, 2)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# Save model state (recommended)
torch.save(model.state_dict(), 'model_weights.pth')
model.load_state_dict(torch.load('model_weights.pth', weights_only=True))

# Save entire model (includes architecture)
torch.save(model, 'model_full.pth')
model = torch.load('model_full.pth')

# Save checkpoint (training state)
checkpoint = {
    'epoch': 5,
    'model_state': model.state_dict(),
    'optimizer_state': optimizer.state_dict(),
    'loss': 0.123
}
torch.save(checkpoint, 'checkpoint.pth')

# Resume training
checkpoint = torch.load('checkpoint.pth')
model.load_state_dict(checkpoint['model_state'])
optimizer.load_state_dict(checkpoint['optimizer_state'])
epoch = checkpoint['epoch']