PyTorch Training Loop

Introduction

PyTorch training involves DataLoader for batching, optimizer for weight updates, and a training loop.

DataLoader

from torch.utils.data import DataLoader, TensorDataset

# From tensors
X = torch.randn(1000, 784)
y = torch.randint(0, 10, (1000,))
dataset = TensorDataset(X, y)

loader = DataLoader(dataset, batch_size=32, shuffle=True)

# Using custom dataset
class CustomDataset(torch.utils.data.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]

Training Loop

model = SimpleNet()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

num_epochs = 10

for epoch in range(num_epochs):
    for batch_X, batch_y in loader:
        # Forward
        outputs = model(batch_X)
        loss = criterion(outputs, batch_y)

        # Backward
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    print(f"Epoch [{epoch+1}], Loss: {loss.item():.4f}")

Validation

model.eval()  # Dropout, BatchNorm freeze
with torch.no_grad():
    correct = 0
    total = 0
    for batch_X, batch_y in val_loader:
        outputs = model(batch_X)
        _, predicted = torch.max(outputs.data, 1)
        total += batch_y.size(0)
        correct += (predicted == batch_y).sum()

    accuracy = correct / total
    print(f"Validation Accuracy: {accuracy:.4f}")

model.train()  # Resume training mode

GPU Training

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

for batch_X, batch_y in loader:
    batch_X = batch_X.to(device)
    batch_y = batch_y.to(device)
    # ... training code

Practice Problems

1. Create DataLoader from data
2. Implement training loop
3. Add validation
4. Move to GPU
5. Save and load model