PyTorch

Learning Objectives

• Understand what PyTorch is and why it became the dominant ML framework
• Identify PyTorch's key features and how they compare to alternatives
• Evaluate PyTorch's role in the broader AI development ecosystem

What Is PyTorch?

PyTorch is an open-source machine learning framework that provides the core building blocks for training and deploying AI models. Originally created by Meta AI Research in 2016, it is now governed by the PyTorch Foundation under the Linux Foundation (transferred from Meta in September 2022).

PyTorch is the framework most AI researchers and an increasing number of production teams use to build neural networks. When you hear that a team "trained a model," they almost certainly used PyTorch (or occasionally JAX). It provides:

• Tensor computation — GPU-accelerated multi-dimensional array operations (the mathematical foundation of neural networks)
• Automatic differentiation — computes gradients automatically for training via backpropagation
• Neural network modules — pre-built layers, optimizers, loss functions, and training utilities
• GPU acceleration — native CUDA support for NVIDIA GPUs; ROCm for AMD; MPS for Apple Silicon

Why PyTorch Won

PyTorch wasn't the first deep learning framework — TensorFlow (Google, 2015) came first and initially dominated. PyTorch overtook TensorFlow through a key design decision: dynamic computation graphs.

• TensorFlow 1.x required defining the entire computation graph before running it (static graph). This was efficient for deployment but painful for research and debugging.
• PyTorch uses dynamic graphs — you build the computation as you run it, using standard Python control flow. This makes debugging intuitive (use print statements, breakpoints, standard Python tools) and experimentation fast.

By 2020, PyTorch was the dominant framework in AI research papers. By 2024, it had largely won the production battle as well, especially after TensorFlow's Keras team moved to support PyTorch as a backend.

Key Components

torch — Core Tensor Library

GPU-accelerated tensor operations. Tensors are the fundamental data structure — multi-dimensional arrays that flow through neural networks.

torch.nn — Neural Network Modules

Pre-built building blocks:

• Layers: Linear, Conv2d, Transformer, Attention, Embedding
• Activations: ReLU, GELU, SiLU
• Loss functions: CrossEntropy, MSE, BCE
• Containers: Sequential, ModuleList, ModuleDict

torch.optim — Optimizers

Training algorithms that update model weights: Adam, AdamW, SGD, and variants. AdamW is the default choice for most modern AI training.

torch.distributed — Distributed Training

Scale training across multiple GPUs and multiple machines. Essential for training large models — every frontier model (GPT, Claude, Llama, Gemini) uses distributed training.

TorchVision, TorchAudio, TorchText

Domain-specific libraries for computer vision, audio processing, and NLP — providing datasets, model architectures, and data transformations.

PyTorch in the AI Ecosystem

PyTorch sits at the center of a vast ecosystem:

• Hugging Face Transformers — the most popular model library, built on PyTorch
• Lightning — high-level PyTorch wrapper that reduces training boilerplate
• ONNX — export PyTorch models to a portable format for deployment
• TensorRT — NVIDIA's inference optimizer accepts PyTorch models
• ExecuTorch — Meta's framework for deploying PyTorch models on mobile/edge devices
• Torchtune — fine-tuning library for LLMs built natively on PyTorch

PyTorch vs. Alternatives

For most developers, PyTorch is the default choice. JAX is the main alternative for teams in Google's ecosystem or those needing advanced functional programming patterns.

Access

Strengths

• Dominant framework — used by the majority of AI researchers and production teams; largest ecosystem
• Python-native — dynamic computation graphs; standard Python debugging; intuitive API
• Massive ecosystem — Hugging Face, Lightning, TorchVision, thousands of libraries built on PyTorch
• GPU support — NVIDIA CUDA (primary), AMD ROCm, Apple MPS
• Independent governance — Linux Foundation stewardship ensures vendor neutrality
• Free and open source — BSD license; no restrictions on commercial use

Limitations & Considerations

• Learning curve — understanding tensors, autograd, and GPU memory management requires investment
• Verbose for simple tasks — training loops, data loading, and logging require more boilerplate than higher-level tools (Lightning helps)
• NVIDIA-optimized — while ROCm and MPS support exists, CUDA remains the best-supported GPU backend
• Not a deployment framework — PyTorch trains models; deploying them efficiently requires additional tools (TorchServe, TensorRT, ONNX)
• Memory hungry — training large models requires careful memory management; gradient checkpointing and mixed precision are often necessary

Key Takeaways

• PyTorch is the dominant open-source ML framework — the foundation on which most AI models are trained, from research prototypes to frontier models like Llama
• Originally created by Meta AI Research (2016), now independently governed by the PyTorch Foundation under the Linux Foundation
• Dynamic computation graphs, Python-native design, and a massive ecosystem (Hugging Face, Lightning, TorchVision) are the key reasons it overtook TensorFlow
• Free, open-source (BSD license), and the default starting point for any AI development project