Data Pipeline & Advanced Training

Dataset & DataLoader

from torch.utils.data import Dataset, DataLoader
from torchvision import transforms

class ImageDataset(Dataset):
    def __init__(self, file_paths, labels):
        self.file_paths = file_paths
        self.labels = labels
        self.transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])

    def __len__(self):
        return len(self.file_paths)

    def __getitem__(self, idx):
        image = load_image(self.file_paths[idx])
        return self.transform(image), self.labels[idx]

train_loader = DataLoader(ImageDataset(paths, labels), batch_size=32, shuffle=True, num_workers=4)

Transfer Learning

from torchvision import models

model = models.resnet18(weights='IMAGENET1K_V1')

# Freeze the pretrained backbone — no gradients computed/applied for it
for param in model.parameters():
    param.requires_grad = False

# Replace and train only the task-specific head
model.fc = nn.Linear(model.fc.in_features, num_classes)

optimizer = optim.Adam(model.fc.parameters(), lr=1e-3)  # only the head's params

Mixed Precision & Multi-GPU

scaler = torch.cuda.amp.GradScaler()

for images, labels in train_loader:
    optimizer.zero_grad()
    with torch.autocast(device_type='cuda', dtype=torch.float16):
        outputs = model(images)
        loss = criterion(outputs, labels)

    scaler.scale(loss).backward()   # scale up to avoid float16 gradient underflow
    scaler.step(optimizer)          # unscales, then steps
    scaler.update()

# Gradient clipping — prevents exploding gradients from huge updates
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

# DistributedDataParallel — one process per GPU, recommended over DataParallel
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])

# torch.compile — JIT-compiles the graph for a speedup, minimal code change
model = torch.compile(model)

Save, Load & Checkpointing

# Recommended: save just the learned parameters, not the whole object
torch.save(model.state_dict(), 'model.pt')

# Reconstruct the architecture first, then load weights into it
model = MLP()
model.load_state_dict(torch.load('model.pt'))
model.eval()

# Full training checkpoint — resume training later
torch.save({
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss,
}, 'checkpoint.pt')