PyTorch

Overview

PyTorch is a deep learning framework for building and training neural networks with dynamic computation graphs and automatic differentiation. It provides tensor operations with GPU acceleration, nn.Module for defining architectures, DataLoader for efficient data loading, mixed precision training for performance, and export tools (TorchScript, ONNX) for production deployment.

Instructions

• When defining models, subclass nn.Module with __init__ for layers and forward for computation, using nn.Sequential for simple stacks and custom forward logic for complex architectures.
• When training, implement the standard loop: forward pass, loss computation, loss.backward(), optimizer.step(), optimizer.zero_grad(), with gradient clipping via clip_grad_norm_ for stability.
• When loading data, subclass Dataset with __len__ and __getitem__, then use DataLoader with num_workers=4 and pin_memory=True for GPU training throughput.
• When optimizing performance, use torch.compile(model) on PyTorch 2.0+ for 20-50% speedup, mixed precision with torch.amp.autocast() for halved memory and doubled throughput, and DistributedDataParallel for multi-GPU training.
• When doing transfer learning, load pretrained models from torchvision.models or Hugging Face, freeze the backbone, and replace the classifier head for your task.
• When deploying, use torch.export() or torch.jit.trace() for production, torch.onnx.export() for cross-framework compatibility, and torch.quantization for INT8 inference speedup.

Examples

Example 1: Fine-tune a vision model for image classification