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

User request: "Fine-tune a pretrained ResNet for classifying product images"

Actions:

1. Load

   resnet50(weights=ResNet50_Weights.DEFAULT)

   and freeze all layers except the final classifier
2. Replace the classifier head with

   nn.Linear(2048, num_classes)

3. Set up DataLoader with image augmentation transforms (RandomCrop, ColorJitter, Normalize)
4. Train with AdamW, CosineAnnealingLR scheduler, and mixed precision

Output: A fine-tuned image classifier with production-quality accuracy and efficient mixed-precision training.

Example 2: Train a text classification model with Hugging Face

User request: "Build a sentiment analysis model using a pretrained transformer"

Actions:

1. Load

   AutoModel.from_pretrained("bert-base-uncased")

   with a classification head
2. Tokenize the dataset using

   AutoTokenizer

   and create a DataLoader
3. Fine-tune with AdamW, linear warmup scheduler, and gradient clipping
4. Export the trained model with

   torch.export()

   for production serving

Output: A sentiment analysis model fine-tuned on custom data and exported for production inference.

Guidelines

• Use

  torch.compile(model)

  on PyTorch 2.0+ for a free 20-50% speedup with one line.
• Use

  AdamW

  over

  Adam

  for correct weight decay implementation with modern architectures.
• Use mixed precision (

  torch.amp

  ) for any GPU training to halve memory and double throughput.
• Move data to device in the training loop, not in the Dataset, to keep Dataset device-agnostic.
• Use

  model.eval()

  and

  torch.no_grad()

  during inference to prevent unnecessary gradient computation.
• Use

  pin_memory=True

  in DataLoader when training on GPU to speed up CPU-to-GPU data transfer.
• Save

  model.state_dict()

  not the full model since state dicts are portable across code changes.