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 (torch.export, ONNX) for production deployment.

Instructions

• Install with pip from the selector at pytorch.org/get-started/locally (it prints the right index URL for your CUDA version or CPU-only). PyTorch 2.9 and later need Python 3.10+.
• Pick the device once: device = "cuda" if torch.cuda.is_available() else "cpu" (Apple Silicon: "mps").
• When defining models, subclass nn.Module with __init__ for layers and forward for computation, using nn.Sequential for simple stacks.
• When training, run the standard loop: optimizer.zero_grad(), forward pass, loss, loss.backward(), clip_grad_norm_, optimizer.step(). Call model.train() before the loop and model.eval() for validation.
• When loading data, subclass Dataset with __len__ and __getitem__, then use DataLoader with num_workers=4, pin_memory=True and persistent_workers=True for GPU training.
• When optimizing performance, use torch.compile(model) (first call is slow while it compiles), mixed precision with torch.autocast(device_type="cuda", dtype=torch.bfloat16), and DistributedDataParallel launched with torchrun for multi-GPU. With float16 also use torch.amp.GradScaler("cuda"); bfloat16 does not need it.
• When doing transfer learning, load pretrained weights with the weights= enum (ResNet50_Weights.DEFAULT; the old pretrained=True was removed) or from Hugging Face, freeze the backbone, and replace the classifier head.
• When saving, store model.state_dict(); load with torch.load(path, map_location=device, weights_only=True). weights_only=True is the default since 2.6 and rejects arbitrary pickled objects.
• When deploying, use torch.export.export(model, example_args) for a graph you can ship, and torch.onnx.export(model, example_args, "model.onnx", dynamo=True) for ONNX (the dynamo exporter is the default since 2.9). TorchScript (torch.jit.trace/script) is deprecated; keep it only for legacy code.
• For INT8 or lower-precision inference, use the torchao package (torchao.quantization.quantize_). torch.ao.quantization is being phased out in favour of it.