DeepSeek-V3 from Scratch: Building the Architecture in PyTorch

Across the first 4 lessons, we have carefully constructed the building blocks of DeepSeek-V3: starting with its configuration and Rotary Position Embeddings (RoPE), advancing through Multihead Latent Attention (MLA), scaling capacity with the Mixture of Experts (MoE), and using Multi-Token Prediction (MTP) to improve training efficiency and enable faster inference through speculative decoding. Each of these innovations has added a vital piece to the architecture, preparing us for the next major milestone: bringing everything together into a unified system.

In this lesson, we will focus on assembling the full DeepSeek-V3 architecture. This is where theory and implementation converge: integrating RoPE, MLA, MoE, and MTP into a cohesive model that reflects the design principles of DeepSeek-V3. We will walk through how these components interact, the structural choices that make them synergize effectively, and the practical considerations for building a scalable, efficient language model. By the end, you will have a complete architecture blueprint, ready for the final lesson in the series, where we will implement the DeepSeek trainer and bring the model to life through training.

This lesson is the 5th in the 6-part series on Building DeepSeek-V3 from Scratch:

1. DeepSeek-V3 Model: Theory, Config, and Rotary Positional Embeddings
2. Build DeepSeek-V3: Multi-Head Latent Attention (MLA) Architecture
3. DeepSeek-V3 from Scratch: Mixture of Experts (MoE)
4. Autoregressive Model Limits and Multi-Token Prediction in DeepSeek-V3
5. DeepSeek-V3 from Scratch: Building the Architecture in PyTorch (this tutorial)
6. Lesson 6

To learn about DeepSeek-V3 and build it from scratch, just keep reading.

Looking for the source code to this post?

Jump Right To The Downloads Section

The Transformer Block: Combining MLA and MoE

Now that we have our sophisticated components, we need to integrate them into the classic transformer architecture. A DeepSeek transformer block follows the pre-norm residual pattern:

This is the pre-layer normalization (Pre-LN) variant used in modern transformers. The residual connections allow gradients to flow directly through the network during backpropagation, addressing the vanishing gradient problem in deep networks. The layer norms stabilize the activations before they enter the attention or MoE layers.

Mathematically, each block implements a function

where is the layer index. The full model is a composition of such blocks:

where denotes the sum of the token and position embeddings. Each layer refines the representations, with early layers typically learning syntax and surface patterns while deeper layers learn semantics and abstract concepts.

The Complete Model: From Tokens to Logits

The full DeepSeek model architecture (Figure 1) consists of:

• Token Embedding: maps token identifiers (IDs) to dense vectors
• Position Embedding: provides positional information
• Input Dropout: Regularization applied to embedded representations
• Transformer Blocks: blocks combining MLA and MoE with layer normalization
• Final Layer Norm: Stabilizes representations before output projection
• Language Modeling Head: Projects to the vocabulary:
• Multi-Token Prediction Heads: Optional heads for future token prediction

The forward pass computes:

where contains the token indices and denotes the position indices.

Weight Tying: Sharing Token Embeddings and Output Weights

An important optimization is weight tying between the token embedding matrix and the language modeling head:

This reduces parameters significantly (for our vocabulary of 50,259 and embedding dimension of 256, we save approximately 12.9 million parameters) and often improves performance. The intuition is elegant: if token has embedding vector , then the similarity measures how much the hidden representation aligns with token . This is exactly what we want for predicting that token.

Mathematically, weight tying implements a form of parameter sharing that encourages the input and output spaces to be aligned. The embedding space becomes jointly optimized for both encoding tokens into context and decoding context into tokens.

Multi-Token Prediction Integration

During training, after computing the main hidden states , we process them through MTP heads:

Each head takes the previous layer’s output and the embedding of the next actual token (ground truth during training), processes them through its mini-transformer, and projects the result to the vocabulary. We compute losses for all predictions and combine them:

where and are cross-entropy losses, are MTP weights, and is the MoE load-balancing loss with coefficient .

Implementing the DeepSeek-V3 Architecture in PyTorch

Let us implement the full architecture:

class DeepSeekBlock(nn.Module):
    """DeepSeek transformer block with MLA and MoE"""

    def __init__(self, config: DeepSeekConfig):
        super().__init__()
        self.config = config

        # Layer norms (pre-norm architecture)
        self.ln1 = nn.LayerNorm(config.n_embd, bias=config.bias)
        self.ln2 = nn.LayerNorm(config.n_embd, bias=config.bias)

        # Attention - MLA
        self.attn = MultiheadLatentAttention(config)

        # MoE feedforward
        self.moe = MixtureOfExperts(config)

Lines 1-16: Block Structure: The DeepSeekBlock class encapsulates a single transformer layer. We use 2 layer norms (ln1 and ln2) for the pre-norm architecture, which normalizes inputs before they enter the attention and feedforward (MoE) sublayers. This has proven more stable for training than post-norm architectures. The block contains our custom MLA mechanism and MoE feedforward layer, both of which we have already implemented.

    def forward(self, x: torch.Tensor, attention_mask: Optional[torch.Tensor] = None):
        # Attention with residual connection
        x = x + self.attn(self.ln1(x), attention_mask)

        # MoE with residual connection
        moe_output, router_logits  = self.moe(self.ln2(x))
        x = x + moe_output
        return x, router_logits

Lines 18-25: Forward Pass with Residual Connections: The forward method implements the classic “attention then feedforward” pattern with residual connections. First, we normalize the input, pass it through attention, and add it back to the original input (residual connection). Then we normalize again, pass through MoE, and add another residual. Importantly, we return both the output and router_logits from MoE. These router logits are needed for computing auxiliary losses during training. The residual connections are crucial: they allow gradients to flow directly through the network, helping mitigate vanishing gradients in deep models.

class DeepSeek(nn.Module):
    """Complete DeepSeek model for children's story generation"""

    def __init__(self, config: DeepSeekConfig):
        super().__init__()
        assert isinstance(config, DeepSeekConfig)
        self.config = config

        # Embeddings and transformer blocks
        self.transformer = nn.ModuleDict(dict(
            wte=nn.Embedding(config.vocab_size, config.n_embd),
            wpe=nn.Embedding(config.block_size, config.n_embd),
            drop=nn.Dropout(config.dropout),
            h=nn.ModuleList([DeepSeekBlock(config) for _ in range(config.n_layer)]),
            ln_f=nn.LayerNorm(config.n_embd, bias=config.bias),
        ))

        # Output heads
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)

        # Multi-Token Prediction heads
        if config.multi_token_predict > 0:
            self.mtp_heads = nn.ModuleList([
                MultiTokenPredictionHead(config, depth)
                for depth in range(1, config.multi_token_predict + 1)
            ])
        else:
            self.mtp_heads = None

        # Weight tying (share embeddings and output projection)
        self.transformer.wte.weight = self.lm_head.weight

        # Initialize weights
        self.apply(self._init_weights)

        # Special initialization for residual projections
        for pn, p in self.named_parameters():
            if pn.endswith(('o_proj.weight', 'down_proj.weight')):
                nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * config.n_layer))

Lines 27-42: Model Initialization: The DeepSeek class constructor sets up the complete model architecture. We use an nn.ModuleDict to organize the transformer components: wte for token embeddings, wpe for positional embeddings, drop for input dropout, h as a list of transformer blocks, and ln_f for the final layer norm. This organization makes the model structure clear and allows easy access to each component.

Lines 45-57: Output Heads and Weight Tying: We create the language modeling head (lm_head) that projects from the hidden dimension to the vocabulary size. If multi-token prediction is enabled, we create a list of MTP heads, one for each future depth. The crucial line self.transformer.wte.weight = self.lm_head.weight implements weight tying: the same weight matrix is used for both embedding tokens and predicting them. This reduces parameters and improves training.

Lines 62-65: Special Initialization for Residual Projections: The special initialization targets 2 types of output projections that feed into residual additions: o_proj in Multihead Latent Attention and down_proj in each SwiGLU expert. It scales their initialization by to account for the accumulation of residual connections through layers, helping prevent activation explosion.

def _init_weights(self, module):
        """Initialize model weights"""
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
        elif isinstance(module, nn.LayerNorm):
            nn.init.ones_(module.weight)
            if module.bias is not None:
                nn.init.zeros_(module.bias)

Lines 66-77: Weight Initialization: The _init_weights method implements careful initialization following best practices. Linear layers and embeddings use normal initialization with a standard deviation of 0.02. Layer normalization weights are initialized to 1, and their biases are initialized to 0 when present.

       def forward(self, input_ids: torch.Tensor, targets: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, **kwargs):
        """Forward pass with Multi-Token Prediction

        Main prediction: h[i] predicts token[i+1]
        MTP heads: Each depth d predicts token[i+d+1] using h[i] and embed[i+d]
        """
        device = input_ids.device
        batch_size, seq_len = input_ids.size()
        assert seq_len <= self.config.block_size

        # Get embeddings
        pos = torch.arange(0, seq_len, dtype=torch.long, device=device)
        tok_emb = self.transformer.wte(input_ids)
        pos_emb = self.transformer.wpe(pos)
        x = self.transformer.drop(tok_emb + pos_emb)
       
        # Forward through transformer blocks
        router_logits_list = []
        for block in self.transformer.h:
            x, router_logits = block(x, attention_mask=attention_mask)
            router_logits_list.append(router_logits)

        # Final layer norm
        x = self.transformer.ln_f(x)  # [B, seq_len, n_embd]
 

Lines 79-102: Forward Pass – Embeddings and Blocks: The forward method implements the complete forward pass. We first get token and position embeddings, sum them, and apply dropout. Then we iterate through all transformer blocks, collecting router logits from each MoE layer (needed for auxiliary losses). Finally, we apply the final layer norm. The hidden states x now encode the full context for each position. The forward signature also accepts an optional attention_mask of shape [B, T], which marks real tokens with 1 and padding with 0. Here, B is the batch size and T is the sequence length. We thread this mask down into every transformer block so that Multihead Latent Attention can exclude padded key positions from its attention scores.

        # Main language modeling head
        main_logits = self.lm_head(x)
        main_loss = None

        if targets is not None:
            # Compute main loss (standard next-token prediction)
            # Shift: h[i] predicts target[i+1]
            # main_logits[:, :-1, :] are predictions from h[0] to h[seq_len-2]
            # targets[:, 1:] are actual tokens at positions [1] to [seq_len-1]
            shift_logits = main_logits[:, :-1, :].contiguous()  # [B, seq_len-1, vocab_size]
            shift_targets = targets[:, 1:].contiguous()     # [B, seq_len-1]

            main_loss = F.cross_entropy(
                    shift_logits.view(-1, shift_logits.size(-1)),
                    shift_targets.view(-1),
                    ignore_index=-100
                )

Lines 105-120: Main Loss Computation: For training (when targets are provided), we compute the standard language modeling loss. We shift logits and targets by 1 position because we are predicting the next token: main_logits[:, i] should predict targets[:, i+1]. The cross-entropy loss with ignore_index=-100 excludes target positions labeled -100, including padding positions if they use that label. This is the base training objective.

             # Multi-Token Prediction
            mtp_loss = None
            if self.mtp_heads is not None:
                mtp_losses = []
                current_hidden = x

                for depth, mtp_head in enumerate(self.mtp_heads, 1):
                    # Check if we have enough future tokens
                    if seq_len > depth:
                        # Get future token embeddings
                        future_indices = input_ids[:, depth:]
                        future_embeds = self.transformer.wte(future_indices)

                        # Pad or truncate to match current_hidden sequence length
                        if future_embeds.size(1) < current_hidden.size(1):
                            pad_size = current_hidden.size(1) - future_embeds.size(1)
                            padding = torch.zeros(
                                batch_size, pad_size, self.config.n_embd,
                                device=device, dtype=future_embeds.dtype
                            )
                            future_embeds = torch.cat([future_embeds, padding], dim=1)
                        elif future_embeds.size(1) > current_hidden.size(1):
                            future_embeds = future_embeds[:, :current_hidden.size(1)]

                        # Process through MTP head
                        current_hidden = mtp_head(current_hidden, future_embeds, attention_mask=attention_mask)
                        mtp_logits = self.lm_head(current_hidden)

                        # Compute loss for this depth
                        # mtp_logits[:, i] predicts target[i + depth + 1]
                        if seq_len > depth + 1:
                            shift_logits = mtp_logits[:, :-(depth+1), :].contiguous()
                            shift_labels = targets[:, depth+1:].contiguous()

                            if shift_labels.numel() > 0:
                                mtp_loss_single = F.cross_entropy(
                                    shift_logits.view(-1, shift_logits.size(-1)),
                                    shift_labels.view(-1),
                                    ignore_index=-100
                                )
                                mtp_losses.append(mtp_loss_single)

                # Average MTP losses
                if mtp_losses:
                    mtp_loss = torch.stack(mtp_losses).mean()

Lines 123-166: Multi-Token Prediction: If MTP heads exist, we compute additional losses for future token prediction. For each depth, we get embeddings of the future tokens, process them through the MTP head combined with current hidden states, and compute predictions. The key insight: each MTP head receives both the context (via current_hidden) and embeddings of future ground-truth tokens (via future_embeds). We carefully handle sequence length issues by padding or truncating, and we pass the attention_mask into each MTP head, since every head runs its own mini-transformer (MLA and MoE) and must ignore padded positions for the same reason the main blocks do. All MTP losses are averaged and weighted by 0.3 relative to the main loss.

            # Add MoE auxiliary loss
            aux_loss = 0.0
            if router_logits_list:
                for i, block in enumerate(self.transformer.h):
                    aux_loss += block.moe._complementary_sequence_aux_loss(router_logits_list[i], seq_mask=attention_mask)
                aux_loss = aux_loss / len(router_logits_list)

            # Combine losses
            total_loss = main_loss
            if mtp_loss is not None:
                total_loss = total_loss + 0.3 * mtp_loss  # Weight MTP loss
           
            total_loss = total_loss + self.config.aux_loss_weight * aux_loss

            return main_logits, total_loss
        else:
            # Inference mode
            logits = self.lm_head(x[:, [-1], :])
            return logits, None

Lines 167-181: Auxiliary Loss and Combination: We compute the complementary sequence-wise auxiliary loss from all MoE layers and average across layers. This encourages load balancing among experts. We pass seq_mask=attention_mask so that padded positions are excluded when measuring expert load; without it, padding would be counted as real tokens and would skew the balancing signal toward whichever experts happen to absorb it. The total loss combines the main prediction loss, MTP loss (if enabled), and auxiliary loss (weighted by self.config.aux_loss_weight). This multi-objective training improves model quality while maintaining computational efficiency.

Lines 184 and 185: Inference Mode: When no targets are provided (inference), we simply compute logits for the last position and return them. This is used during generation when we predict 1 token at a time. No MTP heads are used during inference. They have already served their purpose in improving training.

    @torch.no_grad()
    def generate(self, input_ids: torch.Tensor, max_new_tokens: int = 100,
                 temperature: float = 1.0, top_k: Optional[int] = None):
        """Generate text autoregressively"""
        for _ in range(max_new_tokens):
            # Crop to context window
            idx_cond = input_ids if input_ids.size(1) <= self.config.block_size else input_ids[:, -self.config.block_size:]

            # Forward pass (no targets, so inference mode)
            logits, _ = self(idx_cond)  # [B, 1, vocab_size]
            logits = logits[:, -1, :] / temperature  # [B, vocab_size]

            # Apply top-k filtering
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = -float('Inf')

            # Sample next token
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)  # [B, 1]
            input_ids = torch.cat((input_ids, idx_next), dim=1)

        return input_ids

Lines 186-208: Autoregressive Generation: The generate method implements text generation. For each new token, we crop the input to the context window (if it has grown too large), compute logits, apply temperature scaling and top-k filtering, sample from the distribution, and append to the input. This continues for max_new_tokens iterations. Temperature controls randomness (lower values make sampling more deterministic), while top-k prevents sampling very unlikely tokens that might break coherence.

Architectural Design Patterns

Several design patterns in our architecture are worth highlighting:

Pre-Norm vs Post-Norm: We use pre-normalization (normalization before the sublayer) rather than post-normalization (normalization after the sublayer). Research has shown pre-norm is more stable for training deep networks, though post-norm can sometimes achieve slightly better performance if training succeeds. The stability-performance tradeoff generally favors pre-norm for modern large models.

Residual Connection Placement: Every sublayer (attention, MoE) has a residual connection. This is crucial for gradient flow in deep networks. Without residuals, gradients would have to flow through many sequential transformations, leading to vanishing or exploding gradients. With residuals, gradients have a direct path to earlier layers.

Dropout Placement: We apply dropout in 3 places: input embeddings, within attention and MoE sublayers (implemented in those modules), and in MTP heads. This multi-level dropout provides regularization at different stages of computation. Too much dropout can hurt performance; too little can cause overfitting. Our 0.1 rate is moderate.

Weight Tying: Sharing weights between embeddings and output projection is an elegant constraint that reduces parameters and often helps. The theoretical justification is that both mappings operate in the same semantic space: the space of token meanings.

Parameter Count Analysis

Let us analyze where parameters are allocated:

• Embeddings: parameters
• Position Embeddings: parameters
• Per Layer (MLA and MoE): Approximately 2-3 million parameters
• Total (6 layers): Approximately 30-35 million parameters
• MTP Heads: An additional 2-3 million parameters per head

Most parameters are in the embeddings and the transformer blocks. The MLA compression reduces attention parameters compared to standard transformers, while MoE increases feedforward parameters. The balance results in a model that is larger than a standard transformer of the same compute cost, but not proportionally to the number of experts.

Memory and Computation Footprint

During training, memory usage includes:

• Model Parameters: Approximately 30-35 million parameters at 4 bytes each require 120-140 megabytes (MB) in 32-bit floating-point format (FP32)
• Gradients: Require the same memory as the parameters: 120-140 MB
• Optimizer States: Require twice the parameter memory with AdamW: 240-280 MB
• Activations: Memory usage depends on batch size and sequence length
• Key-Value (KV) Cache: Not needed during training

Total training memory is roughly 500-600 MB plus activations. For a batch size of 4 and sequence length of 1024, activations add approximately 200-300 MB, giving us approximately 800 MB total. This fits comfortably on modern graphics processing units (GPUs).

During inference, we do not need gradients or optimizer states, and we can use 16-bit floating-point format (FP16), roughly halving memory. This implementation does not use a KV cache; it recomputes the context during each generation step.

With all components assembled, we have a complete, working DeepSeek-V3 model. It combines 4 innovations into a coherent architecture: configuration management with RoPE, MLA for efficient attention, MoE for sparse scaling, and MTP for richer training. In the next lesson, we will train this model on real data and see it generate text.

What's next? We recommend PyImageSearch University.

Course information:
120+ total classes • 115+ hours of on-demand code walkthrough videos • Last updated: October 2026
★★★★★ 4.84 (128 Ratings) • 16,000+ Students Enrolled

I strongly believe that if you had the right teacher you could master computer vision and deep learning.

Do you think learning computer vision and deep learning has to be time-consuming, overwhelming, and complicated? Or has to involve complex mathematics and equations? Or requires a degree in computer science?

That’s not the case.

All you need to master computer vision and deep learning is for someone to explain things to you in simple, intuitive terms. And that’s exactly what I do. My mission is to change education and how complex Artificial Intelligence topics are taught.

If you're serious about learning computer vision, your next stop should be PyImageSearch University, the most comprehensive computer vision, deep learning, and OpenCV course online today. Here you’ll learn how to successfully and confidently apply computer vision to your work, research, and projects. Join me in computer vision mastery.

Inside PyImageSearch University you'll find:

• ✓ 120+ courses on essential computer vision, deep learning, and OpenCV topics
• ✓ 94+ Certificates of Completion
• ✓ 115+ hours of on-demand video
• ✓ Brand new courses released regularly, ensuring you can keep up with state-of-the-art techniques
• ✓ Pre-configured Jupyter Notebooks in Google Colab
• ✓ Run all code examples in your web browser — works on Windows, macOS, and Linux (no dev environment configuration required!)
• ✓ Access to centralized code repos for all 540+ tutorials on PyImageSearch
• ✓ Easy one-click downloads for code, datasets, pre-trained models, etc.
• ✓ Access on mobile, laptop, desktop, etc.

Click here to join PyImageSearch University

Summary

In this 5th lesson of our DeepSeek-V3 From Scratch series, we bring everything together by assembling the full DeepSeek-V3 architecture. We start with an overview of how the different components (MLA, MoE, and RoPE) fit into the broader design, and then move into the construction of the transformer block, where Multihead Latent Attention and Mixture of Experts are combined to form the model’s core computational unit. This sets the stage for understanding how the complete model processes information, from raw tokens all the way to logits.

We then explore key architectural innovations such as weight tying, which shares weights between the token embedding layer and the language modeling head, and multi-token prediction integration, which improves efficiency and predictive power. These design choices are not just theoretical. They directly impact how the model learns and generalizes. The implementation section walks us through building the complete DeepSeek model step by step, showing how each piece connects seamlessly into a unified system.

Finally, we analyze the architectural design patterns, parameter counts, and the memory and computation footprint of DeepSeek-V3. This helps us evaluate trade-offs between scalability and efficiency, and understand how the model balances complexity with performance. By the end of this lesson, we have not only assembled the architecture but also gained insight into the design decisions that make DeepSeek-V3 both powerful and practical.

Citation Information

Mangla, P. “DeepSeek-V3 from Scratch: Building the Architecture in PyTorch,” PyImageSearch, S. Huot, A. Sharma, and P. Thakur, eds., 2026, https://pyimg.co/36gvk

@incollection{Mangla_2026_deepseek-v3-from-scratch-building-architecture-in-pytorch,
  author = {Puneet Mangla},
  title = {{DeepSeek-V3 from Scratch: Building the Architecture in PyTorch}},
  booktitle = {PyImageSearch},
  editor = {Susan Huot and Aditya Sharma and Piyush Thakur},
  year = {2026},
  url = {https://pyimg.co/36gvk},
}

To download the source code to this post (and be notified when future tutorials are published here on PyImageSearch), simply enter your email address in the form below!

Download the Source Code and FREE 17-page Resource Guide

Enter your email address below to get a .zip of the code and a FREE 17-page Resource Guide on Computer Vision, OpenCV, and Deep Learning. Inside you'll find my hand-picked tutorials, books, courses, and libraries to help you master CV and DL!

Download the code!

Website

The post DeepSeek-V3 from Scratch: Building the Architecture in PyTorch appeared first on PyImageSearch.

DVC Pipelines for MLOps: Build a Reproducible ML Pipeline

In this lesson you will learn how to build a fully reproducible machine learning pipeline using DVC, starting from raw data and moving through preprocessing, model training, and evaluation, all connected inside a clean and automated workflow that you can run with a single command.

This lesson is the last in a 2-part series on Data and Model Versioning with DVC:

1. DVC for MLOps: Versioning Your Data and Models the Right Way
2. DVC Pipelines for MLOps: Build a Reproducible ML Pipeline (this tutorial)

To learn how to design, run, and reproduce a multi-stage pipeline using DVC’s powerful workflow engine, just keep reading.

Looking for the source code to this post?

Jump Right To The Downloads Section

Introduction

What We Did in Lesson 1

In the previous lesson, we learned how to use DVC to version datasets and model artifacts without polluting Git with large files.

You tracked a raw IMDb dataset, generated a dummy checkpoint, and pushed everything to a remote using DVC metadata files instead of storing actual data in Git.

We also covered how DVC caches files, how .dvc metadata files point to real data in the cache, and why this workflow makes ML projects easier to share, update, and reproduce.

By the end of Lesson 1, you had a clean, Git-friendly approach to managing raw data and model outputs across a team.

Why Pipelines Matter in MLOps

Versioning data is useful, but real machine learning (ML) systems need reproducibility across entire workflows, not just individual files.

Once you have preprocessing scripts, training code, and evaluation logic, you need a way to connect them so they run in the right order and re-run only when necessary.

This is where DVC pipelines come in.

Pipelines let you define each step of your ML workflow as a stage with its own inputs, outputs, and logic, so DVC can trace dependencies and rebuild only the stages affected by changes.

In production teams, pipelines enforce repeatability, traceability, and automation. This structure allows any teammate or continuous integration (CI) job to rebuild your results exactly, even months later.

Put simply, pipelines transform ad hoc scripts into deterministic ML systems.

What We Will Build in This Lesson (3-Stage Iris Pipeline)

In this lesson, you will build a fully reproducible, 3-stage ML pipeline for the classic Iris dataset.

Each stage will be defined in dvc.yaml and orchestrated with dvc repro, making it easy to rebuild the workflow from scratch.

Here is the exact pipeline we will construct:

• Preprocess: load iris.csv, split into train/test, and save processed data
• Train: train a logistic regression classifier using the processed training data
• Evaluate: compute accuracy, generate a classification report, and save metrics and a human-readable report

By the end of this lesson, you will have a production-style pipeline where any team member (or CI server) can:

git clone <repo>
pip install -r requirements.txt
dvc pull
dvc repro

You can then reproduce the outputs, metrics, and model without running each stage manually.

Would you like immediate access to 3,457 images curated and labeled with hand gestures to train, explore, and experiment with … for free? Head over to Roboflow and get a free account to grab these hand gesture images.

Project Setup

Reviewing the Lesson 2 Repository Structure

Before writing any pipeline logic, let us quickly walk through the repository so you know where every component lives.

This project uses a simple, production-friendly layout that separates raw data, scripts, outputs, and reports.

dvc-lesson2/
├── data/
│   └── raw/
│       └── iris.csv                  # Raw dataset
├── src/
│   ├── preprocess.py                 # Stage 1: Preprocessing
│   ├── train.py                      # Stage 2: Training
│   └── evaluate.py                   # Stage 3: Evaluation
├── outputs/
│   ├── processed/                    # Stage 1 output
│   ├── model/                        # Stage 2 output
│   └── metrics.json                  # Stage 3 output
├── reports/
│   └── accuracy.txt                  # Human-readable evaluation report
├── dvc.yaml                          # Pipeline definition
├── requirements.txt                  # Python dependencies
├── .gitignore                        # Files ignored by Git
└── .dvcignore                        # Files ignored by DVC

The structure mirrors real ML projects. Every stage writes outputs that become inputs for the next stage, and DVC tracks these relationships automatically through dvc.yaml.

Installing Dependencies

This lesson uses only 3 core packages: DVC, pandas, and scikit-learn.

Everything is lightweight so you can focus on learning pipelines, not juggling toolchains.

Install dependencies:

pip install -r requirements.txt

requirements.txt contains:

dvc==3.48.4
pandas==2.1.4
scikit-learn==1.3.2

Once installed, you have everything needed to preprocess data, train a model, evaluate it, and let DVC orchestrate the entire workflow.

Initializing Git and DVC

Just like Lesson 1, the pipeline project begins with Git and DVC initialization.

If this repository is fresh on your machine, run:

git init

Now initialize DVC:

dvc init

This will create:

• .dvc/: internal DVC configuration
• .dvcignore: paths ignored by DVC
• Git hooks for DVC tracking

Commit the setup:

git add .dvc .dvcignore
git commit -m "Initialize DVC for pipeline project"

At this point, the repository is ready for multi-stage pipelines, and DVC is prepared to track every dependency, output, and command we define in dvc.yaml.

Need Help Configuring Your Development Environment?

All that said, are you:

• Short on time?
• Learning on your employer’s administratively locked system?
• Wanting to skip the hassle of fighting with the command line, package managers, and virtual environments?
• Ready to run the code immediately on your Windows, macOS, or Linux system?

Then join PyImageSearch University today!

Gain access to Jupyter Notebooks for this tutorial and other PyImageSearch guides pre-configured to run on Google Colab’s ecosystem right in your web browser! No installation required.

And best of all, these Jupyter Notebooks will run on Windows, macOS, and Linux!

Understanding the Pipeline Architecture

A DVC pipeline is simply a series of stages that depend on each other. Each stage has a script, inputs, outputs, and a command. DVC connects those stages into a reproducible workflow that can be re-run automatically whenever something changes.

Our project implements a classic 3-stage ML workflow: preprocess → train → evaluate.

Let us break down how this works.

The 3-Stage Iris Pipeline

The Iris workflow in this project consists of 3 connected stages, each defined in dvc.yaml.

Stage 1: preprocess

Script: src/preprocess.py

• Loads the raw iris.csv
• Splits it into train/test sets
• Saves cleaned data to outputs/processed/

data/raw/iris.csv → outputs/processed/

Stage 2: train

Script: src/train.py

• Reads processed training data
• Fits a logistic regression model
• Saves the model and label encoder to outputs/model/

outputs/processed/ → outputs/model/

Stage 3: evaluate

Script: src/evaluate.py

• Loads the trained model and test data
• Computes metrics (accuracy and classification report)
• Saves results to outputs/metrics.json and reports/accuracy.txt

outputs/model/ → outputs/metrics.json, reports/accuracy.txt

These 3 stages form a clean, linear ML pipeline. Data flows from the raw dataset to the processed dataset, model, and evaluation metrics.

How DVC Uses dvc.yaml

DVC connects all 3 stages through a single file: dvc.yaml.

This file tells DVC:

• What command to run (cmd)
• Which files a stage depends on (deps)
• What outputs it produces (outs)

Your project’s dvc.yaml looks like this:

stages:
  preprocess:
    cmd: python src/preprocess.py
    deps:
      - data/raw/iris.csv
      - src/preprocess.py
    outs:
      - outputs/processed/

  train:
    cmd: python src/train.py
    deps:
      - outputs/processed/
      - src/train.py
    outs:
      - outputs/model/

  evaluate:
    cmd: python src/evaluate.py
    deps:
      - outputs/model/
      - src/evaluate.py
    outs:
      - outputs/metrics.json
      - reports/accuracy.txt

In one place, DVC now knows:

• Changing iris.csv forces the entire pipeline to re-run
• Changing preprocess.py forces preprocess → train → evaluate to re-run
• Changing train.py forces only train → evaluate to re-run
• Changing evaluate.py forces only evaluate to re-run

This is what makes ML pipelines reproducible and efficient.

Reading the Directed Acyclic Graph

DVC can visualize the workflow as a Directed Acyclic Graph (DAG):

dvc dag

You will see an ASCII diagram like this:

          +-------------+
          | data/raw/... |
          +-------------+
                 *
                 *
                 *
          +-------------+
          | preprocess  |
          +-------------+
                 *
                 *
                 *
             +--------+
             | train  |
             +--------+
                 *
                 *
                 *
           +-------------+
           |  evaluate   |
           +-------------+

This DAG tells you:

• The pipeline begins with iris.csv
• Every stage depends on the previous one
• Training will not run until preprocessing finishes
• Evaluation will not run until a model exists

It is a perfect, linear end-to-end representation of your ML system.

Stage 1: Dataset Preprocessing

The first stage in our pipeline handles dataset preparation. This includes loading the raw Iris comma-separated values (CSV) file, splitting it into training and test sets, and writing clean structured files into a dedicated output directory. DVC uses this stage as the foundation for all downstream work. If preprocessing changes, every downstream stage automatically reruns.

Reviewing preprocess.py

Your preprocessing script is located at:

src/preprocess.py

Here is what it does, step by step:

df = pd.read_csv("data/raw/iris.csv")

It loads the raw Iris dataset from data/raw/iris.csv.

X = df.drop("species", axis=1)
y = df["species"]

It separates features and target labels.

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)

It creates an 80/20 train-test split while preserving class proportions (stratify=y).

output_dir = "outputs/processed"
os.makedirs(output_dir, exist_ok=True)

It prepares the output directory where DVC will store processed datasets.

train_df.to_csv(os.path.join(output_dir, "train.csv"), index=False)
test_df.to_csv(os.path.join(output_dir, "test.csv"), index=False)

Finally, the preprocessed data is saved as:

• outputs/processed/train.csv
• outputs/processed/test.csv

These are the cleaned, structured files the rest of the ML pipeline depends on.

Inputs and Outputs

Inputs (DVC Dependencies)

Listed under deps in dvc.yaml:

deps:
  - data/raw/iris.csv
  - src/preprocess.py

Meaning:

• Changing the raw dataset reruns Stage 1
• Changing the script itself reruns Stage 1

Outputs (DVC Tracked Artifacts)

Listed under outs:

outs:
  - outputs/processed/

This directory contains 2 files:

• train.csv
• test.csv

DVC tracks the entire directory as a single output artifact.

How DVC Tracks Directory Outputs (outputs/processed/)

DVC treats outputs/processed/ as one logical output, not just a folder.

When you run:

dvc repro

DVC:

• Computes hashes of the directory contents
• Stores processed files in .dvc/cache/
• Creates a reference entry in dvc.lock
• Recreates (or updates) the directory when needed

DVC does not version each individual CSV separately. Instead, it versions the directory state, making preprocessing reproducible and lightweight.

If you open dvc.lock, you will see a section such as the following:

outs:
  - path: outputs/processed
    md5: <hash>
    size: <bytes>

This hash represents the exact combination of train.csv and test.csv.

What Changes Trigger Re-Runs

DVC re-runs the preprocessing stage only when necessary.

The following changes trigger Stage 1 to run again:

Modify the raw dataset

Example: editing data/raw/iris.csv

Modify the preprocessing script

Example: changing the test split from 0.2 to 0.3:

test_size=0.3

This triggers:

preprocess → train → evaluate

because downstream stages depend on processed data.

Running dvc repro again without changes

DVC outputs:

Stage 'preprocess' didn't change, skipping

This shows why DVC pipelines are efficient: no unnecessary computation.

Stage 2: Model Training and Versioning with DVC

The training stage consumes the preprocessed dataset from Stage 1, fits a model, and writes all model artifacts into a DVC-tracked output directory. This stage represents the “learning” portion of the pipeline, and any upstream change automatically forces retraining.

Reviewing train.py

Your training script lives at:

src/train.py

Here is what it does:

Load processed training data

train_path = "outputs/processed/train.csv"
train_df = pd.read_csv(train_path)

It reads the output from Stage 1, meaning preprocessing must complete successfully before training can run.

Split Features and Labels

X_train = train_df.drop("species", axis=1)
y_train = train_df["species"]

Encode Target Labels

label_encoder = LabelEncoder()
y_train_encoded = label_encoder.fit_transform(y_train)

The Iris dataset uses string labels (setosa, versicolor, and virginica).

The model expects numeric classes, so they are encoded.

Train a Logistic Regression Model

model = LogisticRegression(
    max_iter=200,
    random_state=42,
    solver='lbfgs',
    multi_class='multinomial'
)
model.fit(X_train, y_train_encoded)

This is a simple, lightweight model suitable for demonstrating pipelines, dependencies, and reproducibility.

Report Training Accuracy

train_accuracy = model.score(X_train, y_train_encoded)
print(f"📈 Training accuracy: {train_accuracy:.4f}")

This provides quick feedback but is not saved as a metric because evaluation happens later.

How It Loads Stage 1 Output

The key detail in this stage is the direct dependency on Stage 1’s output directory:

train_path = "outputs/processed/train.csv"

That file is generated by the preprocessing step.

In dvc.yaml:

deps:
  - outputs/processed/

This means:

• If any file inside outputs/processed/ changes
• If the preprocessing script changes
• If the raw data changes

Stage 2 then automatically reruns when you run the following command:

dvc repro

This ensures reproducibility and tightly coupled data lineage.

Saving the Model and Label Encoder

After training, the script saves 2 artifacts:

• The trained model
• The label encoder

This is important because predictions on new data would not map to the correct class names without the encoder.

Create output directory

output_dir = "outputs/model"
os.makedirs(output_dir, exist_ok=True)

Write the model and encoder

model_path = os.path.join(output_dir, "model.pkl")
encoder_path = os.path.join(output_dir, "label_encoder.pkl")

with open(model_path, "wb") as f:
    pickle.dump(model, f)

with open(encoder_path, "wb") as f:
    pickle.dump(label_encoder, f)

These artifacts drive the evaluation stage and are vital for any downstream inference.

Why Model Outputs Are Tracked as DVC Artifacts

In dvc.yaml, Stage 2 defines:

outs:
  - outputs/model/

This ensures DVC:

• Versions your model files deterministically
• Treats the entire outputs/model/ directory as an artifact
• Stores the model in .dvc/cache/ using content hashing
• Adds reproducibility guarantees (via dvc.lock)
• Re-runs Stage 3 only if the model changes

DVC tracks only metadata in Git rather than the .pkl files. This approach keeps your repository lightweight while preserving reproducibility.

This setup also mirrors real production workflows:

• Lightning or PyTorch generates checkpoints
• DVC versions them
• Teams share them without pushing binaries to Git

Stage 3: Model Evaluation and Metrics Tracking with DVC

The evaluation stage is the final step in the pipeline. It loads the trained model and label encoder from Stage 2, evaluates performance on the test dataset from Stage 1, generates structured metrics, produces a human-readable report, and writes both outputs to DVC-tracked locations.

This stage demonstrates how DVC can treat metrics as first-class citizens in a reproducible ML workflow.

Reviewing evaluate.py

This script lives at:

src/evaluate.py

Here is a walkthrough of what it does.

Load Test Data

test_path = "outputs/processed/test.csv"
test_df = pd.read_csv(test_path)

Just as training depends on processed training data, evaluation depends on processed test data from Stage 1.

Split Features and Labels

X_test = test_df.drop("species", axis=1)
y_test = test_df["species"]

Load the Trained Model and Label Encoder

model_path = "outputs/model/model.pkl"
encoder_path = "outputs/model/label_encoder.pkl"

with open(model_path, "rb") as f:
    model = pickle.load(f)

with open(encoder_path, "rb") as f:
    label_encoder = pickle.load(f)

This is where pipeline dependency chaining becomes visible:

• Stage 1: generates processed data
• Stage 2: generates the model and encoder
• Stage 3: loads output from both

Encode Labels and Run Predictions

y_test_encoded = label_encoder.transform(y_test)
y_pred = model.predict(X_test)

Compute Primary Metric (Accuracy)

accuracy = accuracy_score(y_test_encoded, y_pred)
print(f"📈 Test Accuracy: {accuracy:.4f}")

The script also prints a classification report and confusion matrix to help users understand model behavior.

Create Output Directories

os.makedirs("reports", exist_ok=True)

This sets the stage for saving the structured and human-readable results.

Generating Metrics (metrics.json)

The evaluation script produces a machine-readable metrics file:

metrics = {
    "accuracy": float(accuracy),
    "test_samples": len(X_test),
    "classification_report": report
}

metrics_path = "outputs/metrics.json"
with open(metrics_path, "w") as f:
    json.dump(metrics, f, indent=2)

Why this matters

• DVC can compare metrics across runs (e.g., before and after hyperparameter tuning)
• Continuous integration and continuous delivery (CI/CD) pipelines can read this JavaScript Object Notation (JSON) file automatically
• You can integrate these metrics into dashboards or experiment tracking tools

In dvc.yaml, this file appears under:

outs:
  - outputs/metrics.json

(You may optionally move this to the metrics: section for metric-specific behavior.)

Generating Human-Readable Reports (accuracy.txt)

While metrics.json is for machines, humans often want plain text.

The script generates:

accuracy_report_path = "reports/accuracy.txt"
with open(accuracy_report_path, "w") as f:
    f.write("Model Evaluation Report\n")
    f.write(f"{'='*50}\n\n")
    f.write(f"Test Accuracy: {accuracy:.4f}\n")
    f.write(f"Test Samples: {len(X_test)}\n\n")
    f.write(f"Classification Report:\n")
    f.write(f"{'-'*50}\n")
    f.write(classification_report(...))
    f.write("\nConfusion Matrix:\n")
    f.write(f"{'-'*50}\n")
    f.write(f"{cm}\n")

This produces:

• Header
• Accuracy summary
• Classification report
• Confusion matrix

It stores the report in the following location:

reports/accuracy.txt

Because this is listed under outs in dvc.yaml, DVC automatically tracks it.

Marking Metrics and Plots in DVC (optional but recommended)

DVC has special sections for:

• metrics: numeric evaluation values
• plots: visualizations such as confusion matrices

You can enhance the evaluate stage with the following configuration:

Example configuration

stages:
  evaluate:
    cmd: python src/evaluate.py
    deps:
      - outputs/model/
      - src/evaluate.py
    metrics:
      - outputs/metrics.json:
          cache: false
    outs:
      - reports/accuracy.txt

Benefits

• dvc metrics show: displays the current metrics
• dvc metrics diff: compares accuracy between commits
• cache: false: metrics.json is not stored in the DVC cache because metrics files are lightweight and can be tracked directly by Git

Optional: Track Plots

If you ever generate confusion_matrix.png or similar:

plots:
  - outputs/confusion_matrix.png

Then run:

dvc plots show

Running and Reproducing the Pipeline

Once your pipeline is defined in dvc.yaml, everything else becomes beautifully simple.

Instead of manually running preprocess.py, train.py, and evaluate.py, you can let DVC orchestrate the entire workflow. DVC reruns only the stages that need to run.

This is where DVC starts feeling like magic.

Running the Full Pipeline (dvc repro)

Your pipeline has 3 stages:

preprocess → train → evaluate

To run all of them in the correct order, execute:

dvc repro

The first run always executes all stages because no cached outputs exist yet.

Expected output

Running stage 'preprocess':
> python src/preprocess.py
📊 Starting data preprocessing...
...
💾 Saved train data to: outputs/processed/train.csv
💾 Saved test data to: outputs/processed/test.csv
✅ Preprocessing complete!

Running stage 'train':
> python src/train.py
🚀 Starting model training...
...
💾 Saved model to: outputs/model/model.pkl
💾 Saved label encoder to: outputs/model/label_encoder.pkl
✅ Training complete!

Running stage 'evaluate':
> python src/evaluate.py
📊 Starting model evaluation...
...
📈 Test Accuracy: 0.9333
📄 Saved report to: reports/accuracy.txt
✅ Evaluation complete!

Behind the scenes, DVC is:

• hashing each dependency
• hashing each output
• saving everything in .dvc/cache/
• writing a reproducible lock file

That brings us to the next part.

Understanding dvc.lock

After your first dvc repro run, DVC generates:

dvc.lock

This file records the exact state of the pipeline when it ran, similar to package-lock.json or poetry.lock.

Open it, and you will see entries such as the following:

stages:
  preprocess:
    cmd: python src/preprocess.py
    deps:
      - path: data/raw/iris.csv
        md5: <hash>
      - path: src/preprocess.py
        md5: <hash>
    outs:
      - path: outputs/processed
        md5: <hash>
        size: <int>

  train:
    cmd: python src/train.py
    deps:
      - path: outputs/processed
        md5: <hash>
      - path: src/train.py
        md5: <hash>
    outs:
      - path: outputs/model
        md5: <hash>

  evaluate:
    cmd: python src/evaluate.py
    deps:
      - path: outputs/model
        md5: <hash>
      - path: src/evaluate.py
        md5: <hash>
    outs:
      - path: outputs/metrics.json
      - path: reports/accuracy.txt

Why This File Matters

• It guarantees reproducibility across machines.
• It tells DVC which outputs match which inputs.
• It allows DVC to detect when nothing changed.
• It must be committed to Git.

Anyone pulling your repository can run:

dvc repro

They can then reproduce the pipeline.

Incremental Re-Runs (Partial Pipeline Rebuild)

DVC’s biggest advantage is incrementality.

It runs only the stages affected by changes.

Let us walk through real examples.

Example 1: You Change src/preprocess.py

# Change test split from 0.2 → 0.3
dvc repro

DVC will detect the following changes:

• preprocess changed: must run
• outputs/processed/ changed: train must run
• model changed: evaluate must run

Console output:

Running stage 'preprocess' (changed deps):
Running stage 'train' (changed deps):
Running stage 'evaluate' (changed deps):

Example 2: You Change Only src/train.py

dvc repro

DVC will detect the following:

• preprocess unchanged: skip
• train changed: run
• model changed: evaluate must run

Console output:

Stage 'preprocess' didn't change, skipping
Running stage 'train'
Running stage 'evaluate'

Example 3: Nothing Changed

dvc repro

Output:

Stage 'preprocess' didn't change, skipping
Stage 'train' didn't change, skipping
Stage 'evaluate' didn't change, skipping
Data and pipelines are up to date.

This means DVC checked all dependencies, hashes, and outputs. Everything matches the lock file.

Zero wasted compute. Zero unnecessary runs.

Visualizing and Inspecting Pipelines

Once your pipeline is running, you will want ways to inspect, debug, understand, and compare different runs.

DVC includes a suite of visualization and inspection tools that help you answer questions such as the following:

• What stages are in my pipeline?
• Which files changed since the last run?
• Do I need to re-run anything?
• What metrics improved or regressed?

This section walks you through the most useful tools for pipeline introspection.

Visualizing the Pipeline DAG (dvc dag)

The fastest way to understand your ML workflow is to visualize it as a DAG (Directed Acyclic Graph).

DVC generates a clean ASCII graph of your pipeline:

dvc dag

Expected output for your exact pipeline:

        +----------------------+
        |   data/raw/iris.csv |
        +----------------------+
                    *
                    *
             +------------+
             | preprocess |
             +------------+
                    *
                    *
                +-------+
                | train |
                +-------+
                    *
                    *
              +----------+
              | evaluate |
              +----------+

This confirms the 3-stage flow:

raw data → preprocess → train → evaluate

Use cases:

• Quickly verify pipeline structure
• Check whether a stage is connected properly
• Ensure there are no missing dependencies

Bonus: Mermaid DAG Output

For documentation, slides, or PyImageSearch blog posts:

dvc dag --mermaid

Produces a Mermaid flowchart:

flowchart TD
    preprocess --> train
    train --> evaluate

Paste this output into a Markdown renderer that supports Mermaid to produce a visual diagram.

Checking the Pipeline Status (dvc status)

dvc status tells you whether any dependency or output has changed since the last dvc repro run.

Run it anytime:

dvc status

If nothing changed:

Data and pipelines are up to date.

If something changed (e.g., you modified src/preprocess.py):

preprocess:
    changed deps:
        modified: src/preprocess.py

This means the following:

• preprocess must rerun
• which triggers reruns for downstream stages (train and evaluate)

If the dataset changed:

preprocess:
    changed deps:
        modified: data/raw/iris.csv

Use this before running large pipelines; it tells you what will execute.

Comparing Pipeline Versions (dvc diff)

dvc diff helps you inspect changes between 2 Git commits or between the current workspace and a Git commit.

It is especially useful when you are debugging why metrics changed.

To compare your current workspace to the last committed run:

dvc diff

Example 1: Preprocessing Logic Changed

Modified outs:
  outputs/processed:
    size: 4.2KB -> 4.5KB

Example 2: Model Weights Changed (e.g., parameter update)

Modified outs:
  outputs/model:
    size: 12KB -> 15KB

Example 3: Metrics Improved or Regressed

Modified outs:
  outputs/metrics.json:
    diff:
      accuracy: 0.9333 -> 0.9666

This is the easiest way to track performance shifts after code or data changes.

Inspecting Metrics (dvc metrics show, metrics diff)

In your pipeline:

• outputs/metrics.json: contains structured evaluation metrics
• reports/accuracy.txt: contains the human-readable report

DVC can read JSON metrics directly.

Show metrics from the latest run

dvc metrics show

Expected output:

Path: outputs/metrics.json
accuracy: 0.9333
test_samples: 30

Compare Metrics Between Commits

dvc metrics diff

Example output:

Path: outputs/metrics.json
metric: accuracy
    old: 0.9333
    new: 0.9666
    diff: +0.0333

This is extremely powerful during model development:

• See whether your model improved
• Detect regressions immediately
• Track hyperparameter sensitivity
• Log performance across experiments

Note: Your pipeline currently lists outputs/metrics.json as a regular output in dvc.yaml. You can move it to the metrics section to mark it explicitly as a metrics file.

If you want to tag them explicitly:

stages:
  evaluate:
    metrics:
      - outputs/metrics.json:
          cache: false

Extending the Pipeline (Optional Enhancements)

Your current DVC pipeline already handles a clean 3-stage workflow: preprocess, train, and evaluate.

However, real-world ML pipelines rarely stay this small. As projects grow, you will need hyperparameter tuning, metrics tracking, and new stages like feature engineering or postprocessing.

This section shows how to evolve your pipeline without breaking its reproducibility, using DVC’s extensibility features.

Using params.yaml for Hyperparameters

Right now, your training script (src/train.py) has hyperparameters hardcoded:

model = LogisticRegression(
    max_iter=200,
    random_state=42,
    solver='lbfgs',
    multi_class='multinomial'
)

To make these configurable through DVC, create a params.yaml file:

train:
  max_iter: 200
  solver: lbfgs
  multi_class: multinomial
  random_state: 42

split:
  test_size: 0.2
  random_state: 42

Step 1: Modify train.py to read params

import yaml

params = yaml.safe_load(open("params.yaml"))

cfg = params["train"]

model = LogisticRegression(
    max_iter=cfg["max_iter"],
    random_state=cfg["random_state"],
    solver=cfg["solver"],
    multi_class=cfg["multi_class"]
)

Step 2: Update dvc.yaml to declare params

stages:
  train:
    cmd: python src/train.py
    deps:
      - outputs/processed/
      - src/train.py
    params:
      - train.max_iter
      - train.solver
      - train.multi_class
      - train.random_state
    outs:
      - outputs/model/

Result:

• Changing any value in params.yaml triggers only the train and evaluate stages.
• You get clean, reproducible hyperparameter experiments.

Adding Metrics and Plots Tracking

Your pipeline already produces:

• outputs/metrics.json: machine-readable metrics
• reports/accuracy.txt: human-readable report

But DVC can track these outputs as first-class metrics.

Mark metrics in dvc.yaml

stages:
  evaluate:
    cmd: python src/evaluate.py
    deps:
      - outputs/model/
      - src/evaluate.py
    metrics:
      - outputs/metrics.json:
          cache: false
    outs:
      - reports/accuracy.txt

Now try:

dvc metrics show
dvc metrics diff

DVC displays accuracy differences across commits, which makes the command useful for comparing experiments.

Optional: Track Plots

If you add visualization code to evaluate.py:

import matplotlib.pyplot as plt
from sklearn.metrics import ConfusionMatrixDisplay

disp = ConfusionMatrixDisplay(confusion_matrix=cm)
disp.plot()
plt.savefig("outputs/confusion_matrix.png")

Update dvc.yaml:

plots:
  - outputs/confusion_matrix.png

Now you can compare model performance visually:

dvc plots show
dvc plots diff

Adding More Pipeline Stages

Your pipeline is intentionally simple, but DVC makes it easy to scale.

Here are 3 realistic extension patterns:

Option A: Feature Engineering Stage

Create src/feature_engineering.py:

def engineer():
    df = pd.read_csv("outputs/processed/train.csv")
    df["sepal_ratio"] = df["sepal_length"] / df["sepal_width"]
    df.to_csv("outputs/features/train_fe.csv", index=False)

Add to dvc.yaml:

stages:
  feature_engineering:
    cmd: python src/feature_engineering.py
    deps:
      - outputs/processed/train.csv
      - src/feature_engineering.py
    outs:
      - outputs/features/

Update train dependencies:

train:
  deps:
    - outputs/features/

DVC automatically inserts the new stage in the DAG.

Option B: Hyperparameter Search Stage

You can add a Python script that runs multiple training configurations.

For example, src/hpt.py can invoke the equivalent of the following Shell commands:

python src/train.py --max_iter=100
python src/train.py --max_iter=200
python src/train.py --max_iter=500

Add a stage:

stages:
  hpt:
    cmd: python src/hpt.py
    deps:
      - src/hpt.py
      - src/train.py
      - params.yaml
    outs:
      - outputs/hpt_results/

Option C: Postprocessing Stage

Example: Generate a final report that combines the metrics and confusion matrix.

stages:
  report:
    cmd: python src/final_report.py
    deps:
      - outputs/metrics.json
      - reports/accuracy.txt
      - src/final_report.py
    outs:
      - reports/final_report.md

What's next? We recommend PyImageSearch University.

Course information:
120+ total classes • 115+ hours of on-demand code walkthrough videos • Last updated: October 2026
★★★★★ 4.84 (128 Ratings) • 16,000+ Students Enrolled

I strongly believe that if you had the right teacher you could master computer vision and deep learning.

Do you think learning computer vision and deep learning has to be time-consuming, overwhelming, and complicated? Or has to involve complex mathematics and equations? Or requires a degree in computer science?

That’s not the case.

All you need to master computer vision and deep learning is for someone to explain things to you in simple, intuitive terms. And that’s exactly what I do. My mission is to change education and how complex Artificial Intelligence topics are taught.

If you're serious about learning computer vision, your next stop should be PyImageSearch University, the most comprehensive computer vision, deep learning, and OpenCV course online today. Here you’ll learn how to successfully and confidently apply computer vision to your work, research, and projects. Join me in computer vision mastery.

Inside PyImageSearch University you'll find:

• ✓ 120+ courses on essential computer vision, deep learning, and OpenCV topics
• ✓ 94+ Certificates of Completion
• ✓ 115+ hours of on-demand video
• ✓ Brand new courses released regularly, ensuring you can keep up with state-of-the-art techniques
• ✓ Pre-configured Jupyter Notebooks in Google Colab
• ✓ Run all code examples in your web browser — works on Windows, macOS, and Linux (no dev environment configuration required!)
• ✓ Access to centralized code repos for all 540+ tutorials on PyImageSearch
• ✓ Easy one-click downloads for code, datasets, pre-trained models, etc.
• ✓ Access on mobile, laptop, desktop, etc.

Click here to join PyImageSearch University

Summary

In this lesson, you learned how to turn a simple training workflow into a fully reproducible, multi-stage ML pipeline using DVC. Instead of running scripts manually and hoping results remain consistent, you now have a system that captures every dependency, data transformation, and output in a transparent, version-controlled workflow.

We began by breaking the Iris example into 3 stages: preprocessing, training, and evaluation. We then saw how DVC connects them using dvc.yaml, dvc.lock, and a DAG. You ran the pipeline end-to-end with dvc repro, inspected which stages changed, and watched DVC rebuild only the parts affected by modifications. This is the cornerstone of reproducible, efficient ML engineering.

You also learned how to inspect pipelines visually, track metrics across runs, and optionally extend the workflow with hyperparameters, plots, and new stages. These enhancements prepare your pipelines for real-world growth, where experiments evolve and projects become more complex.

By the end, you should have a clear understanding of how DVC pipelines bring structure, repeatability, and discipline to ML projects. They turn ad hoc experimentation into a maintainable engineering process that scales with your team and codebase.

Citation Information

Singh, V. “DVC Pipelines for MLOps: Build a Reproducible ML Pipeline,” PyImageSearch, S. Huot, A. Sharma, and P. Thakur, eds., 2026, https://pyimg.co/v78x4

@incollection{Singh_2026_dvc-pipelines-mlops-build-reproducible-ml-pipeline,
  author = {Vikram Singh},
  title = {{DVC Pipelines for MLOps: Build a Reproducible ML Pipeline}},
  booktitle = {PyImageSearch},
  editor = {Susan Huot and Aditya Sharma and Piyush Thakur},
  year = {2026},
  url = {https://pyimg.co/v78x4},
}

To download the source code to this post (and be notified when future tutorials are published here on PyImageSearch), simply enter your email address in the form below!

Download the Source Code and FREE 17-page Resource Guide

Enter your email address below to get a .zip of the code and a FREE 17-page Resource Guide on Computer Vision, OpenCV, and Deep Learning. Inside you'll find my hand-picked tutorials, books, courses, and libraries to help you master CV and DL!

Download the code!

Website

The post DVC Pipelines for MLOps: Build a Reproducible ML Pipeline appeared first on PyImageSearch.

DVC for MLOps: Versioning Your Data and Models the Right Way

In this lesson, you will learn how to use DVC (Data Version Control) to version datasets, track model checkpoints, and create reproducible ML (machine learning) pipelines without bloating your Git repository. You will see how DVC integrates seamlessly with modern MLOps (machine learning operations) workflows and how it solves the versioning challenges that Git alone cannot handle.

This lesson is the 1st in a 2-part series on Data and Model Versioning with DVC:

1. DVC for MLOps: Versioning Your Data and Models the Right Way (this tutorial)
2. Lesson 2

To learn how to version your data, manage checkpoints, configure remotes, and reproduce pipelines with confidence, just keep reading.

Looking for the source code to this post?

Jump Right To The Downloads Section

DVC for MLOps: Data and Model Versioning Explained

Modern ML systems rely on large datasets, evolving model checkpoints, and pipeline stages that must be rerun reliably. Git alone was never built for multi-GB datasets or binary artifacts, and that is exactly where DVC (Data Version Control) becomes essential. In this introduction, we will explore why ML projects need a tool like DVC, how it elevates MLOps workflows, and what you will build by the end of this lesson.

What DVC Is and Why Git Alone Is Not Enough

Git excels at tracking text files, but it breaks down when you try to version datasets or model checkpoints:

• Large files slow down Git: A single dataset or checkpoint can be hundreds of megabytes (MB) or several gigabytes (GB), making Git repositories heavy and slow.
• Git cannot store multiple versions efficiently: Binary diffs do not compress well, so Git history becomes huge over time.
• Teams end up sharing data manually: Emailing zip files, keeping final_v3_really_final.csv, or relying on someone’s local drive.

DVC solves all these problems:

• Stores data outside Git, while Git tracks only lightweight metadata (.dvc files).
• Guarantees reproducibility using hashes and dependency tracking.
• Makes datasets and checkpoints shareable across the team via remotes (local, Amazon S3 (Amazon Simple Storage Service), GCS (Google Cloud Storage), Azure, SSH (Secure Shell)).
• Integrates naturally with existing Git workflows without changing how you commit or create branches.

Think of it as “Git for data and models”, tailor-made for machine learning.

How DVC Fits Into Modern MLOps

MLOps has 3 pillars:

• Reproducible code
• Reproducible environments
• Reproducible data and model artifacts

Docker or virtual environments solve the first two.

DVC completes the third.

With DVC, teams can:

• Version datasets the same way they version code
• Recreate any training run using DVC’s hashes and pipeline engine
• Store multi-GB artifacts in efficient, cache-backed, cloud-friendly storage
• Collaborate without duplicating files or manually syncing assets
• Build DAG (directed acyclic graph)-style ML pipelines (dvc.yaml) for preprocessing, training, and evaluation

Most importantly, DVC integrates smoothly with:

• PyTorch Lightning (for checkpoints)
• FastAPI and Flask apps (for deployed models)
• CI/CD (continuous integration and continuous delivery) systems (GitHub Actions, GitLab, Jenkins)
• Cloud storage providers (S3, GCS, Azure, etc.)

In short: DVC brings software engineering discipline to ML artifact management.

What We Will Build in This Lesson

In this lesson, you will build a complete DVC-enabled workflow around the provided repository:

• Dataset Versioning: You will track the sample IMDb dataset (data/raw/imdb_sample.csv) using dvc add, inspect its .dvc metadata, and understand how DVC moves the real file into its cache.
• Model Checkpoint Versioning: You will generate a dummy model checkpoint by running src/noop_train.py, then version it with DVC in the same way that you would track real PyTorch Lightning checkpoints.
• Local Remote Storage: You will configure a .dvc_remote/ folder to simulate real cloud remotes, learning how DVC pushes and pulls artifacts.
• A Mini Training Pipeline: You will run the auto-generated DVC pipeline (dvc repro) that recreates model artifacts and updates dvc.lock.

By the end, you will have a fully working data and model versioning workflow, complete with remotes, metadata, pipeline execution, and team-ready reproducibility.

It is simple enough to understand but realistic enough to scale into a real MLOps project.

Would you like immediate access to 3,457 images curated and labeled with hand gestures to train, explore, and experiment with … for free? Head over to Roboflow and get a free account to grab these hand gesture images.

Project Setup

Before we begin versioning datasets and model artifacts, let us walk through the repository structure, install dependencies, and initialize Git and DVC. This gives us a consistent, reproducible foundation for everything we will build later in the lesson.

Overview of the Lesson 1 Repository Structure

Your dvc-lesson1/ project is intentionally small, mirroring a real ML workflow but without unnecessary complexity. Here is the structure you will be working with:

dvc-lesson1/
├── data/
│   └── raw/
│       └── imdb_sample.csv
├── models/
├── src/
│   └── noop_train.py
├── .dvc/
├── .dvcignore
├── .gitignore
├── dvc.yaml
├── requirements.txt
└── README.md

What each part does:

• data/raw/: Contains your raw dataset (IMDb sample).
• models/: Stores model checkpoints generated by the training script.
• src/noop_train.py: A dummy training script that outputs model.ckpt.
• dvc.yaml: Defines a single DVC pipeline stage (generate_model).
• requirements.txt: Dependencies: DVC + pandas.
• .dvcignore and .gitignore: Ensure Git and DVC track the right files.
• .dvc/: Will hold DVC metadata and cache configuration once initialized.

This structure is simple but realistic: data, models, and code are cleanly separated, making it easy to scale toward multi-stage pipelines later.

Installing Dependencies (requirements.txt)

Your project uses only 2 core dependencies: DVC and pandas:

# DVC for data and model versioning
dvc==3.48.4

# Python basics
pandas==2.1.4

Install them using:

pip install -r requirements.txt

This installs:

• DVC core (local filesystem support)
• Optional cloud backends (S3, GCS, Azure, and SSH) can be added later using extras such as the following:

pip install 'dvc[s3]'

After installation, verify that DVC is available:

dvc --version

You should see output similar to the following:

3.48.4

If DVC installs successfully, you are ready for initialization.

Initializing Git and DVC (git init, dvc init)

DVC is designed to work on top of Git, not instead of it.

Let us initialize both:

Step 1: Initialize Git

git init

This creates the .git/ folder and prepares the repository for tracking metadata.

Step 2: Initialize DVC

dvc init

DVC sets up:

• .dvc/: internal DVC configuration and cache references
• .dvcignore: patterns for files DVC should ignore
• Git hooks for pipeline reproducibility

Commit them:

git add .dvc .dvcignore
git commit -m "Initialize DVC"

After this step, your project is now:

• Git-tracked: for code and metadata
• DVC-enabled: for datasets and model artifacts
• Pipeline-ready: via the included dvc.yaml

You now have everything required to start tracking data and model files the right way.

Need Help Configuring Your Development Environment?

All that said, are you:

• Short on time?
• Learning on your employer’s administratively locked system?
• Wanting to skip the hassle of fighting with the command line, package managers, and virtual environments?
• Ready to run the code immediately on your Windows, macOS, or Linux system?

Then join PyImageSearch University today!

Gain access to Jupyter Notebooks for this tutorial and other PyImageSearch guides pre-configured to run on Google Colab’s ecosystem right in your web browser! No installation required.

And best of all, these Jupyter Notebooks will run on Windows, macOS, and Linux!

Dataset Versioning with DVC

Your dataset is the first artifact you will version with DVC. In this section, you will inspect the raw IMDb dataset, add it to DVC tracking, understand the resulting .dvc metadata file, and see how DVC uses its cache behind the scenes. By the end, you will understand precisely why Git stores the metadata while DVC manages the actual data files.

Inspecting the Raw IMDb Dataset

Your dataset lives here:

data/raw/imdb_sample.csv

It contains 10 short IMDb movie reviews labeled positive or negative.

Here are the first few rows:

text,label
"This movie was absolutely fantastic! I loved every minute of it.",positive
"Terrible waste of time. Do not watch this movie.",negative

Even though this dataset is tiny, the workflow you are learning works identically for:

• gigabyte-scale comma-separated values (CSV) files
• multi-GB image folders
• terabyte-scale corpora

DVC applies the same versioning workflow to all these data types.

Tracking Data with DVC

To version this dataset, run:

dvc add data/raw/imdb_sample.csv

DVC automatically performs several steps:

• Calculates an MD5 hash of the file
• Stores the CSV file content in the DVC cache (.dvc/cache/)
• Makes models/model.ckpt available in the workspace as a link or copy, depending on the configuration
• Generates a metadata file:

data/raw/imdb_sample.csv.dvc

• Updates (or creates) a .gitignore inside data/raw/ so Git does not track the actual CSV.

Now add the generated metadata to Git:

git add data/raw/imdb_sample.csv.dvc data/raw/.gitignore
git commit -m "Track IMDB dataset with DVC"

Git now tracks the dataset version, not the dataset itself.

Understanding DVC Metadata Files

Let us inspect the metadata:

cat data/raw/imdb_sample.csv.dvc

You will see configuration similar to the following:

outs:
- md5: a1b2c3d4e5f6...
  size: 512
  path: imdb_sample.csv

What each field means:

• md5: fingerprint of the exact dataset content
• size: file size
• path: original location relative to its .dvc file

DVC uses the MD5 checksum to:

• detect changes
• ensure reproducibility
• retrieve the correct version from the cache or remote

This .dvc file is a small text file that represents a specific version of your dataset.

How DVC Uses Its Cache

When you ran dvc add, the physical CSV file was moved into:

.dvc/cache/

Inside the cache, every file is stored using its MD5 hash:

.dvc/cache/a1/b2c3d4e5f6...

Why?

• Hash-based storage allows deduplication
• Identical files can reference the same cached object
• Cached or remotely stored versions can be restored when needed
• Remote storage (S3, GCS, local) becomes trivial

Your working directory (data/raw/) contains either:

• a hardlink
• a symlink
• a simple copy

depending on the operating system (OS) and configuration.

If the dataset changes, DVC computes a new hash and creates a new cached version, without overwriting the old one.

This is what makes dataset versioning possible.

Why Git Tracks Metadata, Not Data

Git is excellent at managing:

• small files
• code
• configuration

However, Git is not designed to efficiently manage:

• large files
• binary blobs
• constantly evolving datasets
• large model checkpoints

Git repositories can become slow and difficult to manage when they contain large or frequently changing binary files.

DVC solves this by splitting responsibilities:

When you:

git checkout <commit>

Git restores the correct .dvc files.

Then:

dvc pull

DVC restores the exact dataset version associated with that commit.

This workflow allows Git and DVC to coordinate source-code and dataset versions.

Model Versioning with DVC

Datasets evolve over time. Model artifacts do too. Checkpoints, weights, embeddings, tokenizers, and trained models must be versioned with the same rigor as your data. In this section, you will use DVC to track a dummy checkpoint file produced by noop_train.py, and you will see how this workflow directly maps to real PyTorch Lightning training workflows.

Reviewing the Dummy Training Script

Your repository includes a tiny training script designed to simulate a real ML training job:

# src/noop_train.py
def generate_model():
    print("🚀 Starting model generation...")
    time.sleep(1)
    os.makedirs("models", exist_ok=True)

    model_path = "models/model.ckpt"
    with open(model_path, "wb") as f:
        dummy_data = b"PYTORCH_LIGHTNING_CHECKPOINT_v1.0\n" * 30
        f.write(dummy_data)

    print(f"✅ Model checkpoint generated at: {model_path}")
    print(f"📦 File size: {os.path.getsize(model_path)} bytes")

When executed, it:

• Creates a models/ directory (if it does not exist)
• Writes a 1KB dummy checkpoint called model.ckpt
• Produces stable output so you can track it with DVC

This script plays the role of your “training job” in an actual ML workflow.

Generating the Model Checkpoint

Run the dummy training script:

python src/noop_train.py

You should see:

🚀 Starting model generation...
✅ Model checkpoint generated at: models/model.ckpt
📦 File size: 960 bytes

At this point, your project structure now includes:

models/model.ckpt

This file is exactly what you would typically receive from:

• PyTorch Lightning
• Hugging Face Trainer
• Keras and TensorFlow training loops
• Custom PyTorch training scripts

In every real-world ML project, these checkpoints must be versioned.

Adding Model Artifacts to DVC

To version the checkpoint, run:

dvc add models/model.ckpt

Behind the scenes, DVC:

• Computes an MD5 checksum of the checkpoint
• Moves the physical .ckpt file into DVC’s cache (.dvc/cache/*)
• Recreates a link or copy at models/model.ckpt
• Generates:

models/model.ckpt.dvc

• Updates models/.gitignore so Git no longer tracks model.ckpt

Next, commit the metadata:

git add models/model.ckpt.dvc models/.gitignore
git commit -m "Version model checkpoint with DVC"

Your Git history now records:

• when the checkpoint changed
• why it changed (commit message)
• how to reproduce it (via DVC pipeline or scripts)

Git remains lightweight; DVC handles the binary file safely.

What the DVC Metadata Looks Like

outs:
- md5: e3b0c44298fc...
  size: 960
  path: model.ckpt

This metadata supports:

• Reproducibility: DVC knows the exact checksum of the checkpoint.
• Integrity: If the file changes, the hash changes too.
• Recoverability: DVC can restore the checkpoint from cache or remote.

This makes your ML experiments auditable and reversible.

Integrating with Real PyTorch Lightning Workflows (Conceptual)

In real-world ML pipelines, PyTorch Lightning generates checkpoints automatically:

lightning_logs/
    version_0/
        checkpoints/
            epoch=3-step=500.ckpt
            best.ckpt

To version these with DVC, you would:

1. Train Your Lightning Model

python train.py

2. Identify the Checkpoint

Lightning saves checkpoints to:

lightning_logs/version_X/checkpoints/

Pick a file (e.g., best.ckpt)

3. Add It to DVC

dvc add lightning_logs/version_0/checkpoints/best.ckpt
git add lightning_logs/version_0/checkpoints/best.ckpt.dvc
git commit -m "Track best Lightning checkpoint"

4. Push It to Your Remote

dvc push

5. Teammates or Deployment Servers Restore It

git pull
dvc pull

This gives every engineer:

• the exact model version used for inference
• the exact dataset used for training
• a reproducible pipeline state

In production settings (e.g., serving through FastAPI, TorchServe, vLLM, BentoML, or Lambda), DVC helps ensure that the intended model version is available.

Configuring DVC Remotes

Up to this point, you have tracked datasets and model artifacts with DVC. However, everything still lives locally inside your .dvc/cache/ directory. For real collaboration across team members, machines, or deployment environments, you need a remote storage backend.

A DVC remote is simply a storage location (e.g., a local directory, S3 bucket, GCS bucket, Azure Blob container, or SSH server) where DVC stores cached versions of your datasets and model files. Git tracks the metadata, while the remote stores the actual bytes.

Let us break it down.

What Are DVC Remotes and Why They Matter?

DVC remotes solve the biggest challenge in ML collaboration:

How do we store and share datasets and model checkpoints without bloating Git?

A DVC remote gives you:

Centralized artifact storage

All team members sync to the same dataset and model versions.

Lightweight Git history

Git stores only .dvc files, not large binaries.

Reproducible experiments

Anyone can restore the exact dataset and model checkpoint for a given commit.

Separation of concerns

• Git: Versions metadata
• DVC remote: versioning large files

This is essential for training consistency, debugging, deployments, and CI/CD.

Setting Up a Local DVC Remote

Your lesson uses a local folder remote, which is perfect for beginners or offline environments.

Step 1: Create the Remote Directory

Run this at the project root:

mkdir -p .dvc_remote

Step 2: Add It as a DVC Remote

dvc remote add -d local_remote .dvc_remote

Here is what this command does:

Step 3: Commit the New DVC Configuration

git add .dvc/config
git commit -m "Configure local DVC remote for dataset and model storage"

Now your project has a fully functional remote for storing versioned artifacts.

Understanding .dvc/config

After configuring the remote, open the file:

cat .dvc/config

You will see configuration similar to the following:

['core']
    remote = 'local_remote'

['remote "local_remote"']
    url = '.dvc_remote'

Let us decode this:

1. [core]

Specifies the default remote used whenever you run:

dvc push
dvc pull

2. remote “local_remote”

Defines the remote’s properties:

If you were using S3, GCS, or Azure, this section would include the applicable paths and protocol details.

Where Does the Actual Data Go?

Inside .dvc_remote/, DVC stores files under a hash-based directory structure:

.dvc_remote/
└── aa/
    └── bbcd1234...

Each stored artifact is:

• Broken into chunks (optional)
• Named by checksum (similar to .git/objects/)

This ensures deduplication and integrity.

Optional: Cloud Remotes (S3, GCS, Azure, SSH)

Even though this lesson uses a local remote, most real MLOps pipelines use cloud remotes so training jobs, teammates, and CI systems can all access the same data.

The following examples show commands for each storage option.

AWS S3 Remote

dvc remote add -d myremote s3://mybucket/dvc-storage
dvc remote modify myremote access_key_id <AWS_ACCESS_KEY>
dvc remote modify myremote secret_access_key <AWS_SECRET_KEY>

DVC will store:

s3://mybucket/dvc-storage/aa/bbcd123...

Google Cloud Storage Remote

dvc remote add -d myremote gs://mybucket/dvc-storage

Authenticate with:

gcloud auth application-default login

Azure Blob Storage Remote

dvc remote add -d myremote azure://mycontainer/dvc-storage
dvc remote modify myremote connection_string "<AZURE_CONNECTION_STRING>"

SSH and SFTP Remote (Self-Hosted)

dvc remote add -d myremote ssh://user@server:/path/to/storage
dvc remote modify myremote keyfile ~/.ssh/id_rsa

Great for on-premise teams.

When Should You Use Cloud Remotes?

Using DVC Push and Pull to Sync Data and Model Artifacts

So far, you have tracked:

• A dataset (data/raw/imdb_sample.csv)
• A model checkpoint (models/model.ckpt)
• A pipeline stage (dvc.yaml → generate_model)

However, all of these artifacts still live locally in your .dvc/cache/ directory.

To actually share them, back them up, and restore them anywhere, you need to sync them with your configured DVC remote.

This is where dvc push and dvc pull come in.

Uploading Artifacts to a DVC Remote

Once a remote is configured (e.g., .dvc_remote/), DVC can upload any tracked artifacts to it.

Command

dvc push

What happens?

DVC performs several steps under the hood:

• Scans all .dvc files and dvc.lock to determine which files are tracked.
  Example tracked files:

data/raw/imdb_sample.csv
models/model.ckpt

• Locates these files in the local .dvc/cache/: Remember: DVC stores the real data in the cache using checksum names.
• Copies the cache objects to the remote: For example, a file might be uploaded to:

.dvc_remote/aa/bbcd1234...

• Skips uploading duplicates: If the checksum already exists in the remote, DVC does not upload it again.

Output Example

When pushing your dataset and model:

$ dvc push
100%|████████████████████████|1/1 [00:00<00:00, 101.23file/s]

This confirms your artifacts are safely stored in the remote.

Restoring Artifacts (dvc pull)

Imagine a teammate clones your Git repository:

git clone <your-repo>
cd dvc-lesson1

They now have:

• Your code
• Your .dvc metadata files
• Your dvc.yaml pipeline definition
• Your .dvc/config pointing to the remote

But they do not have the data or model artifacts.

To restore them, they run:

dvc pull

With this single command, DVC:

• Reads your .dvc files to know which artifacts should exist
• Downloads their cache objects from the remote into .dvc/cache/
• Recreates the original directory structure:

data/raw/imdb_sample.csv
models/model.ckpt

Real Output Example

$ dvc pull
100%|█████████████████████|1/1 [00:00<00:00, 250.50file/s]

Now your teammate has:

• The exact same dataset you used.
• The same model checkpoint.
• A fully reproducible environment.

How Git and DVC Work Together

The magic of DVC comes from how it pairs with Git while keeping responsibilities clean.

Git Stores:

• Code (src/)
• Pipeline definitions (dvc.yaml, dvc.lock)
• Metadata (*.dvc)
• Remote configuration (.dvc/config)
• Docs and scripts

Git never stores large binaries.

DVC Remote Stores:

• Actual dataset files
• Actual model checkpoints
• Other large ML artifacts (e.g., features, embeddings, plots, and metrics)

All files are stored by checksum, ensuring:

• Deduplication
• Integrity
• Fast restoration

Workflow Example for a Team of 3

Let us imagine you run:

python src/noop_train.py
dvc add models/model.ckpt
git add models/model.ckpt.dvc
git commit -m "Add v1 checkpoint"
dvc push
git push

Now the remote contains the checkpoint, and Git contains its metadata.

Teammate Workflow

Your teammate runs:

git pull
dvc pull

They instantly get:

• The same model checkpoint
• The same dataset
• The same pipeline configuration

They can:

• Reproduce training
• Evaluate the model
• Extend experiments
• Debug issues using the same data

Why This Is a Game-Changer for ML Teams

Building a Reproducible Machine Learning Pipeline with DVC

One of DVC’s superpowers is that it can turn arbitrary commands (e.g., training scripts, preprocessing steps, feature generation, evaluation) into reproducible pipelines.

Even though this lesson uses a simple “noop” model generator, the ideas scale directly to real PyTorch Lightning workflows.

In this section, you will learn how DVC automatically:

• Creates a pipeline stage (dvc.yaml)
• Tracks script dependencies
• Tracks model artifacts
• Rebuilds outputs only when something changes
• Guarantees reproducibility using dvc.lock

Understanding the Auto-Generated DVC Pipeline Configuration

When you ran:

dvc add models/model.ckpt

DVC tracked the artifact, but did not create a pipeline.

To turn model generation into a reproducible pipeline, you used:

dvc run -n generate_model -d src/noop_train.py -o models/model.ckpt python src/noop_train.py

Alternatively, a prewritten dvc.yaml file can define the stage.

Your repository now contains the following file:

dvc.yaml

stages:
  generate_model:
    cmd: python src/noop_train.py
    deps:
      - src/noop_train.py
    outs:
      - models/model.ckpt

Let us break this down:

Why This Matters

If you modify src/noop_train.py, DVC knows the dependency changed and will rerun the pipeline.

If nothing changed, DVC skips execution.

It is like make but designed for ML.

Running the DVC Pipeline Stage

To execute the pipeline stage, run:

dvc repro

Example output:

$ dvc repro

Running stage 'generate_model':
> python src/noop_train.py
🚀 Starting model generation...
✅ Model checkpoint generated at: models/model.ckpt
📦 File size: 870 bytes

Updating lock file 'dvc.lock'

What happened?

• DVC checked whether src/noop_train.py changed.
• Since it is the first run, DVC executed the command.
• The checkpoint was generated.
• DVC updated dvc.lock with:
  • Hash of the script
  • Hash of the output
  • Command used

Now the pipeline is fully reproducible.

Rerunning Without Changes

Run again:

dvc repro

Output:

Stage 'generate_model' didn't change, skipping

DVC intelligently avoids unnecessary work.

Modify the Script and Rerun

Add a comment or modify noop_train.py:

# Added comment
def generate_model():
    ...

Run:

dvc repro

Now DVC sees the dependency changed, so it rebuilds the stage automatically.

Inspecting the DVC Lock File

Next to dvc.yaml, you will find the auto-generated lock file:

dvc.lock

stages:
  generate_model:
    cmd: python src/noop_train.py
    deps:
      - path: src/noop_train.py
        md5: <hash>
        size: <bytes>
    outs:
      - path: models/model.ckpt
        md5: <hash>
        size: <bytes>

This file stores:

• Exact command used
• Checksums of dependencies
• Checksums of outputs
• File sizes

Why This Is Important

If model.ckpt changes, its MD5 checksum will also change, signaling a new version.

Verify the Lock State

dvc status

Example output:

Data and pipelines are up to date.

Change the script and the output becomes:

Stage 'generate_model' changed.

Full Real-World Workflow with Lightning, DVC, and GitHub

Everything you have done so far, including tracking datasets, tracking checkpoints, and defining a DVC pipeline, maps to a real-world workflow.

This section stitches it all together into a workflow where:

• Lightning handles training
• DVC handles dataset and checkpoint versioning
• GitHub handles code and metadata
• A remote store (e.g., S3, GCS, or Azure) stores large files

This is the pattern used by most mature MLOps teams.

Mapping Lightning Checkpoints to DVC Artifacts

In DVC Lesson 1, you generated a dummy checkpoint (model.ckpt) using noop_train.py.

In real projects, PyTorch Lightning generates checkpoints automatically:

lightning_logs/
└── version_0/
    └── checkpoints/
        ├── epoch=4-step=500.ckpt
        ├── last.ckpt
        └── best.ckpt

These files can easily reach hundreds of MB and change frequently.

This makes them terrible candidates for Git, but perfect candidates for DVC.

How Lightning and DVC Connect

Your workflow looks like this:

• Lightning trains the model
• Lightning saves a checkpoint in checkpoints/
• You tell DVC to track it:

dvc add lightning_logs/version_0/checkpoints/best.ckpt
git add lightning_logs/version_0/checkpoints/best.ckpt.dvc
git commit -m "Track best Lightning checkpoint for experiment A"

• Push to remote storage:

dvc push

• Commit code and metadata to GitHub:

git push

Now teammates can reproduce your training output exactly:

git pull
dvc pull  # downloads the checkpoint

No need to upload .ckpt files manually, no need to Slack someone a 500 MB file, and no more “which model version did you use?” chaos.

What Files Go to Git, DVC, and Remote Storage

A clean ML workflow is all about putting each file in the right place.

Table 8 shows what goes where:

The Core Rule: Git tracks the metadata, DVC tracks the large files, a remote storage backend stores the bytes.

This separation prevents Git from bloating and keeps your repo lightning fast.

Suggested Team Workflow (Industry-Grade)

Below is the recommended daily workflow for ML teams using Lightning, DVC, and GitHub.

This is exactly how large production teams work.

Step 1: Pull Code and Artifacts

When starting on a fresh machine:

git pull
dvc pull

You now have:

• Latest code
• Latest dataset
• Latest checkpoints

Step 2: Train Using Lightning

python train.py

Lightning writes:

lightning_logs/version_3/checkpoints/best.ckpt

Step 3: Version the New Checkpoint

dvc add lightning_logs/version_3/checkpoints/best.ckpt
git add lightning_logs/version_3/checkpoints/best.ckpt.dvc
git commit -m "Add checkpoint for experiment: bigger batch size"

Step 4: Push Artifacts and Git Metadata

dvc push   # Uploads large files to DVC remote
git push   # Pushes code + metadata

Everyone on the team can now reproduce your run.

Workflow Diagram (Textual)

Lightning → produces checkpoint.ckpt  
        ↓  
DVC add → creates .dvc file  
        ↓  
git add *.dvc → store metadata  
        ↓  
dvc push → upload model to remote  
        ↓  
git push → sync metadata + code

Perfect reproducibility, every time.

Bonus: Branch-Based Experimentation Pattern

Teams often create experiment branches:

experiments/
    ├── exp-bigger-lr
    ├── exp-more-layers
    ├── exp-new-augmentation

Each branch has:

• Code changes
• DVC-tracked checkpoints
• Git-tracked .dvc metadata

Merging into main requires:

git merge <branch>
dvc pull

DVC then restores the artifacts referenced by the merged metadata.

What's next? We recommend PyImageSearch University.

Course information:
120+ total classes • 115+ hours of on-demand code walkthrough videos • Last updated: October 2026
★★★★★ 4.84 (128 Ratings) • 16,000+ Students Enrolled

I strongly believe that if you had the right teacher you could master computer vision and deep learning.

Do you think learning computer vision and deep learning has to be time-consuming, overwhelming, and complicated? Or has to involve complex mathematics and equations? Or requires a degree in computer science?

That’s not the case.

All you need to master computer vision and deep learning is for someone to explain things to you in simple, intuitive terms. And that’s exactly what I do. My mission is to change education and how complex Artificial Intelligence topics are taught.

If you're serious about learning computer vision, your next stop should be PyImageSearch University, the most comprehensive computer vision, deep learning, and OpenCV course online today. Here you’ll learn how to successfully and confidently apply computer vision to your work, research, and projects. Join me in computer vision mastery.

Inside PyImageSearch University you'll find:

• ✓ 120+ courses on essential computer vision, deep learning, and OpenCV topics
• ✓ 94+ Certificates of Completion
• ✓ 115+ hours of on-demand video
• ✓ Brand new courses released regularly, ensuring you can keep up with state-of-the-art techniques
• ✓ Pre-configured Jupyter Notebooks in Google Colab
• ✓ Run all code examples in your web browser — works on Windows, macOS, and Linux (no dev environment configuration required!)
• ✓ Access to centralized code repos for all 540+ tutorials on PyImageSearch
• ✓ Easy one-click downloads for code, datasets, pre-trained models, etc.
• ✓ Access on mobile, laptop, desktop, etc.

Click here to join PyImageSearch University

Summary

In this lesson, you learned why Git alone cannot support modern machine-learning workflows and how DVC solves that gap by versioning large datasets and model checkpoints without bloating your repository. You saw how DVC keeps Git lightweight by storing only metadata while placing real artifacts in a structured cache and optional remote storage.

You then worked through a hands-on workflow: inspecting a real dataset, adding it to DVC with dvc add, generating a model checkpoint through a dummy training script, and versioning that checkpoint just like data. Along the way, you learned how .dvc files, .dvc/cache/, and auto-generated .gitignore entries work together to make data tracking seamless.

Next, you configured remotes for artifact storage. You started with a local directory and then examined how easily DVC integrates with cloud backends (e.g., S3, GCS, Azure, and SSH). You pushed and pulled artifacts, simulating team collaboration, and explored how Git manages code while DVC synchronizes the corresponding data.

Finally, you ran a mini DVC pipeline using the auto-generated dvc.yaml, reproduced it with dvc repro, and inspected dvc.lock to understand how DVC captures reproducibility. You also connected these concepts back to real PyTorch Lightning workflows, mapping where datasets, checkpoints, and code naturally fit into a Git and DVC ecosystem.

Together, these steps form a strong foundation for reproducible ML projects. In the next lesson, you will expand this foundation into multi-stage DVC pipelines, experiment tracking, and full workflow automation.

Citation Information

Singh, V. “DVC for MLOps: Versioning Your Data and Models the Right Way,” PyImageSearch, S. Huot, A. Sharma, and P. Thakur, eds., 2026, https://pyimg.co/bu1ya

@incollection{Singh_2026_dvc-for-mlops-versioning-data-models-right-way,
  author = {Vikram Singh},
  title = {{DVC for MLOps: Versioning Your Data and Models the Right Way}},
  booktitle = {PyImageSearch},
  editor = {Susan Huot and Aditya Sharma and Piyush Thakur},
  year = {2026},
  url = {https://pyimg.co/bu1ya},
}

To download the source code to this post (and be notified when future tutorials are published here on PyImageSearch), simply enter your email address in the form below!

Download the Source Code and FREE 17-page Resource Guide

Enter your email address below to get a .zip of the code and a FREE 17-page Resource Guide on Computer Vision, OpenCV, and Deep Learning. Inside you'll find my hand-picked tutorials, books, courses, and libraries to help you master CV and DL!

Download the code!

Website

The post DVC for MLOps: Versioning Your Data and Models the Right Way appeared first on PyImageSearch.

YOLO26 Open-Vocabulary Object Detection with YOLOE-26

In this lesson, you will learn how YOLOE evolved into YOLOE-26, how open-vocabulary detection works, and how to use text prompts, visual prompts, and prompt-free inference with Ultralytics.

This lesson is the 1st in a 2-part series on YOLOE-26 and open-vocabulary detection:

1. YOLO26 Open-Vocabulary Object Detection with YOLOE-26 (this tutorial)
2. Lesson 2

To learn how to use YOLOE-26 for open-vocabulary object detection with text prompts, visual prompts, and prompt-free inference, just keep reading.

Looking for the source code to this post?

Jump Right To The Downloads Section

In the last YOLO26 lesson, we stayed in the closed-set world. We loaded a modern YOLO detector, ran it on images and video, and saw just how fast and polished the Ultralytics pipeline has become.

But closed-set detection has a hard limit. A model can only detect the classes it learned during training.

That sounds obvious until you hit it in practice.

You may want to detect a multimeter, a barcode scanner, a microscope slide box, a soldering iron, or a specific logo in a warehouse photo. A standard YOLO model cannot suddenly understand those categories just because you typed their names. If the class was not part of training, the model either misses it or maps it to the closest thing it already knows.

That is the problem YOLOE was built to solve.

This lesson is the foundation piece for the two-part series. We are not fine-tuning anything yet. We are first building the mental model you need before the workflow gets more ambitious in Lesson 2.

In this lesson, you will learn:

• what open-vocabulary detection actually means in plain language
• what YOLOE introduced to the YOLO family
• how YOLOE-26 extends that idea into the YOLO26 generation
• how to run text-prompted, visual-prompted, and prompt-free inference
• when YOLOE-26 is the right tool and when a standard YOLO26 model is still the better choice

Understanding Closed-Set YOLO Object Detection

Let us begin with the constraint that defines standard YOLO models.

A regular YOLO26 detector is excellent at recognizing the categories it was trained on. But it still operates inside a fixed label space. If the class is outside that label space, the model has no mechanism to dynamically add it at inference time.

That is the itch this lesson scratches.

The first example below runs a standard YOLO26 detector on the familiar bus.jpg sample image that ships with Ultralytics:

from ultralytics import YOLO
from ultralytics.utils import ASSETS

baseline_model = YOLO("yolo26s.pt")
baseline_results = baseline_model.predict(ASSETS / "bus.jpg", conf=0.25)
baseline_results[0].show()

In that image, YOLO26 does exactly what you would expect. It finds the obvious closed-set classes such as people and the bus itself.

The problem shows up when your application stops looking like a benchmark and starts looking like the real world. In production, the target object list is often messy, narrow, and unstable. Maybe your robotics pipeline needs to find a clamp that appears in only one manufacturing line. Maybe your e-commerce workflow needs to find a specific product family whose packaging changed last quarter. Maybe your lab automation setup needs to distinguish one kind of equipment tray from another.

In all of those scenarios, a closed-set model puts you in one of 2 boxes:

• the class already exists in the model, so you are fine
• the class does not exist, so you need a new data collection and training workflow

That second path is expensive. You need images. You need labels. You need training time. You need evaluation. You probably need iteration because the first pass will not be clean enough.

This is why open-vocabulary detection matters. It does not replace training forever, but it gives you a much more flexible first step.

This is the key mental shift for the rest of the article. YOLOE does not just make YOLO “a bit better.” It changes how you specify what the model should look for.

Where YOLO Fits Among Object Detection Models

Before going further, it helps to know where YOLO sits among object detection models generally, since not every detector works the same way under the hood.

Most object detection models fall into 1 of 2 architectural families. Two-stage detectors (e.g., Faster Region-based Convolutional Neural Network (R-CNN) implementations available through Detectron2) first propose candidate regions in an image, then classify each region separately. This tends to produce strong accuracy, but the two-pass design costs speed. One-stage detectors, the family YOLO belongs to, predict boxes and classes in a single forward pass. That speed advantage helped make YOLO a common choice for real-time applications, historically trading some accuracy for substantially faster inference.

That architectural split is only half the picture, though, and it is the half most tutorials stop at. A second, independent axis matters just as much for this lesson: whether the class list is closed or open. Whether a detector is one-stage or two-stage, the overwhelming majority share the same constraint: a fixed set of classes, usually drawn from a benchmark like Common Objects in Context (COCO), locked in at training time. To detect anything outside that list, you have to collect new images, label them, and train an object detection model from scratch on the updated data. That process works, but it is slow, and it has to be repeated every time your target categories change.

YOLO26 is a fast, one-stage, closed-set detector. It is excellent at what it was trained to recognize but is bound by that same COCO-style constraint everywhere else. YOLOE-26 keeps YOLO’s one-stage speed but breaks the second constraint. That is the shift the rest of this lesson is about.

Why Open-Vocabulary and Zero-Shot Object Detection Matter

Before we touch YOLOE itself, let us make the term open vocabulary concrete.

In a closed-set detector, the label list is fixed ahead of time. In an open-vocabulary detector, the model can work from language prompts, visual references, or a built-in broader vocabulary instead of a single hardcoded class list.

The practical benefit is simple: you can ask the detector for new concepts without rebuilding the whole model every time the problem changes. That is especially useful when the object is outside everyday benchmark classes, when the label list changes often, or when you want to bootstrap a later fine-tuning workflow.

This is closely related to what is called zero-shot detection in the research literature: the ability to detect object categories the model was never explicitly trained on. Open-vocabulary detection is the broader, more flexible version of that idea, built around accepting arbitrary prompts at inference time rather than a fixed set of unseen classes.

Earlier systems (e.g., Grounded Language-Image Pre-training (GLIP), Open-World Localization Vision Transformer (OWL-ViT), and Grounding DINO) proved that open-vocabulary detection works, but they also made clear how expensive the usual vision-language path can become.

YOLOE matters because it tries to keep that promptable behavior while staying in the real-time YOLO regime.

For PyImageSearch readers who have already trained detectors before, the right framing is this: YOLOE-26 is not a replacement for every older YOLO workflow. It is a new capability layer on top of familiar YOLO deployment patterns.

What YOLOE Introduced

YOLOE, introduced in the paper YOLOE: Real-Time Seeing Anything, takes the familiar YOLO workflow and adds open-vocabulary behavior on top of it.

Instead of being locked to a fixed class list, YOLOE can operate through 3 different prompting modes:

• Text prompting: tells the model what classes to look for using words
• Visual prompting: shows the model an example object and asks it to find more like it
• Prompt-free mode: uses a built-in open vocabulary without supplying your own prompts at inference time

This is the core idea you should carry forward. YOLOE is still part of the YOLO family. It still feels like Ultralytics. But it replaces the “fixed class list only” assumption with a more flexible prompting interface.

The reason that matters is simple. In a text-prompted setting, the detector is no longer just asking, “Does this image contain one of the classes already in its label space?” It is also asking, “How well does this region align with the prompt concept it was given?” That is the conceptual jump.

You do not need to understand every detail of the paper to benefit from it. For this lesson, the important point is that YOLOE creates a bridge between image features and promptable concepts. Sometimes those concepts come from text. Sometimes they come from visual examples. Sometimes they come from a prebuilt broader vocabulary.

At this point, the most useful question is not “How does every module work?” It is “What new things can we now ask the model to do?” The answer is that YOLOE gives you multiple ways to specify the object of interest, which is exactly what a closed-set detector lacks.

How YOLOE-26 Extends Open-Vocabulary Detection to YOLO26

Earlier YOLOE releases were shipped in YOLOv8-based and YOLO11-based model families. YOLOE-26 brings the same open-vocabulary idea into the YOLO26 family.

The following 2 points matter here. YOLO26 introduced a cleaner end-to-end design with native non-maximum suppression (NMS)-free inference, simpler heads, and a stronger accuracy-latency balance. YOLOE-26 inherits that foundation while keeping the promptable modes that made YOLOE interesting in the first place.

Architecturally, it still follows the familiar YOLO pattern: a backbone for feature extraction, a neck for multi-scale fusion, and a prediction head that turns fused features into boxes, masks, and region-level representations.

According to the Ultralytics docs, YOLOE-26:

• extends the YOLOE family onto the YOLO26 backbone
• inherits the NMS-free end-to-end design of YOLO26
• supports text prompting, visual prompting, and prompt-free inference
• is available across 5 scales: n, s, m, l, and x

Model Family and Checkpoints

The naming convention is easy to miss when you first see the model files:

• yoloe-26n-*: nano variant
• yoloe-26s-*: small variant
• yoloe-26m-*: medium variant
• yoloe-26l-*: large variant
• yoloe-26x-*: extra-large variant

For this lesson, we are using the s scale so the workflow stays light enough to reproduce easily. We are also using -seg checkpoints because the overlays are easier to inspect visually and align well with the official Ultralytics examples. The released pretrained YOLOE checkpoints are segmentation-first, which is why the lesson uses segmentation weights even when the conceptual discussion focuses on detection behavior.

Choosing the Right Checkpoint

In practice, there are really 2 checkpoint families to keep straight:

• yoloe-26*-seg.pt: for text prompting and visual prompting
• yoloe-26*-seg-pf.pt: for prompt-free inference

The scale suffix then controls the usual size-speed-accuracy tradeoff:

• n: for the smallest and lightest deployment target
• s: for a practical small-model starting point
• m and l: when you can spend more compute for stronger quality
• x: when you want the highest-end model in the family

For Lesson 1, the s variant is the right teaching checkpoint. It is large enough to show the behavior clearly, but still light enough that readers can reproduce the examples without needing a heavyweight setup. The segmentation-first checkpoints also help because masks make qualitative inspection easier, even when our discussion is mostly about the detection logic.

This is also where Lesson 2 starts to come into focus. We are not reproducing the full YOLOE training pipeline here, but we will borrow the same high-level idea later: use an open-vocabulary model to help create labels for a narrower downstream detector.

The following 3 architecture details are worth calling out while that figure is on screen:

• Re-parameterizable Region-Text Alignment (RepRTA): supports text-prompted detection.
• Semantic-Activated Visual Prompt Encoder (SAVPE): supports visual-prompted detection.
• Lazy Region-Prompt Contrast (LRPC): supports prompt-free open-vocabulary inference.

You do not need to memorize those names. The practical takeaway is that YOLOE adds dedicated machinery for each prompting mode while preserving the familiar YOLO deployment feel.

How the YOLOE-26 Flow Works

At this point, it helps to slow down and build a slightly more mechanical mental model of what is happening inside the model.

At a high level, YOLOE-26 still begins like a normal YOLO pipeline. The input image goes through a backbone that extracts visual features, then through a neck that fuses those features across scales so small, medium, and large objects can all be represented well. Up to that point, the story still feels very familiar to anyone who has used YOLO before.

The open-vocabulary twist happens after those visual features have been built.

Conceptually, you can think of YOLOE-26 as doing 5 stages:

• Input image to visual features: the image is converted into multi-scale feature maps by the backbone.
• Feature fusion: the neck combines information across scales so object evidence is available at multiple resolutions.
• Region prediction: the model predicts boxes and, in the segmentation checkpoints used here, masks as well.
• Prompt alignment: the model also needs a way to compare each candidate region against some notion of “what you are asking for.”
• Similarity scoring to final detections: region-level visual features are matched against prompt embeddings, and the best matches become named detections.

Another useful way to think about the head is this: boxes and masks still come from a normal YOLO-style prediction path, but class scoring is no longer a fixed set of logits over a closed label list. Instead, each candidate region is compared against prompt embeddings, and that similarity score takes over the role that fixed class scores play in a standard detector.

If you want a simple intuition, think of standard YOLO as answering, “Where are the objects already in the label space?” YOLOE-26 answers two questions at once: “Where are the candidate objects?” and “Which of these candidate regions best matches the prompt concept it was given?”

One important note: the simplified flow diagram below is a conceptual teaching figure, not a literal one-to-one tensor graph from the paper. That is intentional. For a blog lesson, the goal is to make the pipeline intuitive before readers dive into implementation details.

What Each Prompt Mode Changes

The cleanest way to understand the architecture is to ask what changes between the 3 prompting modes.

For text prompting, YOLOE uses RepRTA. According to the YOLOE paper, this module refines pretrained text embeddings through a lightweight auxiliary network so the prompt representation aligns better with the detector’s visual region features. The important deployment detail is that this lightweight refinement can be folded back into the model at inference time, which is why text prompting does not introduce the kind of heavy runtime penalty many vision-language models do.

For visual prompting, YOLOE uses SAVPE. Instead of starting from words, it starts from a reference object. The prompt encoder uses semantic and activation cues from that reference so the detector can look for visually similar regions elsewhere. That is why visual prompting feels closer to one-shot retrieval than to ordinary closed-set classification.

For prompt-free mode, YOLOE uses LRPC. Here the model does not wait for an external text prompt at all. It uses a built-in vocabulary and specialized internal embeddings, then scores candidate regions against that internal vocabulary. The Ultralytics docs describe this as open-set recognition using internal embeddings trained on large vocabularies, which is what lets prompt-free mode run without an external prompt encoder at inference time.

This also explains an important checkpoint detail for readers. Text and visual prompting use the same main YOLOE checkpoint family, while prompt-free mode uses separate -pf weights because those models are trained as built-in large-vocabulary variants rather than prompt-conditioned ones.

Where YOLO26 Changes the Base

Now add the YOLO26 side of the story.

The earlier YOLOE families were built on prior YOLO backbones, and YOLOE-26 inherits the lighter, native end-to-end design of YOLO26. According to the YOLO26 paper, that includes NMS-free end-to-end inference, a lighter head with Distribution Focal Loss (DFL) removed, and a training recipe designed to better match the inference-time head. In practice, that means YOLOE-26 is not just “YOLOE with a new name.” It is YOLOE running on a cleaner deployment-oriented detector backbone.

The key intuition is this: YOLOE contributes the promptable alignment machinery, and YOLO26 contributes the faster, simpler end-to-end detector base.

How YOLOE-26 Trains for Open-Vocabulary Object Detection

Readers usually do not need the full training recipe, but they do benefit from understanding what is trained differently.

How Text, Visual, and Prompt-Free Training Diverge

For text prompting, the paper explains that pretrained text embeddings are refined by a lightweight auxiliary network before being aligned with visual region features. That refinement step is a big part of why YOLOE can work with prompts effectively without dragging a large language branch through the full deployment path.

For visual prompting, the model learns how to turn a reference object into a useful prompt representation rather than treating the crop as a raw patch match. That is what SAVPE is doing conceptually: learning a compact visual prompt that can be compared against candidate regions.

For prompt-free mode, the model is trained to work against a built-in vocabulary and internal embedding space, so it can still perform open-vocabulary recognition even when no external prompt is supplied.

Again, the lesson-level takeaway is not every training detail. The takeaway is that YOLOE is trained to bring regions and prompts into a comparable embedding space, then YOLOE-26 places that behavior on top of the more deployment-friendly YOLO26 detector.

Would you like immediate access to 3,457 images curated and labeled with hand gestures to train, explore, and experiment with … for free? Head over to Roboflow and get a free account to grab these hand gesture images.

Configuring Your Development Environment

To follow this guide, you need to install the Ultralytics package.

Luckily, Ultralytics is pip-installable:

$ pip install -U ultralytics

If you want to be a little safer for local runs, you can use:

$ pip install -U ultralytics opencv-python matplotlib

If you need help configuring your development environment for OpenCV, we highly recommend reading our pip install OpenCV guide. It will have you up and running in minutes.

Need Help Configuring Your Development Environment?

All that said, are you:

• Short on time?
• Learning on your employer’s administratively locked system?
• Wanting to skip the hassle of fighting with the command line, package managers, and virtual environments?
• Ready to run the code immediately on your Windows, macOS, or Linux system?

Then join PyImageSearch University today!

Gain access to Jupyter Notebooks for this tutorial and other PyImageSearch guides pre-configured to run on Google Colab’s ecosystem right in your web browser! No installation required.

And best of all, these Jupyter Notebooks will run on Windows, macOS, and Linux!

YOLOE-26 Object Detection Benchmarks and Performance

Before writing code, let us calibrate expectations.

Ultralytics reports that YOLOE-L and YOLOE26-L preserve near-identical inference speed to their underlying closed-set counterparts, while adding open-vocabulary capability. The same documentation also reports stronger Large Vocabulary Instance Segmentation (LVIS) open-vocabulary performance for the YOLO26-based branch.

The important pattern is not just that YOLOE-26 scores higher on LVIS. It does so while preserving the same reported T4 latency as YOLO11-L and YOLOE-L, which is exactly why the open-vocabulary story is practical rather than just academic.

The docs also make an important practical claim: in the regular closed-set case, the open-world additions in YOLOE can be re-parameterized back into a standard YOLO-style path, so you do not pay extra inference cost just for carrying the capability.

It is also useful to unpack the benchmark language briefly so readers do not treat the numbers like magic:

• LVIS is a long-tail detection benchmark, which makes it a more meaningful place to discuss open-vocabulary behavior than a short everyday-class benchmark
• AP is average precision, so higher is better
• T4 latency gives you a rough sense of runtime cost on a standard GPU reference point

The benchmark story here is narrower and more useful than hype. YOLOE-26 improves the open-vocabulary side of the problem while staying in the performance neighborhood practitioners expect from modern YOLO models.

Hands-On with Text Prompting

Now let us move from concept to code.

The most beginner-friendly way to meet YOLOE-26 is through text prompting. You load a YOLOE model, call set_classes() once, and then run predict() just like you would with a normal Ultralytics model.

from ultralytics import YOLOE
from ultralytics.utils import ASSETS

text_model = YOLOE("yoloe-26s-seg.pt")
text_model.set_classes(["person", "bus"])
text_results = text_model.predict(ASSETS / "bus.jpg", conf=0.25)
text_results[0].show()

The following 2 points are worth noticing. The application programming interface (API) is almost boringly simple, and the class list is no longer hardwired into the checkpoint in the same way a standard closed-set detector is. You are telling the model what to care about at inference time.

In this first text-prompt example, we are using person and bus because they are stable and easy to reproduce in the official sample image. Once you understand the workflow, you can swap in more interesting prompts that match your own data.

That basic example is intentionally conservative. It is not trying to prove that YOLOE-26 magically does something a closed-set COCO model never could. It is showing you the mechanics in the simplest reproducible way.

Once that is clear, the real value comes from changing the prompt list.

How to Think About Prompt Design

This is where readers often make a subtle mistake. They treat prompting like keyword search. Open-vocabulary detection is not just string matching, so prompt quality matters.

A few practical rules help:

• start with short, concrete nouns
• avoid vague multi-object phrases
• avoid over-describing the object unless needed
• try synonyms if the first prompt underperforms
• if possible, test singular and plural forms only after trying the simplest base term

For example, soldering iron is likely a better prompt than small metal electronics repair tool with handle. The latter contains more words, but that does not automatically make it better.

You should also expect some prompts to fail for perfectly normal reasons:

• the object may be tiny
• the prompt may be semantically broad
• the object may be heavily occluded
• the prompt may describe a category the model has weak visual grounding for

This is simply the cost of asking a flexible model to generalize beyond a fixed short class list.

Good Prompts vs. Weak Prompts

A few concrete examples make this easier to internalize:

• soldering iron: stronger than tool
• traffic light: stronger than street object
• barcode scanner: stronger than electronics device
• bus: stronger than large road vehicle with windows

The pattern is consistent. Good prompts are usually short, concrete, and visually grounded. Weak prompts are often too broad, too abstract, or too wordy. If the first prompt underperforms, try a nearby synonym before you conclude the model cannot handle the category at all.

Reading the Result, Not Just Looking At It

Annotated images are useful, but the result object tells you more than the picture does. Printing the top class names and confidence values makes it easier to inspect the output programmatically.

Open-vocabulary work often involves a short loop:

• try a prompt
• inspect detections
• revise the prompt
• re-run inference

This feedback loop is why an interactive workflow works well for Lesson 1. You are not yet building a packaged application. You are exploring the detector’s behavior iteratively.

Inspecting the Result Object

If you want to move one step beyond screenshots, inspect the prediction object directly:

result = text_results[0]

print(result.names)
print(result.boxes.xyxy[:3])
print(result.boxes.conf[:3])
print(result.boxes.cls[:3])

if result.masks is not None:
    print(result.masks.data.shape)

That quick inspection tells you almost everything you need for downstream work:

• result.names: maps class identifiers (IDs) to label names
• result.boxes.xyxy: stores box coordinates
• result.boxes.conf: stores confidence scores
• result.boxes.cls: stores predicted class IDs
• result.masks: appears because we are using segmentation-first checkpoints, so mask tensors are available alongside the boxes

This is the easiest bridge from “nice demo” to “usable building block.” Once readers understand where the predictions live, they can start saving detections, filtering them, or passing them into a larger pipeline.

Why Visual Prompting Is the Real Superpower

Text prompts are the easiest entry point, but visual prompting is where YOLOE-26 starts to feel genuinely different from a normal detector.

Sometimes words are not enough. Maybe the object is a very specific industrial part. Maybe the text label is ambiguous. Maybe the thing you want to find is easier to show than describe.

That is what visual prompting is for.

How Visual Prompts Are Structured

With visual prompts, you give the model one or more bounding boxes around reference objects. YOLOE-26 then uses those examples to find visually similar instances.

The next example uses the exact same bus.jpg sample image from the docs. To keep the first visual-prompt workflow stable and easy to reproduce, we place 1 prompt box around a person.

import numpy as np

from ultralytics import YOLOE
from ultralytics.models.yolo.yoloe import YOLOEVPSegPredictor
from ultralytics.utils import ASSETS

visual_model = YOLOE("yoloe-26s-seg.pt")
visual_model.set_classes(["person"])
visual_prompts = dict(
   bboxes=np.array(
       [
           [221.52, 405.8, 344.98, 857.54],
       ]
   ),
   cls=np.array([0]),
)

visual_result = visual_model.predict(
   ASSETS / "bus.jpg",
   visual_prompts=visual_prompts,
   predictor=YOLOEVPSegPredictor,
   conf=0.25,
)[0]
visual_result.names = {0: "person"}
visual_result.show()

The only slightly unusual part here is the visual_prompts dictionary:

• bboxes: contains the reference box
• cls: contains a sequential ID that associates that box with the prompt category

These are not global COCO class IDs. They are temporary identifiers for the prompt session.

For visualization, we also relabel that temporary prompt ID so the plotted output says person instead of object0.

This is the mode many tutorials skip, but it is extremely practical. If you can show the model what an object looks like once, you can turn that into a lightweight one-shot retrieval workflow.

When Visual Prompting Beats Text Prompting

Visual prompting is most useful when the object is easier to show than to describe. Suppose you are looking for a very specific wrench, connector, or control knob. A text label like metal connector may be too broad, but a carefully drawn reference box can tell the model, “Find more objects that look like this.” In practice, make the prompt box tight, representative, and light on background clutter. Very small prompt regions can be brittle, so start with larger, visually distinctive examples when you are learning the workflow.

Optional Extension: Prompting from a Separate Reference Image

One optional demo is also worth mentioning: the prompt does not have to come from the same image as the target. The refer_image argument handles that case. This is closer to a real retrieval workflow, where you have one known example and want to search a different frame set or batch for similar objects.

If you want to include this extension, here is a minimal working example:

import numpy as np

from ultralytics import YOLOE
from ultralytics.models.yolo.yoloe import YOLOEVPSegPredictor
from ultralytics.utils import ASSETS

reference_model = YOLOE("yoloe-26s-seg.pt")
reference_model.set_classes(["person"])

reference_prompts = dict(
   bboxes=np.array([[221.52, 405.8, 344.98, 857.54]]),
   cls=np.array([0]),
)

reference_result = reference_model.predict(
   ASSETS / "zidane.jpg",
   refer_image=ASSETS / "bus.jpg",
   visual_prompts=reference_prompts,
   predictor=YOLOEVPSegPredictor,
   conf=0.10,
)[0]

reference_result.names = {0: "person"}
reference_result.show()

Treat this as an advanced extension, not a required part of the first pass through the lesson.

YOLOE-26 Prompt-Free Open-Vocabulary Object Detection

YOLOE-26 also has prompt-free variants. These models come with a built-in open vocabulary and do not require your own text prompts or visual prompts at inference time.

How Prompt-Free Inference Works in Practice

Here is the corresponding code:

from ultralytics import YOLOE
from ultralytics.utils import ASSETS

prompt_free_model = YOLOE("yoloe-26s-seg-pf.pt")
prompt_free_results = prompt_free_model.predict(ASSETS / "bus.jpg", conf=0.25)
prompt_free_results[0].show()

This mode is useful when you want a larger built-in vocabulary with minimal setup. According to the Ultralytics documentation, the prompt-free models use a built-in large vocabulary and internal embeddings for open-set recognition.

But there is a tradeoff, and it is important enough to say clearly: prompt-free convenience costs accuracy.

Prompt-free mode is appealing because it feels effortless. Load the model, run inference, and get a wide-vocabulary result without deciding on prompts first. In practice, it is best for exploratory analysis and quick discovery passes, not for the cases where you need the strongest possible precision on one narrow concept. If text prompting is you saying “look for this,” prompt-free mode is closer to saying “show me what the built-in vocabulary thinks is here.”

Where YOLOE-26 Beats YOLO26, and Where It Does Not

YOLOE-26 beats standard YOLO26 whenever the target category is dynamic, long-tail, or simply outside the fixed closed-set label space. If your environment changes often, or if you need to search for unusual objects without retraining a detector from scratch, YOLOE-26 is the more flexible tool.

It also has a strong headline benchmark story. Ultralytics reports that the x model reaches:

• 40.6 AP: on LVIS minival with text prompts
• 38.5 AP: with visual prompts
• 31.1 AP: in the prompt-free non-E2E setting

That is a meaningful spread. Prompt-free mode is convenient, but you pay for that convenience in accuracy.

The pattern is straightforward: the more explicit guidance you give YOLOE-26, the better it performs. Prompt-free mode is still useful, but it should be treated as a convenience mode, not the default path when accuracy matters most.

YOLO26 and YOLOE-26 at a Glance

A Practical Decision Framework

At the same time, YOLOE-26 is not automatically “better” than YOLO26 for every production workload. If your class list is fixed, your latency budget is tight, and you already know exactly what you need to detect, a standard fine-tuned YOLO26 model is often the simpler production artifact.

Here is a clean decision framework:

• YOLO26: choose when the classes are fixed and you care most about a lean production artifact
• YOLOE-26 text prompting: choose when you know the concept but want the flexibility to change classes without retraining
• YOLOE-26 visual prompting: choose when the object is easier to show than to describe
• YOLOE-26 prompt-free: choose when you want broad exploratory coverage and are willing to trade accuracy for convenience

This decision framework is exactly why Lesson 1 comes before Lesson 2. You need to understand where YOLOE-26 shines before you can use it intelligently as a labeling engine later.

Common Failure Modes and How to Debug Them

Open-vocabulary detection is powerful, but it is not magic. Here are the most common reasons your first run may disappoint you:

The Prompt Is Too Broad

A prompt (e.g., tool or electronics device) may be semantically valid but visually broad. The model has too many ways to satisfy the prompt, so detections may become noisy.

Fix: start narrower. Use the most concrete noun you can.

The Prompt Is Too Obscure

Sometimes the opposite happens. You use a very domain-specific term that the model has weak grounding for.

Fix: try a simpler synonym or a more common parent concept first.

The Object Is Too Small

Tiny objects are hard for almost every detector. Open-vocabulary capability does not remove that challenge.

Fix: use a larger input size, crop the image, or test on closer examples before judging the prompt.

The Visual Prompt Box Is Poor

If your reference box includes too much background or cuts off the object, you are teaching the model the wrong visual concept.

Fix: redraw the prompt box tightly and try again.

You Are Expecting Prompt-Free Mode to Behave Like Curated Prompting

Prompt-free is convenient, not optimal. If the results are underwhelming, that does not mean YOLOE-26 as a whole is weak. It may simply mean you should move to text or visual prompting.

Fix: switch to a more controlled prompt mode before drawing conclusions.

Your Prompts Overlap Semantically

Some prompt sets are too close to each other. If you ask for overlapping concepts (e.g., person, pedestrian, and worker, or tool, hand tool, and screwdriver), the detector may produce messy or merged behavior because multiple prompts can plausibly match the same region.

Fix: start with clearly separated categories. Add finer-grained prompts only after the coarse categories are behaving the way you expect.

The right question is not “Did one random prompt work perfectly on the first try?” It is “How much controllability do we get, and how quickly can we improve the result through prompting?”

One Deployment Detail You Should Not Miss

One operational detail from the Ultralytics docs deserves explicit attention. When you export a YOLOE model, the configured classes are baked into the exported weights. After that, you cannot keep swapping prompt classes on the exported artifact. To change them, you need to re-export from the original checkpoint.

During experimentation, this is not a big deal. You just change the prompts and rerun the code. But if you are preparing a deployable artifact, the prompt configuration becomes part of what you are exporting.

This is one of the clearest ways to understand the line between experimentation and production:

• exploration mode: YOLOE-26 is highly flexible
• deployment mode: some of that flexibility gets frozen into the exported artifact

That detail will matter a lot in Lesson 2, because Lesson 2 is about what happens when you stop exploring and start converging on a narrower detector workflow.

What's next? We recommend PyImageSearch University.

Course information:
120+ total classes • 115+ hours of on-demand code walkthrough videos • Last updated: October 2026
★★★★★ 4.84 (128 Ratings) • 16,000+ Students Enrolled

I strongly believe that if you had the right teacher you could master computer vision and deep learning.

Do you think learning computer vision and deep learning has to be time-consuming, overwhelming, and complicated? Or has to involve complex mathematics and equations? Or requires a degree in computer science?

That’s not the case.

All you need to master computer vision and deep learning is for someone to explain things to you in simple, intuitive terms. And that’s exactly what I do. My mission is to change education and how complex Artificial Intelligence topics are taught.

If you're serious about learning computer vision, your next stop should be PyImageSearch University, the most comprehensive computer vision, deep learning, and OpenCV course online today. Here you’ll learn how to successfully and confidently apply computer vision to your work, research, and projects. Join me in computer vision mastery.

Inside PyImageSearch University you'll find:

• ✓ 120+ courses on essential computer vision, deep learning, and OpenCV topics
• ✓ 94+ Certificates of Completion
• ✓ 115+ hours of on-demand video
• ✓ Brand new courses released regularly, ensuring you can keep up with state-of-the-art techniques
• ✓ Pre-configured Jupyter Notebooks in Google Colab
• ✓ Run all code examples in your web browser — works on Windows, macOS, and Linux (no dev environment configuration required!)
• ✓ Access to centralized code repos for all 540+ tutorials on PyImageSearch
• ✓ Easy one-click downloads for code, datasets, pre-trained models, etc.
• ✓ Access on mobile, laptop, desktop, etc.

Click here to join PyImageSearch University

Summary

The easiest way to think about this lesson is:

• YOLO26 is your fast closed-set baseline.
• YOLOE introduced promptable open-vocabulary detection to the YOLO family.
• YOLOE-26 brings that same idea into the YOLO26 generation with stronger open-vocabulary performance and the same real-time mindset.

By now, you have seen 3 concrete workflows:

• Text prompting: inference with set_classes()
• Visual prompting: inference with reference boxes
• Prompt-free inference: inference with a built-in vocabulary

You have also seen the deeper point behind those workflows. YOLOE-26 changes how quickly we can move from “we have an object in mind” to “we have a detector doing something useful.”

That is already enough to start experimenting. But it also raises the next practical question.

What if you do not want a permanently flexible open-vocabulary detector? What if you want to use YOLOE-26 to discover and pre-label a niche category, then turn that into a small, fast, deployable closed-set detector?

That is exactly what we will do in Lesson 2.

Citation Information

Singh, V. “YOLO26 Open-Vocabulary Object Detection with YOLOE-26,” PyImageSearch, S. Huot, A. Sharma, and P. Thakur, eds., 2026, https://pyimg.co/mzpx3

@incollection{Singh_2026_yolo26-open-vocabulary-object-detection-yoloe-26,
  author = {Vikram Singh},
  title = {{YOLO26 Open-Vocabulary Object Detection with YOLOE-26}},
  booktitle = {PyImageSearch},
  editor = {Susan Huot and Aditya Sharma and Piyush Thakur},
  year = {2026},
  url = {https://pyimg.co/mzpx3},
}

To download the source code to this post (and be notified when future tutorials are published here on PyImageSearch), simply enter your email address in the form below!

Download the Source Code and FREE 17-page Resource Guide

Enter your email address below to get a .zip of the code and a FREE 17-page Resource Guide on Computer Vision, OpenCV, and Deep Learning. Inside you'll find my hand-picked tutorials, books, courses, and libraries to help you master CV and DL!

Download the code!

Website

The post YOLO26 Open-Vocabulary Object Detection with YOLOE-26 appeared first on PyImageSearch.

Make a Chrome Extension to Digest Webpages with Manifest V3 and Groq API

In this lesson, you will learn how to build a Chrome extension that can read the current webpage, send that content to a large language model (LLM), and stream the answer back into a popup interface.

To learn how to build a Chrome extension that summarizes webpages with Manifest V3 and the Groq API, just keep reading.

Looking for the source code to this post?

Jump Right To The Downloads Section

This is exactly the kind of project we want for a first lesson. The codebase is small enough to understand in one sitting, but practical enough to expose the moving parts that actually matter in a real browser-based artificial intelligence (AI) tool:

• Manifest V3 configuration
• Extension permissions
• Popup user interface (UI) structure
• Persistent storage with chrome.storage.local
• Reading live page content with chrome.scripting.executeScript
• Calling the Groq API with fetch
• Streaming tokens back into the interface
• Managing conversation history across turns

By the end of this lesson, you will understand what each file does, how the pieces fit together, and why this design works so well for a first Chrome LLM extension.

Meet the Project

Our project is called Page Intelligence. When the user clicks the extension icon, Chrome opens a popup. From that popup, the user can:

• Save a Groq application programming interface (API) key
• Summarize the current webpage
• Ask follow-up questions about that page
• Clear the current conversation and start over

At first glance, that workflow feels straightforward. Under the hood, though, it touches 3 separate environments:

• The popup UI, rendered by Chrome as an extension page
• The active tab, which contains the webpage we want to read
• The Groq API, which generates the response

That separation is the key architectural idea in this lesson. The popup cannot directly read the page Document Object Model (DOM) just because both are open in the same browser window. Instead, it has to ask Chrome for permission and inject code into the active tab at the right moment.

Would you like immediate access to 3,457 images curated and labeled with hand gestures to train, explore, and experiment with … for free? Head over to Roboflow and get a free account to grab these hand gesture images.

Configuring Your Development Environment

To follow this lesson, you only need Google Chrome, a Groq API key, and the project files for the extension. Since this is a lightweight Chrome extension project, there are no Python packages or JavaScript dependencies to install.

Then follow these steps:

• Download or open the project folder.
• Get a Groq API key.
• Open chrome://extensions.
• Enable Developer mode.
• Click Load unpacked and select the project folder.

Need Help Configuring Your Development Environment?

All that said, are you:

• Short on time?
• Learning on your employer’s administratively locked system?
• Wanting to skip the hassle of fighting with the command line, package managers, and virtual environments?
• Ready to run the code immediately on your Windows, macOS, or Linux system?

Then join PyImageSearch University today!

Gain access to Jupyter Notebooks for this tutorial and other PyImageSearch guides pre-configured to run on Google Colab’s ecosystem right in your web browser! No installation required.

And best of all, these Jupyter Notebooks will run on Windows, macOS, and Linux!

Project Structure

We first need to review our project directory structure.

Start by accessing this tutorial’s “Downloads” section to retrieve the source code.

From there, take a look at the directory structure:

chrome-llm-extension/
├── manifest.json
├── popup.html
├── popup.js
├── styles.css
└── README.md

The minimal layout is intentional. It keeps the lesson focused. There is no framework, no extra build system, and no separate content script file to chase through. All of the user-facing logic lives in the popup, and the extension reaches into the current page only when it needs to read content.

What Each File Is Responsible For

manifest.json

• Declares that this is a Manifest V3 extension
• Requests the permissions needed to read the active tab, inject a script, and store the API key
• Registers popup.html as the default popup

popup.html

• Defines the structure of the popup interface
• Creates the API key panel, chat area, quick action buttons, and message box

styles.css

• Controls layout, spacing, colors, bubble styles, and the overall feel of the popup

popup.js

• Contains the extension’s runtime logic
• Loads and stores the API key
• Reads page content from the active tab
• Calls the Groq API
• Streams the response token by token
• Maintains chatHistory

Why This Structure Works Well for a Lesson

For a first Chrome extension project, fewer files usually mean faster understanding. Instead of bouncing between a popup, a background worker, and multiple content scripts, you get to follow one direct path:

• Chrome opens the popup.
• The popup loads saved state.
• The popup requests page text when needed.
• The popup sends a streaming LLM request.
• The popup updates the UI in real time.

For a teaching project, that is a great tradeoff. We stay practical without burying the learner in scaffolding.

The Chrome Extension Mental Model

Before we go any further, let us make the execution model crystal clear.

A popup is not the webpage

When you click a Chrome extension icon, Chrome opens a small Hypertext Markup Language (HTML) page that belongs to the extension. That page has its own DOM, JavaScript context, and permissions.

This means:

• popup.js can manipulate elements in popup.html
• popup.js cannot directly access document.body from the currently open website
• To read the website, the extension must ask Chrome to run code inside the active tab

This is exactly why chrome.scripting.executeScript matters in this project.

Storage Belongs to the Extension, Not the Site

There is a second separation we need to keep in mind, and that is storage.

If you used localStorage inside a normal webpage, that data would belong to that page’s origin. Here, we want the API key to belong to the extension itself, regardless of which site the user is visiting.

For that reason, the code uses:

chrome.storage.local

This storage area is managed by Chrome and shared across extension contexts.

Network Requests Must Be Explicitly Allowed

Extensions are permission-driven. If you want to send requests to the Groq API, you declare that in manifest.json using host_permissions.

This is one of the core Chrome extension design principles:

• UI and code live inside the extension
• Privileges are declared in the manifest
• Cross-context access is granted only through the proper APIs

Walking Through manifest.json

Open manifest.json, and you will find the contract between your extension and Chrome.

Here is the full file:

{
  "manifest_version": 3,
  "name": "Page Intelligence",
  "version": "1.0",
  "description": "Chat with any webpage using an LLM. Built for the Agent AI course.",
  "permissions": ["activeTab", "scripting", "storage"],
  "host_permissions": ["https://api.groq.com/*"],
  "action": {
    "default_popup": "popup.html",
    "default_title": "Page Intelligence"
  }
}

manifest_version: 3

Chrome expects modern extensions to use Manifest V3. That affects how extensions are structured, how scripts are injected, and how background logic is handled.

In this project, Manifest V3 gives us a clean, current, production-relevant starting point.

Why These Permissions Matter

activeTab

This permission allows the extension to interact with the tab the user is currently using. In a tool like this, that is essential because the extension must read the active page when the user clicks Summarize Page.

scripting

This permission unlocks chrome.scripting.executeScript, which is the modern Manifest V3 (MV3) way to inject code into a page.

storage

This permission lets the extension persist the Groq API key using chrome.storage.local.

Why host_permissions Matters

The extension calls https://api.groq.com/openai/v1/chat/completions.

Without the proper host permission, that request would be blocked. This line gives the extension permission to communicate with the Groq API host.

Why action.default_popup Matters

Here we define the entry point for the user experience. This tells Chrome, “when the extension icon is clicked, open popup.html.”

One line is all it takes to wire the entire interface into the browser.

Understanding popup.html

Next, open popup.html. If manifest.json is the contract, popup.html is the stage where the user interacts with the extension.

The Popup Is Intentionally Split into Clear UI Regions

The HTML defines 5 major regions:

• Header
• API key section
• Chat history
• Quick actions
• Input area

That separation matters because each region maps directly to behavior in popup.js.

Here is the structural core of the file:

<div class="container">
  <div class="header">...</div>
  <div id="api-key-section" class="api-key-section">...</div>
  <div id="chat-history" class="chat-history">...</div>
  <div class="quick-actions">...</div>
  <div class="input-area">...</div>
</div>

The Header Sets Up Identity and Settings Access

At the top, we have:

<h1>Page Intelligence</h1>
<button id="settings-btn" class="icon-btn">...</button>

The lesson here is simple but important. Even a tiny extension UI needs clear affordances. The settings button gives the user a predictable place to reopen the API key panel after the initial save.

The API Key Section Introduces Stateful UI

The API key input exists in the DOM from the beginning, but JavaScript decides whether it should be visible based on whether a key has already been saved.

A common frontend pattern is at work here:

• Keep the structure in HTML
• Let JavaScript decide the current state

The Chat Area Is the Main Feedback Surface

The #chat-history container starts with a welcome message:

<div class="welcome-msg">
  Ask me anything about this page,<br />
  or click <strong>Summarize</strong> to get started.
</div>

This does more than decorate the popup. It solves the empty-state problem. Before the first interaction, the popup still feels guided and purposeful.

Quick Actions Reduce Friction

The 2 action buttons are:

• Summarize Page
• Clear

These are excellent teaching examples because they show 2 different UI intents:

• Task button: launches a multi-step workflow
• Reset button: clears local state

The Input Area Supports Chat-Like Behavior

The textarea plus send button give the user a second interaction path. They can either:

• Click the guided summary flow first
• Skip directly to a custom question

Practically, this supports both beginners and more confident users.

Understanding styles.css

Now open styles.css. This is where the extension begins to feel polished instead of merely functional.

The Layout Uses a Simple But Effective Flex Column

The 2 most important layout rules are:

body {
  width: 400px;
}

.container {
  display: flex;
  flex-direction: column;
  height: 560px;
}

This gives the popup a fixed footprint and lets the chat area grow while the header, actions, and input stay anchored.

Why the Chat Panel Works Well

The chat container uses:

.chat-history {
  flex: 1;
  overflow-y: auto;
  display: flex;
  flex-direction: column;
  gap: 10px;
}

You will see this pattern often in messaging interfaces:

• flex: 1: lets the panel absorb remaining height
• overflow-y: auto: makes long conversations scrollable
• gap: keeps messages visually separated without hard-to-maintain margins

Message Bubbles Encode Role Visually

The Cascading Style Sheets (CSS) differentiate:

• .message.user
• .message.assistant
• .message.error
• .message.system

This is not just cosmetic. Good styling teaches the user how to read the conversation. With one glance, they can tell which text came from them, which came from the model, and whether a message is informational or an error.

Small Animation Helps the Interface Feel Alive

The fadeIn keyframe is subtle:

@keyframes fadeIn {
  from { opacity: 0; transform: translateY(4px); }
  to   { opacity: 1; transform: translateY(0); }
}

It is also a good reminder that user experience (UX) is part of engineering. Because the model response streams in incrementally, the interface benefits from motion that feels lightweight and responsive.

Input and Actions Are Styled for Fast Iteration

The quick action buttons and textarea focus states give the extension a clean, modern feel without pulling in a UI framework.

This leads to another useful lesson from the project:

• HTML: handles structure
• CSS: handles presentation
• JavaScript: handles behavior

For a beginner-friendly extension, that separation is exactly what we want.

Reading README.md the Right Way

It is easy to ignore README.md, but in a teaching repository, that file matters.

The README Is the Learner’s Runway

This project’s README does 3 useful jobs:

• Explains: what concepts the extension teaches
• Shows: setup steps for getting a Groq API key and loading the extension
• Presents: the codebase structure in a fast, approachable way

So while README.md is not runtime code, it is still part of the learning architecture of the project.

Why This Matters in Real Projects

When you teach or ship developer tooling, code alone is not enough. Learners need:

• Context
• Setup steps
• A map of what to read first

This repository already does a good job of that by aligning the README sections with the files learners will inspect next.

Walking Through popup.js

With the surrounding files covered, we can move to the heart of the extension.

popup.js is where the extension:

• Loads saved state
• Handles button clicks
• Reads the active page
• Calls the Groq API
• Streams the assistant output
• Updates the chat UI

This single file is small enough to follow in one sitting, which makes it ideal for a hands-on lesson.

Configuration and State Live at the Top

The file begins with configuration values:

const GROQ_API_URL = "https://api.groq.com/openai/v1/chat/completions";
const MODEL = "meta-llama/llama-4-scout-17b-16e-instruct";

Then it defines a reusable system prompt and 2 key pieces of runtime state:

let apiKey = "";
let chatHistory = [];

This pattern is worth teaching because it keeps the mental model clean:

• Constants describe how the app talks to the model
• Mutable state tracks what the user has done in this popup session

DOM References Create a Bridge from HTML to Behavior

The next block grabs the key UI elements:

const apiKeySection = document.getElementById("api-key-section");
const apiKeyInput   = document.getElementById("api-key-input");
const saveKeyBtn    = document.getElementById("save-key-btn");
const settingsBtn   = document.getElementById("settings-btn");
const chatHistoryEl = document.getElementById("chat-history");
const summarizeBtn  = document.getElementById("summarize-btn");
const clearBtn      = document.getElementById("clear-btn");
const userInput     = document.getElementById("user-input");
const sendBtn       = document.getElementById("send-btn");

This is why the element identifiers (IDs) in popup.html matter. They are the handles that let JavaScript attach behavior to markup.

DOMContentLoaded Restores Saved Extension State

When the popup opens, this listener runs:

document.addEventListener("DOMContentLoaded", async () => {
  const stored = await chrome.storage.local.get("groq_api_key");

  if (stored.groq_api_key) {
    apiKey = stored.groq_api_key;
    apiKeySection.style.display = "none";
  }
});

Here is one of the first places where theory and practice meet:

• The popup is ephemeral. It opens when the user clicks the extension icon.
• Any in-memory JavaScript state disappears when the popup closes.
• Persistent state must therefore live outside regular variables.

This is exactly why chrome.storage.local is essential here. It gives us persistence without forcing us to introduce a database or a heavier architecture.

Saving the API Key

When the user clicks Save, the extension:

• Reads the input value
• Validates that it is not empty
• Stores it with chrome.storage.local.set(...)
• Copies it into the in-memory apiKey variable
• Hides the API key panel
• Shows a confirmation message

What you are seeing is a tidy example of state synchronization between:

• The DOM
• Extension storage
• In-memory JavaScript state

The Settings Button Reopens the Key Panel

This small handler is worth pointing out:

settingsBtn.addEventListener("click", () => {
  apiKeySection.style.display =
    apiKeySection.style.display === "none" ? "flex" : "none";
});

Again, not every useful behavior needs a framework. For a focused popup, direct DOM manipulation is perfectly reasonable and easy to understand.

How the Extension Reads the Current Webpage

This is the most important architectural jump in the lesson, so take a moment to make sure the flow is clear before moving on.

getPageText() Starts by Locating the Active Tab

The function begins with:

const [tab] = await chrome.tabs.query({ active: true, currentWindow: true });

In plain English, this asks Chrome to return the tab the user is looking at right now.

The Extension Then Injects a Function into That Tab

Here is the core pattern:

const results = await chrome.scripting.executeScript({
  target: { tabId: tab.id },
  func: () => {
    const clone = document.body.cloneNode(true);
    clone.querySelectorAll("script, style, nav, footer, aside")
      .forEach(el => el.remove());

    return {
      title: document.title,
      url: window.location.href,
      text: clone.innerText.replace(/\s+/g, " ").trim().slice(0, 8000),
    };
  },
});

As a teaching example, this snippet is excellent because it shows the exact boundary between extension code and page code.

Why Clone the Page Body First

The code uses:

const clone = document.body.cloneNode(true);

Cloning first is smart because it lets the extension clean the content without altering the actual webpage the user is viewing.

Why Remove Script, Style, Nav, Footer, and Aside

This is a lightweight content-cleaning step. The extension wants the main readable text, not the surrounding noise.

This gives learners a practical introduction to preprocessing before sending data to an LLM.

Why Normalize Whitespace and Cap Text Length

This line matters:

clone.innerText.replace(/\s+/g, " ").trim().slice(0, 8000)

It does 3 useful things:

• Collapses repeated whitespace
• Trims leading and trailing space
• Limits the payload to 8,000 characters

The final limit is especially important. Even though modern models have large context windows, good engineering still means keeping inputs focused and predictable.

The Return Value Comes Back to the Popup

After executeScript finishes, the popup receives:

• The page title
• The page uniform resource locator (URL)
• Cleaned text content

This makes getPageText() the bridge from the browser page to an LLM-ready prompt.

How the Groq Streaming Call Works

Once the extension has page content, it needs to send a chat completion request and render the answer in real time.

streamResponse() Checks for the API Key First

The function opens with a guard clause:

if (!apiKey) {
  apiKeySection.style.display = "flex";
  showSystem("Please set your Groq API key first.", "error");
  return;
}

This is good defensive programming. Before the app makes a network request, it verifies that the minimum required state is present.

The Request Body Sends the Full Conversation

The fetch call includes:

body: JSON.stringify({
  model: MODEL,
  messages: chatHistory,
  stream: true,
  temperature: 0.7,
  max_tokens: 1024,
})

There are 2 especially important details here.

messages: chatHistory

This is how the model gets context. Instead of sending only the newest user message, the extension sends the full running conversation.

As a result, the model can handle:

• Follow-up questions
• Clarifications
• Multi-turn interaction grounded in prior turns

stream: true

This tells the API to return response content instead of waiting for the full answer to be finished.

That is what gives the extension its chat-like feel.

The Response Is Read as a Stream

The code then creates:

const reader  = response.body.getReader();
const decoder = new TextDecoder();

This is where the frontend side starts to feel more advanced. Instead of calling await response.json(), the extension reads raw chunks from the response body.

Server-Sent Events (SSE) Are Parsed Line by Line

The code looks for lines that begin with: data:

It then trims the prefix, checks for [DONE], and parses the JavaScript Object Notation (JSON) payload:

const parsed = JSON.parse(payload);
const token  = parsed.choices?.[0]?.delta?.content ?? "";

If content is present, the function sends it to onChunk(token).

Why the UI Feels Responsive

Every arriving content chunk updates the assistant bubble immediately. That means the user does not stare at a frozen popup waiting for a full paragraph to appear all at once.

There is a practical UX lesson here:

• Streaming improves perceived speed
• Real-time rendering makes the app feel more interactive
• Even a simple UI can feel polished if feedback arrives continuously

One Subtle Engineering Note

This implementation splits each decoded chunk by newline and skips JSON parsing errors when a partial fragment arrives at a chunk boundary.

For a teaching demo, that tradeoff is acceptable. In a more robust production version, you would usually maintain a rolling buffer so partial SSE lines can be reconstructed correctly before parsing.

That is a valuable lesson too: simple implementations are often ideal for learning, even when they are not yet industrial-strength.

How Conversation State Is Managed

A good chat interface is really a state-management problem wearing a friendly UI.

sendMessage() Handles the Main Interaction Loop

At a high level, sendMessage() does the following:

• Cleans the user input
• Pushes the user message into chatHistory
• Renders the user bubble
• Disables the input while the model responds
• Creates an empty assistant bubble
• Streams response content into that bubble
• Saves the final assistant response back into chatHistory
• Re-enables the input

This is the core interaction cycle of the whole extension.

Why Create an Empty Assistant Bubble First

This line is the trick:

const assistantBubble = appendMessage("assistant", "");

Instead of waiting for the final answer, the UI creates a placeholder bubble and fills it as response content arrive. That is what makes streaming visible to the user.

Auto-Scroll Keeps the Newest Response Content in View

During streaming, the code does:

chatHistoryEl.scrollTop = chatHistoryEl.scrollHeight;

This small line is easy to miss, but it matters. Without it, long responses would stream out of view and the interface would feel clumsy.

The System Prompt Is Injected When Needed

When the user clicks Send on a fresh chat, the code makes sure the system prompt is present:

if (chatHistory.length === 0) {
  chatHistory.push({ role: "system", content: SYSTEM_PROMPT });
}

This is a compact but important pattern. It ensures that every conversation begins with the assistant’s behavioral instructions, even if the user skips the summarize flow and goes straight to a custom question.

The Enter Key Is Tuned for Chat UX

This handler:

if (e.key === "Enter" && !e.shiftKey) {
  e.preventDefault();
  sendBtn.click();
}

gives the interface a familiar messaging behavior:

• Enter sends
• Shift + Enter can still be used for a multiline prompt

Small details like this make the project feel thoughtful.

How the Summarize Workflow Works

The Summarize Page button is where all the major ideas in the project come together.

The Button First Updates Its Own UI State

When clicked, it temporarily changes to: "Reading page..."

This gives the user immediate feedback that work has started.

The History Is Reset for a Page-First Conversation

Inside the click handler, the code sets:

chatHistory = [
  { role: "system", content: SYSTEM_PROMPT },
];

This is a deliberate design decision. The summarize flow creates a fresh conversation centered on the current page.

The Page Content Is Converted into a User Message

The extension builds a message like this:

const userMessage =
  `Please summarise the following webpage.\n\n` +
  `Title: ${page.title}\nURL: ${page.url}\n\nContent:\n${page.text}`;

This is a clever and practical technique. Instead of inventing a separate prompt schema, the extension simply turns page context into a normal user message. That means the rest of the chat pipeline can stay unchanged.

Why This Design Is Elegant

The same sendMessage() function handles:

• Free-form user questions
• A page summary request

That reduces duplicated logic and keeps the codebase teachable.

How the Clear Button Resets the Interface

The clear handler is short, but it teaches an important UI principle.

Reset Both State and Visible Output

When the user clicks Clear, the code:

• Empties chatHistory
• Restores the welcome message HTML inside #chat-history

This gives the user a clean slate both logically and visually.

In many small projects, bugs appear because developers reset one layer but forget the other. This extension avoids that trap by resetting both.

End-to-End Flow, From Click to Answer

Let us put the whole lesson together in one concrete sequence.

What Happens When the User Clicks Summarize Page

• Chrome opens the extension popup.
• popup.js loads and restores any saved API key from chrome.storage.local.
• The user clicks Summarize Page.
• getPageText() asks Chrome for the active tab.
• Chrome injects a function into that page using chrome.scripting.executeScript.
• The injected function clones the page body, removes noisy elements, and returns title, URL, and cleaned text.
• The popup resets chatHistory with the system prompt.
• The popup creates a user message containing the page content.
• sendMessage() renders the user bubble and creates an empty assistant bubble.
• streamResponse() sends the chat request to Groq with stream: true.
• The popup reads SSE chunks, extracts response content, and updates the assistant bubble live.
• When streaming ends, the full assistant message is saved into chatHistory.

That is the full extension lifecycle in one pass.

What Happens on the Next Follow-Up Question

The second turn is even more interesting:

• The user asks a follow-up question in the textarea.
• sendMessage() appends that new message to the existing chatHistory.
• The full conversation is sent again to the model.
• The model now has access to both the page context and the prior summary.

That is what makes the extension feel conversational instead of one-shot.

Practical Engineering Takeaways

This extension may be small, but it teaches several patterns that show up in larger AI products too.

Pattern 1: Retrieve First, Generate Second

The extension does not ask the model to guess what is on the page. It retrieves the page content first, then sends that content to the model.

The retrieval step is simple, but it is foundational.

Pattern 2: Separate UI State from Persistent State

The API key lives in chrome.storage.local.

The active conversation lives in memory as chatHistory.

This is a healthy separation because:

• Long-lived secrets survive popup closures
• Short-lived conversations reset naturally when the session changes

Pattern 3: Stream Whenever Responsiveness Matters

Streaming is not just a flashy feature. It changes how the user experiences latency.

Even if the total response time stays similar, streaming makes the system feel faster and more alive.

Pattern 4: Keep the First Version Intentionally Small

This extension could have included:

• A background service worker
• Conversation persistence across popup sessions
• Rich markdown rendering
• Better page extraction heuristics
• More robust streaming buffering

But for a lesson, the current scope is exactly right. It is complete enough to be useful and small enough to fully understand.

Where You Could Take This Project Next

Once learners understand this version, there are several natural improvements worth exploring.

Improve Content Extraction

Right now, the extension removes a few noisy tags and then uses innerText.

A stronger version could:

• Prefer main or article-like containers when present
• Skip repeated sidebar content more aggressively
• Chunk long pages instead of truncating them at 8,000 characters

Improve Streaming Robustness

The current SSE parsing is easy to understand, which is excellent for a first lesson. A more production-oriented version could keep a line buffer across chunk boundaries before parsing JSON.

Improve Conversation Persistence

chatHistory currently lives only in memory. If you close the popup, the conversation disappears.

That is fine for a lesson, but a future version could store conversations in chrome.storage.local or IndexedDB.

Improve Prompting

The current system prompt is short and sensible. A later lesson could show how to:

• Ask for citation-style answers grounded in extracted text
• Detect when the page content is too thin
• Adapt response style for summarization versus question answering

What's next? We recommend PyImageSearch University.

Course information:
120+ total classes • 115+ hours of on-demand code walkthrough videos • Last updated: October 2026
★★★★★ 4.84 (128 Ratings) • 16,000+ Students Enrolled

I strongly believe that if you had the right teacher you could master computer vision and deep learning.

Do you think learning computer vision and deep learning has to be time-consuming, overwhelming, and complicated? Or has to involve complex mathematics and equations? Or requires a degree in computer science?

That’s not the case.

All you need to master computer vision and deep learning is for someone to explain things to you in simple, intuitive terms. And that’s exactly what I do. My mission is to change education and how complex Artificial Intelligence topics are taught.

If you're serious about learning computer vision, your next stop should be PyImageSearch University, the most comprehensive computer vision, deep learning, and OpenCV course online today. Here you’ll learn how to successfully and confidently apply computer vision to your work, research, and projects. Join me in computer vision mastery.

Inside PyImageSearch University you'll find:

• ✓ 120+ courses on essential computer vision, deep learning, and OpenCV topics
• ✓ 94+ Certificates of Completion
• ✓ 115+ hours of on-demand video
• ✓ Brand new courses released regularly, ensuring you can keep up with state-of-the-art techniques
• ✓ Pre-configured Jupyter Notebooks in Google Colab
• ✓ Run all code examples in your web browser — works on Windows, macOS, and Linux (no dev environment configuration required!)
• ✓ Access to centralized code repos for all 540+ tutorials on PyImageSearch
• ✓ Easy one-click downloads for code, datasets, pre-trained models, etc.
• ✓ Access on mobile, laptop, desktop, etc.

Click here to join PyImageSearch University

Summary

In this lesson, you saw how a compact Chrome extension can combine browser APIs and LLM APIs into a practical, real-time workflow.

You learned how:

• manifest.json: declares the extension’s privileges and popup entry point
• popup.html: structures the user interface
• styles.css: makes the UI readable and responsive
• README.md: supports learning and onboarding
• popup.js: ties everything together through storage, page extraction, streaming, and chat state

Most importantly, you learned the core architectural idea behind the project:

• read the current page with the proper Chrome API
• turn that content into model-ready context
• stream the answer back into a lightweight interface

This pattern is simple, useful, and widely applicable. Once you understand it here, you can reuse it in richer browser tools, agent workflows, and AI-powered productivity extensions.

Citation Information

Singh, V. “Make a Chrome Extension to Digest Webpages with Manifest V3 and Groq API,” PyImageSearch, S. Huot, A. Sharma, and P. Thakur, eds., 2026, https://pyimg.co/8q9l7

@incollection{Singh_2026_make-chrome-extension-digest-webpages-manifest-v3-groq-api,
  author = {Vikram Singh},
  title = {{Make a Chrome Extension to Digest Webpages with Manifest V3 and Groq API}},
  booktitle = {PyImageSearch},
  editor = {Susan Huot and Aditya Sharma and Piyush Thakur},
  year = {2026},
  url = {https://pyimg.co/8q9l7},
}

To download the source code to this post (and be notified when future tutorials are published here on PyImageSearch), simply enter your email address in the form below!

Download the Source Code and FREE 17-page Resource Guide

Enter your email address below to get a .zip of the code and a FREE 17-page Resource Guide on Computer Vision, OpenCV, and Deep Learning. Inside you'll find my hand-picked tutorials, books, courses, and libraries to help you master CV and DL!

Download the code!

Website

The post Make a Chrome Extension to Digest Webpages with Manifest V3 and Groq API appeared first on PyImageSearch.

Scaling, Optimizing, and Exporting Transformers with PyTorch Lightning

In this lesson, you will learn how to scale, optimize, and export Transformer models using PyTorch Lightning, from enabling mixed precision and distributed training to generating ONNX and TorchScript exports ready for production inference. You will see how Lightning’s configuration-driven design, combined with Hydra, lets you take the exact same training code from Lesson 1 and push it into a high-performance, deployment-friendly workflow.

This lesson is the last in a 2-part series on PyTorch Lightning:

1. Training with PyTorch Lightning: Structured MLOps Development
2. Scaling, Optimizing, and Exporting Transformers with PyTorch Lightning (this tutorial)

To learn how to scale, optimize, and export Transformer models with PyTorch Lightning, just keep reading.

Looking for the source code to this post?

Jump Right To The Downloads Section

Introduction to Scaling PyTorch Lightning Transformer Training

Scaling Transformer training and preparing models for real-world deployment typically requires complex engineering effort such as distributed training, mixed-precision optimization, and multiple export formats. But thanks to PyTorch Lightning and Hydra, you can achieve all of this without rewriting your codebase or adding low-level boilerplate. In this lesson, we will take the sentiment-classification project you built earlier and upgrade it into a fast, scalable, and deployment-ready workflow.

What We Built Previously

In the previous tutorial, you created a complete sentiment classification pipeline using DistilBERT.

You structured your project with:

• A LightningDataModule to handle dataset loading, tokenization, and data loaders
• A LightningModule containing the model, forward pass, loss, metrics, and optimizer
• A clean Hydra configuration hierarchy that controlled every part of the training run
• A flexible training script that handled seeding, callbacks, logging, and TensorBoard
• An inference script for single, batch, and interactive predictions

This foundation gave you a fully reproducible training workflow that runs the same way on CPU, GPU, or MPS, with settings that can be overridden easily using Hydra’s command-line syntax.

However, training a base model is only the first step in a production ML workflow.

Why Scaling and Exporting Matter

Real-world Transformer workloads often push beyond what a single GPU or even a single machine can handle.

Teams need:

• Mixed precision (FP16 or BF16): for faster training and lower memory usage
• Distributed data parallel training: to use multiple GPUs efficiently
• Gradient accumulation: to simulate large batch sizes
• Model exports (ONNX or TorchScript): for deployment in production systems
• Lightweight inference runtimes: faster than PyTorch eager mode
• Stable artifact folders: ready for DVC, CI/CD, or cloud deployment pipelines

These features turn a research-grade script into an MLOps-grade training system, and that is exactly what you will build in this lesson.

How Lightning Makes It Easy

The best part is that none of this requires rewriting your training loop.

PyTorch Lightning abstracts the hard parts:

• Mixed precision: uses a single configuration key
• Distributed training (DDP or FSDP): requires only a strategy change
• Export logic: can be cleanly added without modifying model code
• Hydra: lets you switch between configurations instantly
• Checkpoints, logs, and metrics: remain fully reproducible

Instead of modifying your DataModule or LightningModule, you will extend the trainer configuration and add a dedicated export script.

That is the power of Lightning: the same codebase now supports single-GPU training, multi-GPU scaling, and deployment-ready exports with minimal changes.

Would you like immediate access to 3,457 images curated and labeled with hand gestures to train, explore, and experiment with … for free? Head over to Roboflow and get a free account to grab these hand gesture images.

Configuring Your Development Environment

Lesson 2 introduces new capabilities (e.g., distributed training, mixed precision, and exporting models to ONNX or TorchScript), so the environment now includes additional dependencies specifically meant for scaling and production-grade inference.

Below is the exact requirements.txt used in this lesson:

# Lesson 2 Requirements
# Additional dependencies for production features

# PyTorch and Lightning
torch>=2.1.0
torchvision>=0.16.0
torchaudio>=2.1.0
pytorch-lightning>=2.1.0

# Hugging Face ecosystem
transformers>=4.35.0
datasets>=2.14.0
tokenizers>=0.15.0

# Configuration management
hydra-core>=1.3.0
omegaconf>=2.3.0

# Metrics and monitoring
torchmetrics>=1.2.0
tensorboard>=2.15.0

# Model export (LESSON 2 specific)
onnx>=1.15.0
onnxruntime>=1.16.0  # CPU version works on all platforms

# Data processing
numpy>=1.24.0
pandas>=2.0.0

# Utilities
tqdm>=4.66.0
pyyaml>=6.0.0

# Note: For NVIDIA GPUs on Linux, you can optionally install:
# onnxruntime-gpu>=1.16.0  # Requires CUDA, Linux only

​​This environment enables 3 major Lesson 2 features:

1. Multi-GPU Distributed Training

Powered by:

• PyTorch Lightning (DDP or FSDP strategies)
• Hydra configurations for trainer selection

2. Mixed Precision

torch >= 2.x unlocks CPU BF16, GPU AMP, and faster matrix kernels.

3. ONNX Export and Runtime Inference

The addition of onnx and onnxruntime enables:

• model export to ONNX format
• cross-platform CPU inference
• benchmarking ONNX vs PyTorch (optional)

If you are using an NVIDIA GPU on Linux, you can optionally install:

pip install onnxruntime-gpu

This is not required for the lesson (we stay framework-agnostic), but readers who want GPU ONNX inference can enable it easily.

Need Help Configuring Your Development Environment?

All that said, are you:

• Short on time?
• Learning on your employer’s administratively locked system?
• Wanting to skip the hassle of fighting with the command line, package managers, and virtual environments?
• Ready to run the code immediately on your Windows, macOS, or Linux system?

Then join PyImageSearch University today!

Gain access to Jupyter Notebooks for this tutorial and other PyImageSearch guides pre-configured to run on Google Colab’s ecosystem right in your web browser! No installation required.

And best of all, these Jupyter Notebooks will run on Windows, macOS, and Linux!

Preparing PyTorch Lightning Models for Scalable Multi-GPU Training

Before we scale training, enable DDP or FSDP, or export models to ONNX or TorchScript, it is important that the project structure and configuration layout are set up correctly. Lesson 2 builds directly on the modular design from Lesson 1, but introduces specialized trainer configs and improved environment setup that unlock multi-GPU training and optimized inference.

Your updated project structure now reflects these goals.

Project Structure Refresher (Updated for Lesson 2)

Below is the exact directory layout used in this lesson:

.
├── configs
│   ├── config.yaml
│   ├── data
│   │   └── imdb.yaml
│   ├── model
│   │   └── distilbert.yaml
│   └── trainer
│       ├── ddp.yaml
│       ├── fsdp.yaml
│       └── single_gpu.yaml
├── README.md
├── requirements.txt
├── RUN.md
├── sample_reviews.txt
└── src
    ├── data_module.py
    ├── inference.py
    ├── model_module.py
    └── train.py

This layout introduces 3 major upgrades compared to Lesson 1:

Dedicated Trainer Configurations

You now have separate Hydra configurations under configs/trainer/ for:

• single_gpu.yaml: standard 1-GPU or CPU training
• ddp.yaml: multi-GPU Distributed Data Parallel (DDP)
• fsdp.yaml: Fully Sharded Data Parallel (memory-efficient large-model training)

This keeps the scaling logic completely outside the Python code (a major MLOps advantage).

Unified Root Configuration (config.yaml)

The root configuration composes the model, data, and trainer settings into a single experiment specification.

Switching between single-GPU and DDP training is as simple as:

python src/train.py trainer=ddp

No code changes or additional flags are required. By default, the project uses the DDP trainer (trainer=ddp), but you can switch to single-GPU training by running python src/train.py trainer=single_gpu.

Clean Separation of Code Files

Your src/ folder remains identical to Lesson 1 (data_module.py, model_module.py, train.py, and inference.py). This reinforces the Lesson 2 philosophy:

“Scaling and exporting should not require modifying your model code.”

Only the configs/ files change, not the implementation.

Hydra Enhancements for Scaling

Lesson 2 is where Hydra truly shines.

Instead of hardcoding distributed training logic in Python, you now have clean trainer profiles:

configs/trainer/
│── single_gpu.yaml
│── ddp.yaml
└── fsdp.yaml

Each YAML file defines:

• accelerator
• devices
• strategy (ddp or fsdp)
• precision
• logging settings

Examples a reader will use later:

python src/train.py trainer=single_gpu
python src/train.py trainer=ddp
python src/train.py trainer=fsdp

Hydra swaps in the right configuration automatically.

This is exactly what you expect in a real MLOps pipeline.

Why This Setup Matters (Before We Scale)

This lightweight structure unlocks all the heavy features coming next:

• Training: can scale from CPU to GPU to multi-GPU without touching code
• Export pipelines (ONNX or TorchScript): work consistently because configs capture preprocessing
• FSDP and mixed precision: require stable config-driven Trainer settings
• Model reproducibility: is guaranteed across environments

Most importantly:

You now have a production-grade pattern:

• Model logic: stays in model_module.py
• Data logic: stays in data_module.py
• Training logic: stays in train.py
• Scaling and exporting logic: stays in configs/*

This is the exact separation used in serious MLOps workflows.

Revisiting the Code Architecture

Before we introduce mixed precision, distributed training, gradient accumulation, and model export, it is important to ground ourselves in the architecture we built in the previous tutorial. Lesson 1 gave us a clean, modular foundation. Lesson 2 builds directly on this foundation without changing most of the code. That is the real strength of PyTorch Lightning and Hydra: scaling does not require rewriting your project.

DataModule, LightningModule, and Hydra Recap

Our project continues to rely on the same 3 core abstractions:

LightningDataModule

Handles all data concerns:

• Dataset: downloading the IMDB dataset
• Tokenization: using a Hugging Face tokenizer
• Data loaders: creating training, validation, and test DataLoaders
• Separation of concerns: keeping preprocessing separate from training logic

This structure remains untouched in Lesson 2. AMP, DDP, and export do not require modifying data_module.py.

LightningModule

Encapsulates the model architecture and training logic:

• Encoder: DistilBERT
• Forward pass: processing model inputs
• Loss and metrics: computing training loss and evaluation metrics
• Optimizer: configuring optimization

Lesson 2 also does not modify the internal model logic.

Instead, we scale training through the Trainer and Hydra configuration.

Hydra

Hydra remains the central controller of the entire workflow.

It composes the configuration from:

• configs/model/*.yaml
• configs/data/*.yaml
• configs/trainer/*.yaml

Lesson 2 extends Hydra with:

• DDP trainer configurations
• FSDP configurations (if desired)
• AMP configurations
• gradient accumulation
• export configurations

But the training code (train.py) stays nearly the same.

This consistent architecture is what makes the next steps (i.e., scaling, exporting, and optimizing) feel incremental rather than overwhelming.

Preparing for Advanced Training

Lesson 2 introduces new capabilities commonly needed in real-world pipelines:

• multi-GPU distributed training
• mixed precision training
• larger effective batch sizes
• ONNX or TorchScript model export
• ONNX Runtime production inference

Here is the key design rule we follow:

We scale the system without touching the model or data code unless absolutely necessary.

Lightning was built exactly for this: features (e.g., AMP and DDP) are injected through the Trainer, not through the model.

Hydra was built to help you swap configurations without modifying Python files. This keeps your codebase stable (a critical requirement for teams and MLOps pipelines).

We continue using Lightning’s ModelCheckpoint and LearningRateMonitor callbacks to automatically track the best model and log learning-rate schedules during training.

Configuration Extensions

Lesson 2 adds new YAML files to support advanced features, but again, no source code changes are required. Each Hydra run is automatically logged to outputs/YYYY-MM-DD/HH-MM-SS/, keeping experiments isolated and reproducible.

Your directory now includes:

configs/
    trainer/
        single_gpu.yaml
        ddp.yaml
        fsdp.yaml

These configurations enable:

• single_gpu.yaml: local training
• ddp.yaml: distributed GPU training
• fsdp.yaml: model sharding for large models
• future extensions (AMP, gradient accumulation, etc.)

Because of Hydra’s override system, switching between them is as simple as:

python src/train.py trainer=ddp
python src/train.py trainer=fsdp
python src/train.py trainer=single_gpu

No code changes. No rewriting loops. No new training scripts.

This is exactly why Lightning and Hydra are so powerful: the same codebase can serve research, training at scale, and export pipelines without branching or duplication.

Enabling Mixed Precision Training with PyTorch Lightning AMP

Mixed precision training is one of the easiest “free wins” you can get when training Transformer models. Instead of doing every operation in full 32-bit floating point (FP32), we let the GPU run most of the math in 16-bit (FP16 or BF16) while keeping a few critical values in FP32 for numerical stability. The result is faster training, lower memory usage, and almost identical accuracy, especially on modern GPUs and Apple Silicon.

PyTorch Lightning wraps all of this inside the Trainer so you do not have to touch autocast contexts or manual gradient scaling. In this lesson, we will enable mixed precision purely through Hydra configurations, and Lightning will handle the rest.

Why Mixed Precision Boosts Transformers

Transformer models are heavy on matrix multiplications, which GPUs are extremely good at accelerating in half precision. When you switch from FP32 to a mixed precision mode (e.g., 16-mixed or bf16-mixed):

• Many tensor operations run 1.5-2.5× faster.
• Activations and gradients consume roughly half the memory.
• You can often increase batch size without running out of video random-access memory (VRAM).

Lightning uses PyTorch’s Automatic Mixed Precision (AMP) under the hood, so you still get stable training through automatic loss scaling. From an MLOps point of view, this is a small configuration change that can dramatically reduce training time and GPU cost without changing your codebase.

Updating Hydra: precision: 16-mixed

You do not need to modify data_module.py or model_module.py. Mixed precision is controlled entirely through the trainer configuration.

Here is the single-GPU trainer configuration in configs/trainer/single_gpu.yaml:

# Single GPU Configuration with Mixed Precision
# For experimentation and development

strategy: auto
devices: 1
accelerator: auto

# Mixed precision for faster training
precision: 16-mixed

# Training configuration
max_epochs: 3

# Performance
deterministic: false
benchmark: true

# Checkpointing
enable_checkpointing: true

# Progress
enable_progress_bar: true
log_every_n_steps: 10

Compared with an FP32 trainer, the required configuration change for mixed precision is:

precision: 16-mixed

Your DDP trainer in configs/trainer/ddp.yaml is already AMP-ready too:

# DDP (Distributed Data Parallel) Strategy
# Best for: Multi-GPU training on single node or multi-node setups
# Use when: Model fits in GPU memory, you want data parallelism
# Core distributed settings
strategy: ddp
devices: auto  # Use all available GPUs
accelerator: auto  # Auto-detect GPU/CPU

# Mixed precision for faster training
precision: 16-mixed  # FP16 automatic mixed precision

# Training configuration
max_epochs: 5  # More epochs for production
accumulate_grad_batches: 1  # Gradient accumulation (increase if OOM)

# Performance optimizations
deterministic: false  # Set to true for full reproducibility (slower)
benchmark: true  # Optimize CUDA kernels for input size

# Checkpointing
enable_checkpointing: true

# Progress tracking
enable_progress_bar: true
log_every_n_steps: 10

# DDP-specific optimizations
# ddp_find_unused_parameters: false  # Uncomment if you get DDP warnings

For FSDP, you are using BF16 mixed precision by default (better stability on very large models):

# FSDP (Fully Sharded Data Parallel) Strategy
# Best for: Very large models that don't fit in single GPU memory
# Use when: Model parameters exceed GPU memory, need memory efficiency

# Core distributed settings
strategy: fsdp
devices: auto  # Use all available GPUs
accelerator: auto

# Mixed precision - BF16 recommended for FSDP
precision: bf16-mixed  # Better stability than FP16 for large models

# Training configuration
max_epochs: 5
accumulate_grad_batches: 1

# Performance settings
deterministic: false
benchmark: true

# Checkpointing
enable_checkpointing: true

# Progress tracking
enable_progress_bar: true
log_every_n_steps: 10
# FSDP-specific settings
# Note: FSDP automatically shards model parameters, gradients, and optimizer states
# This allows training models that wouldn't fit on a single GPU

All of these configurations are passed to the Trainer through:

trainer = pl.Trainer(
    **cfg.trainer,
    callbacks=callbacks,
    logger=logger,
)

So switching between FP32, FP16, and BF16 is just a matter of changing precision in YAML (no changes to the training loop, no new context managers, and no extra boilerplate).

When to Use FSDP (Fully Sharded Data Parallel)

While DDP works well when the entire model fits on one GPU, FSDP is designed for cases where it doesn’t.

Instead of replicating the whole model on each GPU, FSDP shards parameters, gradients, and optimizer states across GPUs, allowing you to train models many times larger than a single GPU’s memory.

Lightning makes this just as simple as DDP:

python src/train.py trainer=fsdp

FSDP uses BF16 mixed precision by default (configured in configs/trainer/fsdp.yaml), which provides better stability for large-scale Transformer models.

Use FSDP if:

• you are hitting OOM even with small batch sizes
• you are experimenting with larger Transformer backbones
• you want memory-efficient training across multiple GPUs

Running AMP on NVIDIA GPUs or Apple Silicon

Once the configurations are in place, running AMP simply requires choosing the right trainer preset from the CLI.

Single GPU with Mixed Precision (NVIDIA GPU or Apple Silicon):

# From the lesson2 directory
python src/train.py trainer=single_gpu

Because trainer=single_gpu already sets:

accelerator: auto
devices: 1
precision: 16-mixed

Lightning will:

• Use CUDA if an NVIDIA GPU is available,
• Use MPS on a Mac with Apple Silicon,
• Fall back to the CPU otherwise (mixed precision provides little benefit there).

If you want to be explicit, you can override precision on the command line:

python src/train.py trainer=single_gpu trainer.precision=32        # FP32 baseline
python src/train.py trainer=single_gpu trainer.precision=16-mixed # FP16 AMP

The same pattern applies to DDP:

# Use all available GPUs with AMP
python src/train.py trainer=ddp

# Or force 2 GPUs with mixed precision
python src/train.py trainer=ddp trainer.devices=2 trainer.precision=16-mixed

For very large models using FSDP:

python src/train.py trainer=fsdp trainer.precision=bf16-mixed

In every case, you are only changing Hydra configuration values, while the underlying Lightning code path remains exactly the same.

What Speedups Should You Expect?

The exact speedup depends on your GPU, batch size, and model size, but for Transformer-style models like DistilBERT, you may see:

• 1.5-2× faster iteration times
• ~50% lower activation memory usage
• Similar validation accuracy to FP32

If you want to quantify performance in your own environment, you can run a simple timing experiment:

# FP32 (baseline)
time python src/train.py trainer=single_gpu trainer.precision=32 trainer.max_epochs=1

# FP16 mixed precision
time python src/train.py trainer=single_gpu trainer.precision=16-mixed trainer.max_epochs=1

You can compare the wall-clock time and VRAM usage between the 2 runs. In real MLOps pipelines, this can translate into lower training costs and the ability to run larger experiments on the same hardware.

Distributed Training with PyTorch Lightning DDP for Multi-GPU Scaling

Scaling your model across multiple GPUs is one of the fastest ways to reduce training time for Transformer architectures. PyTorch Lightning makes this dramatically easier by handling the boilerplate (process spawning, gradient synchronization, device management, and checkpoint coordination). In Lesson 2, we enable Distributed Data Parallel (DDP) using a Hydra configuration file and a single command-line override.

Let us walk through how this works and how you can run multi-GPU training without changing a single line of Python code.

What Distributed Data Parallel (DDP) Is

Distributed Data Parallel (DDP) is PyTorch’s recommended approach for training models across multiple GPUs. Each GPU receives:

• a full copy of the model
• a shard of the dataset
• parallel forward and backward passes

After each backward pass, gradients are synchronized across all GPUs so that every model replica remains synchronized.

The benefit is simple:

More GPUs → Larger global batch size → Faster training

Lightning abstracts the entire DDP workflow, so you do not need to manually write any multiprocessing logic, barrier synchronization, or device placement. All of this is handled by Trainer(strategy="ddp").

Hydra and Lightning Configuration (strategy: ddp)

DDP is enabled entirely through configuration, not code.

Here is the Hydra configuration file you include in Lesson 2:

configs/trainer/ddp.yaml

# Distributed training with DDP
strategy: ddp
accelerator: auto
devices: auto
precision: 16-mixed
max_epochs: 3

# Optional: helps stabilize multi-GPU runs
num_nodes: 1

This file replaces the default trainer settings when you select it from the CLI.

Lightning reads this configuration and automatically enables:

• automatic GPU detection
• multi-process spawning
• distributed samplers for the DataModule
• synchronized batch normalization
• gradient synchronization across devices
• safe checkpointing on rank 0

You do not need to modify train.py, DataModule, or LightningModule.

Running Multi-GPU Training

Use the following command to launch DDP:

$ python src/train.py trainer=ddp

To explicitly choose the number of devices:

$ python src/train.py trainer=ddp trainer.devices=2

To train on all visible GPUs:

$ python src/train.py trainer=ddp trainer.devices=auto

You can also combine overrides:

$ python src/train.py trainer=ddp trainer.precision=bf16-mixed data.batch_size=16

Hydra composes the configurations, Lightning spawns the GPU workers, and DDP executes the training loop.

No additional code is required.

Gradient Accumulation for Large Effective Batch Sizes

Modern Transformer models often benefit from larger batch sizes because they produce smoother gradients, more stable optimization, and sometimes higher accuracy in fewer iterations. However, large batches require more GPU memory, and even a mid-sized GPU may not be able to process them directly.

Gradient accumulation solves this by splitting a large batch across multiple smaller forward passes. Lightning collects gradients over several steps before performing one optimizer update, giving you the effect of large-batch training without increasing memory usage.

Why Gradient Accumulation Helps

Instead of running:

• Batch size = 64 (requires large GPU memory)

…you can simulate it as:

• batch_size = 8
• accumulate_grad_batches = 8

Lightning will:

• Run 8 forward and backward passes
• Accumulate gradients internally
• Call optimizer.step() only once

This produces the same gradient update you would get from a single batch of size 64 while using only the memory required for a batch of 8.

This is especially useful when:

• You hit CUDA OOM errors during DDP or FSDP training
• You want larger effective batch sizes for stability
• You are training with mixed precision, which can benefit from larger batches
• You are using limited GPU hardware, including Apple Silicon

Updating Hydra Configurations

Your Lesson 2 repository supports gradient accumulation through Hydra overrides.

The primary trainer configurations (ddp.yaml, fsdp.yaml, single_gpu.yaml) set:

accumulate_grad_batches: 1

This means accumulation is opt-in and controlled by command-line overrides.

You can also create a dedicated Hydra configuration:

configs/trainer/accum.yaml (optional)

accumulate_grad_batches: 4

However, this is not required because the command-line interface (CLI) override is often clearer and more flexible.

Using Gradient Accumulation via CLI

Because Hydra merges configurations top-down, you can override accumulation using the following options:

Single GPU

python src/train.py trainer=single_gpu trainer.accumulate_grad_batches=4

DDP (multi-GPU)

python src/train.py trainer=ddp trainer.accumulate_grad_batches=4

With Mixed Precision

python src/train.py trainer=ddp trainer.precision=16-mixed trainer.accumulate_grad_batches=8

With Large Batches

python src/train.py trainer=ddp data.batch_size=4 trainer.accumulate_grad_batches=8

In this case:

effective_batch_size = batch_size × accumulate_grad_batches × num_gpus

For example, with:

• batch_size = 4
• accumulate_grad_batches = 8
• num_gpus = 2

the effective batch size is:

4 × 8 × 2 = 64

This configuration normally requires ~16-20 GB of GPU memory but is now possible even on a laptop GPU.

Batch Size vs Memory Tradeoffs

Gradient accumulation directly affects:

• Lower Memory Consumption: Only the small per-step batch must fit in memory.
• Equivalent Gradient Updates: The optimizer receives the same gradient update as full-size training.
• Slightly Longer Training Time: Lightning performs more forward and backward passes before each update. However, this approach is still more efficient than dealing with out-of-memory (OOM) errors or reducing the sequence length or tokenizer settings.

Lightning Integration (Zero Code Changes)

Your Lesson 2 train.py requires zero code modification.

Lightning handles:

• gradient scaling
• accumulation logic
• optimizer stepping
• multi-GPU gradient synchronization
• mixed-precision scaling

You control the behavior through Hydra configurations.

Exporting PyTorch Lightning Transformer Models to ONNX and TorchScript

Training a model is only half the story because real-world ML systems need fast, portable, and framework-agnostic inference. In Lesson 2, you extend your training pipeline to automatically export your best checkpoint into two production-ready formats:

• ONNX: for cross-platform, high-performance inference (ONNX Runtime, TensorRT, OpenVINO)
• TorchScript: for PyTorch-native, C++ or mobile deployment

The key design goal is that none of this export logic lives inside the LightningModule. Instead, the export code is isolated in train.py and driven entirely by Hydra configuration. This keeps your model clean and your training loop reusable in both research and production settings.

Why Export?

Exported models provide several benefits:

Fast inference: ONNX Runtime routinely delivers 2-5× faster CPU inference than PyTorch eager mode.

Portability: A single ONNX file runs on:

• Linux, macOS, and Windows
• mobile devices
• GPU runtimes (e.g., TensorRT)
• serverless platforms (AWS Lambda with ONNX Runtime)

Reproducible Production Pipelines: TorchScript provides a stable, serialized version of your model for:

• C++ backends
• embedded devices
• custom inference servers
• TorchServe

DVC-ready artifact tracking: Exports drop cleanly into:

artifacts/models/
artifacts/metrics/

These artifacts can be versioned and tracked, just like code.

Enabling Export via Hydra

Exports are controlled by your root configuration:

configs/config.yaml

# Export configuration (LESSON 2 feature)
export:
  enabled: true       # Enable model export after training
  onnx: true         # Export to ONNX format
  torchscript: true  # Export to TorchScript format

You can override these settings at runtime:

Export Both Formats

python src/train.py trainer=ddp export.enabled=true

Export Only ONNX

python src/train.py export.enabled=true export.torchscript=false

Export Only TorchScript

python src/train.py export.enabled=true export.onnx=false

Lightning handles training, ModelCheckpoint saves the best checkpoint, and then your script reloads that checkpoint for export.

How the Export Pipeline Works

During export, the script automatically reloads the best checkpoint saved by ModelCheckpoint, ensuring the exported ONNX or TorchScript files always correspond to the best validation score.

Your train.py contains 2 production-grade export utilities:

ONNX Export

export_to_onnx(best_model, onnx_path, cfg.data.max_length)

TorchScript Export

export_to_torchscript(best_model, ts_path, cfg.data.max_length)

Both functions:

• load the best checkpoint
• switch the model to evaluation mode
• move it to the CPU
• construct a dummy input of shape (1, max_length)
• save the exported file under artifacts/models/

This separation ensures:

• no modification to model_module.py
• clean Transformer traceability
• reproducibility across runs

ONNX Export Details

The utility inside train.py uses:

torch.onnx.export(
    model,
    (dummy_input_ids, dummy_attention_mask),
    str(output_path),
    opset_version=14,
    input_names=["input_ids", "attention_mask"],
    output_names=["logits"],
    dynamic_axes={
        "input_ids": {0: "batch_size", 1: "sequence_length"},
        "attention_mask": {0: "batch_size", 1: "sequence_length"},
        "logits": {0: "batch_size"},
    },
)

Key features:

Dynamic axes: Allow variable batch sizes and sequence lengths at inference time.

Opset 14: Compatible with ONNX Runtime and TensorRT.

Framework-agnostic: You can deploy the same .onnx file using:

• ONNX Runtime
• NVIDIA TensorRT
• OpenVINO
• Triton Inference Server
• AWS Lambda (serverless inference)

ONNX Runtime Inference Example

Readers can test the exported model using:

import onnxruntime as ort
import numpy as np

session = ort.InferenceSession("artifacts/models/sentiment_classifier.onnx")

input_ids = np.random.randint(0, 1000, (1, 128), dtype=np.int64)
attention_mask = np.ones((1, 128), dtype=np.int64)

outputs = session.run(None, {
    "input_ids": input_ids,
    "attention_mask": attention_mask
})

print("Logits:", outputs[0])

This should produce logits similar to those from PyTorch inference.

TorchScript Export Details

The utility inside train.py uses:

traced_model = torch.jit.trace(
    model,
    (dummy_input_ids, dummy_attention_mask)
)
traced_model.save(str(output_path))

TorchScript provides:

PyTorch-Native Deployment

Works with:

• TorchServe
• custom C++ services
• mobile runtimes
• embedded systems

Stable, Production-Safe Serialization

Unlike Python pickles, TorchScript provides a stable serialized format for multi-process and multi-node systems.

Artifacts Stored in a DVC-Ready Structure

After export completes, your script creates:

artifacts/
├── models/
│   ├── sentiment_classifier.onnx
│   ├── sentiment_classifier.pt
└── metrics/
    └── metrics.json

Your script also generates artifacts/metrics/metrics.json, which contains the final validation loss and accuracy. Your continuous integration and continuous deployment (CI/CD) pipeline, Data Version Control (DVC), or model registry can consume this file.

This folder is well suited for DVC tracking:

dvc add artifacts/models/
dvc add artifacts/metrics/

End-to-End Export Command

A typical production run:

python src/train.py \
    trainer=ddp \
    trainer.precision=16-mixed \
    export.enabled=true

produces:

Exported: sentiment_classifier.onnx
Exported: sentiment_classifier.pt
Saved metrics.json

Both formats are validated, versioned, and ready for deployment.

Zero Modifications to Model Code

Because the LightningModule remains unchanged, you maintain:

• readability
• testability
• modularity
• compatibility with future lessons (CI/CD, DVC, deployments)

The entire export logic is kept in one place.

What's next? We recommend PyImageSearch University.

Course information:
120+ total classes • 115+ hours of on-demand code walkthrough videos • Last updated: October 2026
★★★★★ 4.84 (128 Ratings) • 16,000+ Students Enrolled

I strongly believe that if you had the right teacher you could master computer vision and deep learning.

Do you think learning computer vision and deep learning has to be time-consuming, overwhelming, and complicated? Or has to involve complex mathematics and equations? Or requires a degree in computer science?

That’s not the case.

All you need to master computer vision and deep learning is for someone to explain things to you in simple, intuitive terms. And that’s exactly what I do. My mission is to change education and how complex Artificial Intelligence topics are taught.

If you're serious about learning computer vision, your next stop should be PyImageSearch University, the most comprehensive computer vision, deep learning, and OpenCV course online today. Here you’ll learn how to successfully and confidently apply computer vision to your work, research, and projects. Join me in computer vision mastery.

Inside PyImageSearch University you'll find:

• ✓ 120+ courses on essential computer vision, deep learning, and OpenCV topics
• ✓ 94+ Certificates of Completion
• ✓ 115+ hours of on-demand video
• ✓ Brand new courses released regularly, ensuring you can keep up with state-of-the-art techniques
• ✓ Pre-configured Jupyter Notebooks in Google Colab
• ✓ Run all code examples in your web browser — works on Windows, macOS, and Linux (no dev environment configuration required!)
• ✓ Access to centralized code repos for all 540+ tutorials on PyImageSearch
• ✓ Easy one-click downloads for code, datasets, pre-trained models, etc.
• ✓ Access on mobile, laptop, desktop, etc.

Click here to join PyImageSearch University

Summary

In this lesson, you transformed a simple sentiment-classification project into a scalable, production-ready training system using PyTorch Lightning and Hydra. You learned how to enable mixed precision to speed up training, use Distributed Data Parallel (DDP) to leverage multiple GPUs, and apply gradient accumulation to simulate large batch sizes without increasing memory usage. All of these upgrades were achieved without modifying your model or data code, demonstrating the benefits of Lightning’s abstraction and Hydra’s configuration-driven design.

You also implemented a clean and reusable export pipeline, allowing the same trained model to be saved as ONNX or TorchScript and used in lightweight, fast inference systems. With a small export utility built directly into the train.py script and a clear configuration structure, you now have a workflow that produces consistent, portable artifacts suitable for CI/CD, edge devices, or cloud deployment.

Taken together, these enhancements provide a robust foundation for real-world ML engineering. You have taken the same codebase from Lesson 1 and upgraded it with improved performance, scalability, and deployment-ready outputs. This prepares you for next steps such as deployment, optimization, monitoring, or integrating your models into full MLOps pipelines.

Citation Information

Singh, V. “Scaling, Optimizing, and Exporting Transformers with PyTorch Lightning,” PyImageSearch, S. Huot, A. Sharma, and P. Thakur, eds., 2026, https://pyimg.co/2xntj

@incollection{Singh_2026_scaling-optimizing-exporting-transformers-pytorch-lightning,
  author = {Vikram Singh},
  title = {{Scaling, Optimizing, and Exporting Transformers with PyTorch Lightning}},
  booktitle = {PyImageSearch},
  editor = {Susan Huot and Aditya Sharma and Piyush Thakur},
  year = {2026},
  url = {https://pyimg.co/2xntj},
}

To download the source code to this post (and be notified when future tutorials are published here on PyImageSearch), simply enter your email address in the form below!

Download the Source Code and FREE 17-page Resource Guide

Enter your email address below to get a .zip of the code and a FREE 17-page Resource Guide on Computer Vision, OpenCV, and Deep Learning. Inside you'll find my hand-picked tutorials, books, courses, and libraries to help you master CV and DL!

Download the code!

Website

The post Scaling, Optimizing, and Exporting Transformers with PyTorch Lightning appeared first on PyImageSearch.