FastAPI for MLOps: Python Project Structure and API Best Practices

In this lesson, you will learn how to structure a Machine Learning (ML) project like a real production system, complete with a src directory layout, layered configuration, environment management, logging, and a FastAPI service that exposes your model through clean Application Programming Interface (API) routes.

This lesson is the 1st of a 2-part series on Software Engineering for Machine Learning Operations (MLOps):

1. FastAPI for MLOps: Python Project Structure and API Best Practices (this tutorial)
2. Lesson 2

To learn how to build reliable, scalable ML software the right way, just keep reading.

Looking for the source code to this post?

Jump Right To The Downloads Section

Introduction

Modern ML systems do not succeed because of models alone — they succeed because of the software engineering wrapped around them. Most real-world failures in MLOps come from poor structure, missing configuration, messy environments, unclear APIs, or nonexistent logging, not from bad ML.

This lesson gives you the engineering foundation you need to build ML systems that are stable, testable, and production-ready. You’ll learn how to structure your project, manage environments, load configurations, build APIs, and prepare your system for future modules like testing, deployment, and automation.

To learn how solid software engineering underpins every ML workflow, just keep reading.

What You Will Build and Learn

In this lesson, you’ll build the backbone of a real ML application: a clean repository layout, environment management with modern tooling, configuration loading via Pydantic, structured logging, a FastAPI interface, and a simple service layer to power prediction.

These concepts form the “foundation layer” every MLOps system relies on — regardless of the model you eventually plug in.

Why Software Engineering Comes First in MLOps Best Practices

ML projects fail not because the model is wrong, but because the plumbing around the model collapses. Scripts turn into spaghetti, notebooks become unmaintainable, configs get scattered, and environments drift until the system becomes impossible to debug.

Good software engineering fixes this by introducing structure, consistency, and predictable behavior. When your API, config, logs, and model code work together cleanly, everything built on top (e.g., testing, serving, scaling, monitoring) suddenly becomes reliable.

Where This Fits in the Overall Curriculum

This lesson is the foundation of the entire MLOps series. Everything that comes next — testing, model integration, deployment workflows, Continuous Integration/Continuous Delivery (CI/CD) automation, monitoring, and scaling — builds on the engineering habits you establish here.

Think of this as your “software engineering base layer.” Once you master this structure, adding real models, adding load testing, or plugging the system into cloud infrastructure becomes far easier.

Python Project Structure Best Practices for MLOps

A well-structured repository is the first sign of a healthy ML system. Before we write any API code or load a model, we need a layout that cleanly separates configuration, services, models, and utilities. This not only prevents chaos — it makes testing, scaling, and future modules dramatically easier.

How to Structure a Python Project with src/ Layout

ML projects quickly become messy if everything sits at the root level. The src/ layout prevents naming collisions, enforces imports that match production structure, and makes it clear where application code actually lives.

This is the same structure used in mature Python services deployed in production environments.

Python Project Structure Explained: Repository Walkthrough

Here’s the repository layout we’re working with in this module (the exact tree will be shown later when you provide it):

sw-eng-mlops/
│
├── src/
│   ├── core/
│   ├── models/
│   ├── services/
│   ├── api/
│   ├── utils/
│   └── config/
│
├── tests/
│   ├── unit/
│   ├── integration/
│   └── performance/
│
├── pyproject.toml
├── README.md
├── setup_env.sh
└── .env.example

This structure is intentionally clean: core/ contains primitives, models/ stores your ML logic, services/ contains business logic, and api/ exposes everything through FastAPI routes.

Python Project Structure Best Practices: Directory Breakdown

core/ — The Application Base Layer

This folder contains shared components such as logging setup, base classes, or utility abstractions. Everything here is meant to be reusable across the whole system.

models/ — ML or Dummy Model Code

Even if you’re starting with a dummy model, isolating model code here makes it easy to swap in real models later.

services/ — The Business Logic Layer

This is where you place the logic that actually powers /predict, not inside the API route. This separation keeps production-grade APIs maintainable.

api/ — FastAPI Endpoints

Routes live here. Each endpoint calls a service, which calls a model.

Tight, clean, and testable.

utils/ — Shared Helpers

Config loaders, yaml readers, or general-purpose helper functions sit here.

If it isn’t domain logic or a model, it goes here.

config/ — Configuration Files

YAML configs, BaseSettings classes, validation logic, and environment overrides.

Centralizing config makes behavior predictable and testable.

How This Structure Scales to Larger ML Systems

This layout scales easily as your ML workload grows:

• Add a new model → create a folder inside models/.
• Add a new prediction workflow → add a service in services/.
• Add new API functionality → add a route in api/.
• Add data pipelines or vector DB logic → expand core/ or services/.

This way, the project grows horizontally, not chaotically.

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.

Managing Python Dependencies with Poetry for ML Projects

Modern MLOps projects rely on predictable, repeatable environments — and this section teaches you how to create exactly that. Before we build APIs or load models, we need a clean, isolated workspace where dependencies are installed, versions are pinned, and tools behave consistently across machines.

To learn how to manage dependencies, virtual environments, and setup scripts in real-world ML projects, just keep reading.

Python Poetry vs PDM vs UV: Choosing a Package Manager for MLOps

There are 3 modern Python toolchains worth knowing:

• Poetry: full-featured dependency + environment + packaging manager.
• PDM (Python Dependency Manager): simpler and faster than Poetry, with PEP-582 support.
• UV: an extremely fast Rust-based package manager from Astral.

All 3 support pyproject.toml, the modern Python standard for dependencies and metadata.

Teams often standardize on a single tool, but your project supports all three, so students can use whichever they prefer.

Understanding pyproject.toml for Python Project Configuration

Your pyproject.toml defines:

• project name, version, description
• dependencies like fastapi, pydantic, pyyaml
• dev tools like pytest (Lesson 2)
• optional entrypoints (start-server = "src.main:main")

In other words, it is the single source of truth for installation and build metadata.

Any tool (Poetry, PDM, UV, pip) reads this file to install exactly what the project needs.

This is how professional ML systems avoid “works on my machine” issues.

Installing Dependencies (Poetry, PDM, UV)

Using Poetry (recommended)

poetry install
poetry shell
poetry run python src/main.py

Poetry creates an isolated virtual environment and resolves all versions deterministically.

Using UV (lightweight + blazing fast)

uv venv
source .venv/bin/activate
uv pip install -e .
python src/main.py

UV is perfect for fast installs and CI systems where speed matters.

Using PDM (simple + modern)

pdm install
pdm run python src/main.py

PDM feels like npm — no venv folder by default; lightweight and straightforward.

Managing Python Virtual Environments for Reproducible MLOps

Regardless of what tool you choose, the goal is the same: isolate project dependencies from the system Python installation.

• Poetry creates its own environment automatically.
• UV uses .venv/ inside your project.
• PDM can create or avoid virtual environments depending on the configuration.

The important principle:

Never install ML dependencies globally.

Environments keep your project reproducible and safe.

Automating MLOps Setup with Python Environment Scripts

Your project includes a helper script:

./scripts/setup_env.sh

This script:

• Detects whether Poetry, UV, or plain pip is available
• Installs dependencies using the detected tool
• Creates or activates the .env file
• Shows the next steps to start the API

This is extremely helpful for teams because it removes all “setup guessing” and gives new developers a consistent starting point.

You now know how environments, dependency managers, and pyproject.toml work together to create a stable foundation for ML systems. With everything installed and configured, you’re ready to build and serve a real API.

Up next, we’ll create your first ML service with FastAPI and connect it to your project’s service layer.

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!

Configuration Management in MLOps: YAML, .env, and Pydantic

How the entire ML system loads, merges, and applies configuration at runtime.

Configuration is one of the most important engineering foundations in any ML system. In Lesson 1, we want students to walk away understanding not only why configuration matters but exactly how this project loads and merges config values. That means stepping through the real code inside src/core/config.py, the .env.example, and configs/config.yaml.

We also want to show how the API, model, and services consume configuration. So when students replace the dummy model with a real one, the pattern already scales.

Let’s walk through it piece by piece.

Using Pydantic Settings for MLOps Configuration Management

Your configuration system starts with a Settings class:

class Settings(BaseSettings):
    api_host: str = "0.0.0.0"
    api_port: int = 8000
    debug: bool = False
    environment: str = "development"
    log_level: str = "INFO"

    class Config:
        env_file = ".env"
        env_file_encoding = "utf-8"

What This Means for MLOps Configuration and System Design

• Pydantic’s BaseSettings automatically reads:
  • environment variables
  • .env file
  • any overrides you pass at runtime
• Defaults are provided in code so the system always works, even if .env is missing.
• Type safety ensures that if someone writes API_PORT=hello, the app will fail fast.

This is the right pattern for ML systems where dozens of environment variables must be synchronized across dev, test, staging, and production.

Loading YAML and Merging Layers

Next comes one of the most important parts of your system:

def load_config() -> Settings:
    settings = Settings()

    config_path = "configs/config.yaml"
    if os.path.exists(config_path):
        yaml_config = load_yaml_config(config_path)

        for key, value in yaml_config.items():
            if hasattr(settings, key):
                setattr(settings, key, value)

    return settings

Why This Is Powerful

You now have layered configuration, which production ML systems use everywhere:

Layer 1: Code defaults

Ensures the app always runs.

Layer 2: YAML (configs/config.yaml)

Great for team-shared configs, model settings, cache sizes, service parameters.

Layer 3: .env file

Local overrides (ports, debug mode, secrets).

Layer 4: Runtime environment variables

Final source of truth in cloud deployments.

This layered system prevents the “hard-coded value” trap and keeps ML infra consistent across environments.

Designing YAML Configs for Scalable MLOps Pipelines

Your YAML file contains deeper structural config:

api_host: "0.0.0.0"
api_port: 8000
debug: true
environment: "development"

log_level: "INFO"

model:
  name: "dummy_classifier"
  version: "1.0.0"
  cache_size: 100

service:
  timeout: 30
  max_retries: 3

Even though Settings does not yet support nested objects for models or services, YAML allows you to introduce new structured configuration later. This is how real ML teams configure:

• model version
• tokenizer version
• max batch size
• timeouts
• cache settings
• experiment IDs

Using .env Files for Secure MLOps Configuration

You also provide .env.example:

API_PORT=8000
API_HOST=0.0.0.0
DEBUG=true
ENVIRONMENT=development
LOG_LEVEL=INFO

Why Configuration Management Matters in MLOps Systems

• .env.example acts as documentation and a template.
• You copy it to .env, fill values, and the system boots.
• This is a best practice in every production ML repo.

How the App Uses Configuration (src/main.py)

Your FastAPI entrypoint reads config like this:

logger.info(f"Starting server on {settings.api_host}:{settings.api_port}")

uvicorn.run(
    "main:app",
    host=settings.api_host,
    port=settings.api_port,
    reload=settings.debug
)

Meaning:

• Change .env to API_PORT=9000: Your app automatically runs on port 9000.
• Change YAML to debug: false: Hot reload turns off.

This is the practical benefit of structured configuration: no hard-coded values are buried inside the code.

How FastAPI Uses Configuration in Production MLOps Systems

Today, your inference service is simple, but in real projects, you might use:

• model name
• version
• batch size
• latency budget
• max retries
• cache settings
• rate limits

All of these come from settings, not hardcoded logic.

In this lesson, you teach the pattern, so when the dummy model is eventually replaced with an Open Neural Network Exchange (ONNX) model, a Hugging Face model, or a custom PyTorch model, the service already has the right structure.

Extending MLOps Configuration Safely in Python Projects

Suppose tomorrow you want:

MODEL_PATH=models/checkpoint.pt
ENABLE_CACHE=true
CACHE_TTL=300

You add:

model_path: str = "models/dummy.pt"
enable_cache: bool = False
cache_ttl: int = 120

Then update .env.example. Then, optionally override in YAML.

The app instantly supports new behavior — no rewrites, no refactoring, no confusion.

This is the level of software engineering maturity we want students to learn.

Logging Best Practices for MLOps and FastAPI Applications

Logging is one of the most underappreciated parts of an ML system. A model prediction might take milliseconds, but diagnosing a production issue without proper logs can take hours. Good logs reduce that time to minutes. In this section, we’ll look at how our lesson’s project initializes a logger, formats log messages, and uses logs consistently across the entire API.

Why Logging Is Critical for ML Systems

ML systems fail in ways traditional software does not.

A model might produce an unexpected prediction, a dependency might break silently, or the environment might load the wrong configuration. Logging gives you the breadcrumbs needed to understand:

• What inputs reached the API
• What model version was used
• What the service did before failing
• How often errors occur
• Whether latency is increasing

Logs are your “black box recorder” when something goes wrong, and they’re equally important when everything seems to be working — because they tell you why things are working.

Logger Initialization

The project defines a single shared logger in src/core/logger.py:

import logging
import sys

logger = logging.getLogger("mlops-lesson1")
logger.setLevel(logging.INFO)

handler = logging.StreamHandler(sys.stdout)
formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
handler.setFormatter(formatter)

if not logger.handlers:
    logger.addHandler(handler)

Here’s what this setup accomplishes:

• A named logger (mlops-lesson1) groups logs for later aggregation (e.g., in Datadog, ELK (Elasticsearch, Logstash, Kibana), OpenTelemetry).
• INFO as the default level ensures we capture meaningful operational details without spamming output.
• A StreamHandler writes logs to stdout — the standard for containerized deployments (Docker, Kubernetes).
• A simple timestamped formatter makes logs human-readable while remaining machine-parseable.
• The if not logger.handlers: guard prevents duplicate logs if modules are reloaded.

This small file gives us a production-friendly logger with minimal overhead.

Log Formatting and Levels

The logger uses this format:

2025-01-01 12:34:56 - INFO - Prediction result: positive

Each part of the log line matters:

• Timestamp: crucial for correlating logs with events or latency spikes.
• Log level: signals severity: INFO, WARNING, ERROR.
• Message: the human-readable explanation.

In MLOps systems, you’ll most commonly use:

• INFO for model loading, API calls, predictions
• WARNING for slow responses, unexpected patterns
• ERROR when something fails

Because FastAPI reloads modules during development, you may see log duplication without safeguards — which is why we include the if not logger.handlers: check.

If you later want structured JSON logs (for cloud log ingestion), this same module is the place to upgrade.

Logging Across the App

The logger is used in multiple places, showing a consistent logging strategy.

Health endpoint (src/main.py)

@app.get("/health")
async def health_check():
    logger.info("Health check requested")
    return {"status": "ok"}

This gives visibility into uptime checks — important when a load balancer or Kubernetes performs probes.

Prediction endpoint (src/services/inference_service.py)

logger.info(f"Making prediction for input: {input_text[:50]}...")
prediction = model.predict(input_text)
logger.info(f"Prediction result: {prediction}")

Here we log:

• The incoming input (truncated to avoid leaking full user data)
• The model’s output
• Any errors

If something goes wrong:

except Exception as e:
    logger.error(f"Error during prediction: {str(e)}")
    raise

This ensures errors appear in the logs before FastAPI converts them into HTTP exceptions.

Server startup (main.py)

logger.info(f"Starting server on {settings.api_host}:{settings.api_port}")

This is important for:

• verifying the config loaded correctly
• ensuring the correct port is used
• debugging environments with conflicting overrides

Together, This Gives Us Structured, Traceable Behavior Across the App

If a user reports:

“The API feels slow today.”

You can immediately look at:

• prediction request timestamps
• whether model loading was triggered again
• whether latency warnings appear
• whether certain inputs correlate with errors

Without logs, you’re flying blind.

FastAPI for MLOps: Building a Production ML API

APIs are the interface between your ML system and the outside world. Whether the consumer is a mobile app, a batch job, another microservice, or a human developer testing in Postman, every interaction eventually flows through an API. In MLOps, your API becomes the stable contract that hides internal details (model type, version, preprocessing, logging) — allowing you to upgrade models without breaking clients.

Why FastAPI Is Ideal for MLOps API Development

FastAPI gives you a fast, typed, and production-ready way to expose ML predictions.

It handles validation, serialization, documentation, and error responses, so your ML logic stays clean and modular.

The goal is simple: your API should stay stable even when everything behind it changes — models, configs, logging, monitoring, infrastructure.

Creating a FastAPI Application for Machine Learning APIs

Your project defines the API inside src/main.py:

from fastapi import FastAPI
app = FastAPI(
    title="ML Service API",
    description="Code Foundations & API Engineering for MLOps",
    version="0.1.0"
)

This initializes a fully documented ML service with:

• A title for the UI
• A description that shows up in Swagger
• A semantic version
• Automatically generated schemas

FastAPI instantly gives you API docs and a clean, declarative way to add endpoints.

Implementing Health Check Endpoints in FastAPI (MLOps)

A health endpoint is the first thing any production system needs.

Kubernetes, AWS Application Load Balancer (ALB), Docker Compose, Jenkins, and uptime monitors all rely on it.

Your implementation:

@app.get("/health")
async def health_check():
    logger.info("Health check requested")
    return {"status": "ok"}

This performs 2 critical functions:

• Confirms the API server is alive
• Confirms logs are working

It also gives you a simple smoke test to verify the environment.

Building a FastAPI Prediction Endpoint for ML Models

The /predict endpoint is where real ML work happens.

@app.post("/predict")
async def predict_route(input: str):
    return {"prediction": predict_service(input)}

This endpoint:

• Accepts a simple string input
• Passes it into the inference service
• Returns a structured JSON prediction

Because prediction logic is isolated in services/inference_service.py, the API stays lightweight and focused on HTTP behavior — not business logic.

Behind This Endpoint Is Your Prediction Engine

from models.dummy_model import DummyModel

model = DummyModel()

def predict(input_text: str) -> str:
    logger.info(f"Making prediction for input: {input_text[:50]}...")
    prediction = model.predict(input_text)
    logger.info(f"Prediction result: {prediction}")
    return prediction

Even though this is a dummy model, the structure mirrors real production design:

• The service layer owns the prediction logic
• The model is instantiated once
• Logging wraps the input and output

When you upgrade to a real transformer or classifier, the API does not need to change.

Deploying FastAPI with Uvicorn for MLOps Applications

The server entrypoint lives at the bottom of main.py:

def main():
    logger.info(f"Starting server on {settings.api_host}:{settings.api_port}")
    uvicorn.run(
        "main:app",
        host=settings.api_host,
        port=settings.api_port,
        reload=settings.debug
    )

A few details matter:

• reload=True reloads on code changes → perfect for development
• host and port come from config → ideal for containers/cloud
• logging is integrated → so you can trace server start behavior

You can run the server with:

poetry run start-server

or

uvicorn src.main:app --reload

Both give you a live API with hot reload.

Auto-Generated API Docs (Swagger, ReDoc)

FastAPI automatically exposes:

• Swagger UI: http://localhost:8000/docs
• ReDoc: http://localhost:8000/redoc
• OpenAPI schema: http://localhost:8000/openapi.json

These docs are invaluable in ML workflows because:

• You can test predictions interactively
• Product, QA, and frontend engineers can explore endpoints
• Payload schemas are always up to date
• No one needs to ask “What does this endpoint expect?”

FastAPI generates this from your Python type hints, which makes documentation essentially free.

MLOps Architecture: Service Layer Design Patterns

The service layer is where your application’s real business logic lives. In an ML system, this includes preprocessing, model selection, inference, error handling, postprocessing, and logging. By keeping this logic out of your API routes, you ensure that your codebase remains modular, testable, and ready for future model upgrades.

Why We Separate Services from Routes

FastAPI routes should only handle HTTP concerns: input validation, request parsing, and response formatting.

They should not know how your model works internally.

Separating logic into a services/ folder gives you:

• Cleaner API routes: easier to read and maintain
• Better testability: you can unit test the inference logic without starting a server
• Loose coupling: upgrading models doesn’t require rewriting routes
• Clear ownership: one layer handles HTTP, the other handles ML logic

This separation is one of the most critical software engineering patterns in MLOps — you want your system flexible enough that models can change, scale, or switch frameworks without touching your API.

Designing an ML Inference Service

Your inference logic lives in:

src/services/inference_service.py

Let’s look at how it’s structured:

from models.dummy_model import DummyModel
from core.logger import logger

# Initialize model
model = DummyModel()
logger.info(f"Loaded model: {model.model_name}")

This loads the model once at startup. In a real ML system, this is where:

• You load a transformer model
• You warm up a GPU
• You hydrate a vector store
• You initialize the tokenizer/preprocessor state

Then comes the prediction function:

def predict(input_text: str) -> str:
    logger.info(f"Making prediction for input: {input_text[:50]}...")
   
    try:
        prediction = model.predict(input_text)
        logger.info(f"Prediction result: {prediction}")
        return prediction
    except Exception as e:
        logger.error(f"Error during prediction: {str(e)}")
        raise

This function represents the business logic of your ML service:

• It trims the input for logging
• Calls the model’s predict()
• Logs errors and output cleanly
• Returns only the result — not HTTP details

This is exactly why we keep services separate: inference is not an HTTP concern, so it does not belong in a route.

Scaling MLOps Systems with Modular Service Architecture

A great design scales. Tomorrow, your system might need:

• SentimentService: for NLP
• RecommendationService: for personalization
• VisionService: that loads YOLO or CLIP
• BatchService: for async workflows
• RetrievalService: for Retrieval-Augmented Generation (RAG) pipelines

You don’t modify main.py or existing endpoints.

You simply add more files under:

src/services/
├── inference_service.py  
├── recommendation_service.py  
├── vision_service.py  
└── retrieval_service.py  

Each service becomes independent, testable, and reusable.

Later in Lesson 2, this design becomes even more powerful because:

• Unit tests: target individual services
• Integration tests: validate routes and services working together
• Load tests: measure the throughput of the /predict pipeline

By the time you add real ML models, this service layer becomes the heart of your system.

Model Abstraction in MLOps: Decoupling ML from APIs

Models change constantly in MLOps. Today you may be serving a dummy classifier; tomorrow it might be a 7B LLM or a YOLOv12 object detector. A good software engineering foundation treats the model as a pluggable, versioned component that can be replaced with minimal friction.

Your current models/ directory demonstrates exactly how this abstraction works.

Designing a Python ML Model Class for MLOps

Your lesson uses a simple placeholder model located at:

src/models/dummy_model.py

The goal of this class isn’t to perform “real” ML — it’s to give you a clean structure that mimics how production model classes are written.

class DummyModel:
    def __init__(self) -> None:
        self.model_name = "dummy_classifier"
        self.version = "1.0.0"
   
    def predict(self, input_data: Any) -> str:
        text = str(input_data).lower()
        if "good" in text or "great" in text:
            return "positive"
        return "negative"

Even in this tiny model, you already see foundational patterns:

• A constructor to load or initialize model state
• A predict() method that defines the inference interface
• model_name and version fields for introspection and tracking

This interface is intentionally minimal: it forces your service and API layers to depend on an abstraction, not on implementation details.

In real MLOps systems, this exact pattern makes it easy to introduce new models without breaking your API.

How to Replace Dummy Models with Production ML Models

Here’s where the abstraction shines.

If tomorrow you decide to replace the dummy model with:

• A Hugging Face transformer
• A PyTorch Lightning checkpoint
• A TensorRT engine
• An ONNX Runtime session
• A vLLM text-generation server
• A YOLO detection model

…all you need to do is drop a new file into:

src/models/

For example:

src/models/
├── dummy_model.py
├── sentiment_model.py
├── llm_generation_model.py
└── object_detector.py

And update your service:

from models.sentiment_model import SentimentModel
model = SentimentModel()

Nothing else changes.

Your FastAPI routes stay the same.

Your service interface stays the same.

Your tests stay the same (except for new model-specific tests).

This is model decoupling.

This is how ML systems avoid turning into tangled spaghetti when models evolve.

Versioning the Model Class

Model versioning is a real production concern, and your dummy model subtly teaches the pattern.

self.version = "1.0.0"

Model versioning matters because:

• You may deploy multiple models at once
• Clients might depend on specific behaviors
• A/B testing needs separate versions
• Rollbacks require deterministic reproducibility
• Monitoring tools (e.g., Prometheus or Langfuse) track model changes

In production, versioning happens in several places:

• version field in the class
• model registry tag (MLflow, SageMaker, Hugging Face Hub)
• Docker image tag
• config.yaml entry
• model card metadata

Your project follows the simplest, clearest entrypoint: a version attribute that propagates everywhere the model is used.

Later in Lesson 2, test cases and load tests will automatically pick up this version, mimicking real-world CI/CD systems that validate each model release.

Building Reusable Utilities in Python MLOps Projects

A well-designed ML system always contains a dedicated utilities layer — small, reusable functions that solve cross-cutting problems without polluting your core logic, service layer, or API routes.

In this project, the src/utils/ folder gives you a clean space to organize those helpers, starting with configuration loading, and is ready to grow as your system becomes more complex.

This layer keeps your codebase maintainable, testable, and extensible.

Loading YAML Configs

Your primary helper is load_yaml_config() found in:

src/utils/helpers.py

Here’s the implementation:

def load_yaml_config(path: str) -> Dict[str, Any]:
    config_path = Path(path)
   
    if not config_path.exists():
        return {}
   
    try:
        with open(config_path, 'r', encoding='utf-8') as file:
            config = yaml.safe_load(file)
            return config if config is not None else {}
    except yaml.YAMLError as e:
        print(f"Error loading YAML config from {path}: {e}")
        return {}
    except Exception as e:
        print(f"Unexpected error loading config from {path}: {e}")
        return {}

This function may look simple, but it embodies 3 production-level lessons:

Separation of concerns

Your application logic (FastAPI, inference services) should not know how a YAML file is parsed. They should only receive clean configuration objects.

Fault tolerance

In real deployments:

• configs may be missing
• YAML indentation may break
• a misconfigured CI pipeline may pass an empty file

Returning {} instead of crashing gives you graceful degradation.

Extensibility

Tomorrow you may add:

• JSON config support
• remote config loading (S3, Google Cloud Storage (GCS), Azure Blob)
• encrypted secrets
• multiple config layers

This helper becomes the foundation.

Inside core/config.py, you saw how load_yaml_config() merges YAML values into your Pydantic settings. This is a real-world pattern used in production MLOps stacks like Airflow, FastAPI microservices, Ray Serve, and MLflow.

Adding New Helper Functions

The utilities layer is designed to grow organically as your system grows.

Common helpers you may introduce later include:

String helpers

• text normalization
• input cleaning
• token counting

File helpers

• safe file writes
• temporary directory management
• checksum calculation for model files

Model helpers

• downloading artifacts from cloud storage
• caching models on disk
• validating model signatures

API helpers

• request validation
• standardized error responses
• retry/backoff wrappers around external calls

Monitoring helpers

• timing decorators
• metrics emitters (Prometheus, StatsD, OpenTelemetry)
• latency buckets

All of these belong in one place:

src/utils/

This prevents your service layer or route handlers from becoming cluttered and ensures that common functionality is implemented once and reused everywhere.

Running a FastAPI MLOps Application Locally

At this point, you have a fully structured ML application: configuration, logging, models, service layer, and a clean FastAPI interface. Now it’s time to actually run the system locally.

This section walks you through running the API with Poetry, UV, or PDM, depending on your setup. We’ll conclude with a quick validation test to ensure everything works end-to-end.

Running via Poetry

If you’re using Poetry (recommended for most workflows), your steps are:

# Install dependencies
poetry install

# Activate the environment
poetry shell

# Start the API server
poetry run python src/main.py

You should see log lines like:

INFO - Starting server on 0.0.0.0:8000
INFO - Loaded model: dummy_classifier

Running via UV

If you prefer UV (super-fast installer by Astral), run:

# Create and activate a virtual environment
uv venv
source .venv/bin/activate

# Install project in editable mode
uv pip install -e .

# Start the API
python src/main.py

This path is great for users who want lightweight dependency management without Poetry’s abstraction.

Running Python MLOps Projects with PDM

If your workflow uses PDM, run:

# Install dependencies
pdm install

# Start the server
pdm run python src/main.py

PDM offers a cleaner pyproject-first workflow and works well for CI/CD pipelines that prefer explicit environment setup.

Testing FastAPI Endpoints: Health Check and Prediction API

Once the server is running, validate the system with 2 quick API calls.

Health Check

Open:

http://localhost:8000/health

Expected response:

{"status": "ok"}

This confirms:

• the API is reachable
• config and logger initialized
• FastAPI routes are registered

Prediction Test

Send a prediction request:

curl -X POST "http://localhost:8000/predict?input=This+is+good"

Expected response:

{"prediction": "positive"}

Under the hood:

• the service layer logs the request
• the dummy model classifies sentiment
• the API returns structured JSON

What's next? We recommend PyImageSearch University.

Course information:
86+ total classes • 115+ hours hours of on-demand code walkthrough videos • Last updated: April 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:

• ✓ 86+ courses on essential computer vision, deep learning, and OpenCV topics
• ✓ 86 Certificates of Completion
• ✓ 115+ hours 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 build a clean, scalable foundation for ML systems using real software-engineering practices. You now understand why ML projects must be structured like production services — not experiments — if they are ever going to ship reliably.

We began by exploring the why: ML code becomes maintainable only when you enforce clear boundaries between configuration, logic, services, and I/O. That idea naturally led to the src/ layout, which gave our project a predictable and extensible shape.

You then learned how to manage dependencies using Poetry, UV, or PDM — ensuring that every ML environment is reproducible, isolated, and easy to rebuild. This solved the classic “it works on my machine” trap that haunts ML teams.

Next, we built a robust configuration system using Pydantic BaseSettings, merging defaults, YAML files, and .env variables into a single typed interface. You now have a configuration pattern used by real-world production ML systems.

We also implemented structured logging, enabling the application to communicate what it’s doing internally — a prerequisite for debugging, observability, and monitoring.

From there, you built your first production-style ML API with FastAPI, complete with /health, /predict, and auto-generated documentation. You learned how to expose ML logic cleanly, and why APIs are the interface between ML systems and the real world.

We introduced the Service Layer, showing how routes should delegate to independent business logic so APIs stay thin and models stay swappable. This design decision is what makes the system testable and future-proof.

You then explored model abstraction, using a simple dummy model to illustrate how real models (PyTorch, TensorFlow, ONNX, vLLM, Transformers) can be slotted in without changing the API layer.

Finally, you saw how helper utilities make the system cleaner, and how to run the full application with Poetry, UV, or PDM. The result is a working ML service that looks, behaves, and organizes itself like production-grade software.

By completing this lesson, you’ve built the foundation required for every advanced MLOps practice: testing, performance monitoring, CI/CD, orchestration, and deployment.

You’re now ready for Lesson 2, where we transform this service into a fully tested, validated, and performance-monitored ML system.

Citation Information

Singh, V. “FastAPI for MLOps: Python Project Structure and API Best Practices,” PyImageSearch, S. Huot, A. Sharma, and P. Thakur, eds., 2026, https://pyimg.co/yn8a5

@incollection{Singh_2026_fastapi-for-mlops-python-project-structure,
  author = {Vikram Singh},
  title = {{FastAPI for MLOps: Python Project Structure and API Best Practices}},
  booktitle = {PyImageSearch},
  editor = {Susan Huot and Aditya Sharma and Piyush Thakur},
  year = {2026},
  url = {https://pyimg.co/yn8a5},
}

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 FastAPI for MLOps: Python Project Structure and API Best Practices appeared first on PyImageSearch.

Agentic AI Vision System: Object Segmentation with SAM 3 and Qwen

This lesson is the 4th and final part of our series on SAM 3. In the previous parts, we built a strong foundation for concept-aware segmentation.

In Part 1, we introduced the fundamentals of SAM 3 and explored how it enables concept-based visual understanding and segmentation. We moved beyond fixed labels and used natural language to describe objects.

In Part 2, we extended this idea by introducing multi-modal prompting and interactive segmentation. We combined text, points, and bounding boxes to gain more precise control over segmentation.

In Part 3, we extended this into the temporal domain. We applied SAM 3 to videos and built systems for concept-aware segmentation and object tracking across frames.

In this final part, we take a major step forward. Instead of treating segmentation as a single-step prediction, we introduce an agentic AI system that can reason, verify, and iteratively refine its outputs.

This lesson is the last of a 4-part series on SAM 3:

1. SAM 3: Concept-Based Visual Understanding and Segmentation
2. Advanced SAM 3: Multi-Modal Prompting and Interactive Segmentation
3. SAM 3 for Video: Concept-Aware Segmentation and Object Tracking
4. Agentic AI Vision System: Object Segmentation with SAM 3 and Qwen (this tutorial)

To learn how to build an Agentic AI Vision System with SAM 3 and Qwen, just keep reading.

Looking for the source code to this post?

Jump Right To The Downloads Section

Why Agentic AI Outperforms Traditional Vision Pipelines

Modern computer vision systems are evolving beyond traditional pipelines.

We designed systems where:

• an image is passed to a vision model
• the model produces a prediction
• the pipeline ends there

This approach works well for clearly defined tasks. However, it struggles when tasks require understanding intent, handling ambiguity, or refining outputs.

To address this, we now transition toward agentic AI systems.

Agentic systems are not limited to a single prediction. Instead, they behave more like an iterative reasoning loop.

They can:

• interpret a user request
• select the appropriate models or tools
• evaluate intermediate outputs
• refine their decisions over multiple steps

This shift allows us to build systems that are adaptive, iterative, and self-correcting.

Why Agentic AI Improves Computer Vision and Segmentation Tasks

Vision tasks are often ambiguous.

For example, consider the instruction:

• “the bag on the leftmost side”

A traditional segmentation model cannot directly handle this:

• it expects fixed labels like “bag”
• it does not understand spatial reasoning like “leftmost”

This is where agentic design becomes powerful.

We introduce a Vision-Language Model (VLM) to:

• understand the instruction
• extract the correct intent
• translate it into a form usable by a segmentation model

Then, instead of trusting the output blindly, we:

• verify the result
• refine the input if needed
• retry the process

This creates a loop where the system continuously improves.

What We Will Build: An Agentic AI Vision and Segmentation System

In this lesson, we build an agentic segmentation system that combines reasoning with perception.

The system takes:

• an image
• a natural language instruction

and produces:

• segmentation masks
• bounding boxes
• confidence scores
• a final overlay visualization

Agentic AI Workflow: Vision-Language Reasoning and Segmentation Loop

The pipeline follows these steps:

• User Input: First, we provide an image along with a natural language instruction.
• Instruction Understanding (VLM): Next, the VLM processes both the image and the text. It extracts the core intent and converts it into a short concept.
• Concept Simplification: The system converts complex instructions into concise phrases. For example:
  • “the bag on the leftmost side” → “leftmost bag”
• Segmentation (SAM3): Then, SAM3 uses this concept to generate:
  • segmentation masks
  • bounding boxes
  • confidence scores
• Verification (VLM): After segmentation, the VLM evaluates whether the output matches the instruction.
• Refinement Loop: If the result is incorrect:
  • the VLM refines the concept
  • SAM3 runs again
  • the process repeats
• This loop continues until the result aligns with the user’s intent.

Agentic AI Architecture: Combining VLMs and SAM 3 for Vision

Before implementing the code, we break down the system into its core components.

Vision-Language Model (VLM): The Reasoning Component

The VLM is the reasoning component of our system. It performs 3 key roles:

Instruction Understanding. It interprets the natural language input in the context of the image.

Concept Generation. It converts long instructions into short, structured phrases. For example:

• “the person wearing a red shirt” → “person red shirt”
• “the car in the background” → “background car”

This step is critical because segmentation models perform better with:

• short
• object-centric
• unambiguous phrases

Result Verification. After segmentation, the VLM checks:

• whether the correct object was segmented
• whether spatial or contextual constraints are satisfied

SAM 3: Open-Vocabulary Object Segmentation

SAM3 acts as the perception component.

Unlike traditional segmentation models, SAM3 supports:

• flexible prompts
• open-vocabulary segmentation

This means we are not restricted to predefined classes.

Given a concept phrase, SAM3 produces:

• pixel-level segmentation masks
• bounding boxes
• confidence scores

This makes SAM3 ideal for integration with a language-based reasoning system.

The Agentic Feedback Loop: Reasoning, Verification, and Refinement

The most important part of this system is the agentic loop.

Instead of a linear pipeline, we build a feedback-driven process.

Step-by-step:

• Generate a segmentation concept
• Run segmentation using SAM3
• Evaluate the output using the VLM

If the output is incorrect:

• identify what went wrong
• refine the concept
• retry segmentation

Why Agentic Segmentation Outperforms One-Shot Models

This loop introduces several important capabilities:

• Self-correction: The system can recover from incorrect predictions
• Robustness: It handles ambiguous or complex instructions better
• Generalization: It works with open-ended language instead of fixed labels
• Improved alignment: Outputs better match user intent over iterations

Final Output: Agentic Vision System with Segmentation and Reasoning

By the end of this tutorial, we build a system that:

• understands natural language instructions
• converts them into structured segmentation concepts
• performs open-vocabulary segmentation
• verifies its own outputs
• improves results through iterative refinement

This represents a shift

from:

• static, one-shot predictions

to:

• dynamic, reasoning-driven vision systems

Key Takeaway: VLM + SAM 3 = Intelligent Vision Agent

The real power of this system is not just segmentation.

It is the collaboration between models:

• the VLM provides reasoning
• SAM3 provides perception
• the loop provides intelligence

Together, they form an agentic vision system that can think, act, and improve.

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 have the following libraries installed on your system.

!pip install -q transformers accelerate pillow torch torchvision bitsandbytes

First, we install the transformers library. This library provides access to a wide range of pretrained models, including the Vision-Language Model we will use in this project.

Next, we install accelerate, which helps efficiently run large models across GPUs and manage device placement automatically.

After that, we install pillow, a lightweight Python library used for image loading and processing. We will use this library to read images and prepare them for model inference.

We also install torch, which serves as the core deep learning framework for this project. Both the Vision-Language Model and the segmentation model rely on torch for tensor computations and GPU acceleration.

Along with torch, we install torchvision, which provides datasets, transforms, and model utilities for computer vision tasks.

Finally, we install bitsandbytes. This library enables efficient memory usage when working with large models by supporting quantization and optimized GPU kernels.

The -q flag runs the installation in quiet mode, reducing unnecessary output in the notebook.

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!

Python Setup and Imports for Agentic AI Vision System

Now that our environment is ready, we import the libraries required to build our agentic vision system. These libraries will help us perform deep learning inference, process images, visualize segmentation outputs, and load the models.

import torch
import numpy as np
import os
import json
from PIL import Image, ImageDraw
import matplotlib
import matplotlib.pyplot as plt
from transformers import (
      AutoProcessor,
   Qwen2_5_VLForConditionalGeneration,
   Sam3Model,
   Sam3Processor,
)

First, we import torch. This is the primary deep learning framework used to run both the Vision-Language Model and the segmentation model. PyTorch handles tensor computations and GPU acceleration during inference.

Next, we import numpy, a popular library for numerical computing in Python. We will use NumPy when working with arrays such as segmentation masks and bounding boxes returned by the segmentation model.

After that, we import the os and json libraries. The os module helps us manage file paths and directories, while the json module allows us to parse structured responses generated by the Vision-Language Model.

Next, we import Image and ImageDraw from the Pillow library. Pillow is a lightweight image processing library that allows us to load, manipulate, and display images. In this project, we will use it to read input images and create segmentation overlays.

Then, we import matplotlib, which we will use to visualize the results. Specifically, we use matplotlib.pyplot to create figures that display the original image, bounding boxes, and segmentation masks.

Finally, we import several classes from the transformers library. These classes allow us to load and run the models used in our system.

• The AutoProcessor class automatically prepares inputs for multimodal models by handling both text and image preprocessing.
• The Qwen2_5_VLForConditionalGeneration class loads the Qwen2.5-VL Vision-Language Model, which will interpret user instructions and generate segmentation prompts.
• The Sam3Model and Sam3Processor classes load the SAM3 segmentation model and prepare its inputs.

Before loading the models, we configure PyTorch to use optimized GPU settings. These settings help improve inference performance, especially when running large multimodal models.

torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype  = torch.bfloat16 if device == "cuda" else torch.float32
print(f"Using device: {device}, dtype: {dtype}")

First, we enable TensorFloat-32 (TF32) support in PyTorch. TF32 is a numerical format supported by modern NVIDIA GPUs. It allows faster matrix multiplications during deep learning inference while maintaining good numerical stability. Since large models perform many matrix operations, enabling TF32 can significantly improve performance.

Next, we determine which device will be used for inference. Here, we check whether a CUDA-enabled GPU is available. If a GPU is detected, the system runs on "cuda". Otherwise, it falls back to the CPU.

After that, we configure the tensor precision. When running on a GPU, we use bfloat16 precision. This reduces memory usage and speeds up computation while preserving enough numerical accuracy for inference tasks.

If the system runs on a CPU, we instead use the standard float32 precision, which ensures compatibility with CPU computations.

Finally, we print the device configuration. This helps confirm whether the system is using the GPU and which precision mode is active. This information is useful when debugging performance or memory issues during model inference.

Loading SAM 3 and Qwen Vision-Language Models in Transformers

Now that the environment is configured, we load the two core models used in our agentic vision system: a Vision-Language Model (VLM) and a segmentation model.

The VLM will interpret the user’s instruction and generate a clean segmentation concept. The segmentation model will then use that concept to detect and segment objects in the image.

VLM_MODEL_ID = "Qwen/Qwen2.5-VL-7B-Instruct"  # swap for Qwen/Qwen3-VL-8B once released in transformers
SAM_MODEL_ID = "facebook/sam3"

print("Loading VLM...")
vlm_processor = AutoProcessor.from_pretrained(VLM_MODEL_ID, trust_remote_code=True)
vlm_model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
   VLM_MODEL_ID,
   device_map="auto",
   torch_dtype=dtype,
   trust_remote_code=True,
)
vlm_model.eval()
print("VLM loaded.")

print("Loading SAM3...")
sam_processor = Sam3Processor.from_pretrained(SAM_MODEL_ID)
sam_model = Sam3Model.from_pretrained(SAM_MODEL_ID, torch_dtype=dtype).to(device)
sam_model.eval()
print("SAM3 loaded.")

First, we define the model identifiers. These identifiers correspond to the pretrained models hosted on the Hugging Face model hub.

The Qwen2.5-VL-7B-Instruct model is a Vision-Language Model capable of understanding both images and text instructions. We will use this model to interpret the user’s request and generate segmentation prompts.

The second model, SAM3, is an open-vocabulary segmentation model that can segment objects based on text prompts.

Next, we load the Vision-Language Model. We first load the processor associated with the model. The processor prepares the inputs required by the VLM, including tokenizing text prompts and preprocessing images.

The trust_remote_code=True argument allows the Transformers library to load custom processing code provided by the model repository.

Next, we load the model itself. The from_pretrained() method downloads the pretrained model weights and initializes the model architecture.

The device_map="auto" argument automatically distributes the model across available devices, which is useful when working with large models that require GPU memory.

We also specify torch_dtype=dtype, which ensures the model runs using the precision we configured earlier: bfloat16 on GPU or float32 on CPU.

After loading the model, we switch it to evaluation mode. Evaluation mode disables training-specific behaviors such as dropout, ensuring consistent inference results.

Next, we load the segmentation model. Similar to the VLM, we first load the Sam3Processor. This processor handles preprocessing tasks such as preparing the input image and formatting segmentation prompts.

Next, we load the SAM3 model. The from_pretrained() function loads the segmentation model weights, and we move the model to the appropriate device using .to(device).

Finally, we set the model to evaluation mode. At this point, both models are fully initialized. The Vision-Language Model will interpret user instructions, while SAM3 will perform open-vocabulary segmentation based on those instructions.

Implementing VLM Inference for Agentic Vision Reasoning with Qwen2.5-VL

Now that our models are loaded, we implement a helper function that allows us to run inference using the Vision-Language Model. This function will take an image and a list of chat messages as input and return the model’s response.

In our agentic pipeline, this function plays a very important role. We will use it to:

• extract a clean segmentation prompt from the user instruction
• refine prompts if segmentation fails
• verify whether the segmentation results match the user intent

def vlm_generate(image: Image.Image, messages: list, max_new_tokens: int = 512) -> str:
   """
   Mirrors: send_generate_request()
   Runs VLM inference given a list of chat messages and returns the reply string.
   """
   text_input = vlm_processor.apply_chat_template(
       messages, tokenize=False, add_generation_prompt=True
   )
   inputs = vlm_processor(
       text=[text_input],
       images=[image],
       return_tensors="pt",
   )
   inputs = {k: v.to(vlm_model.device) for k, v in inputs.items()}
   input_len = inputs["input_ids"].shape[1]

   with torch.no_grad():
       generated_ids = vlm_model.generate(
           **inputs,
           max_new_tokens=max_new_tokens,
           do_sample=False,
       )

   new_tokens = generated_ids[0][input_len:]
   return vlm_processor.tokenizer.decode(new_tokens, skip_special_tokens=True).strip()

First, we define the function vlm_generate. This function takes three inputs:

• image: the input image that the model will analyze
• messages: a list of chat-style prompts used to guide the model
• max_new_tokens: the maximum number of tokens the model can generate

The function returns a string response produced by the Vision-Language Model.

Next, we convert the chat messages into the format expected by the model. Many modern Vision-Language Models use a chat-style interface similar to conversational AI systems. The apply_chat_template() method converts the list of messages into a properly formatted text prompt that the model understands.

The argument add_generation_prompt=True tells the processor that the model should generate a response after the provided messages.

Next, we prepare the inputs for the model. Here, we pass both the text prompt and the image to the processor. The processor converts these inputs into tensors that can be processed by the model. The argument return_tensors="pt" ensures the outputs are returned as PyTorch tensors.

Next, we move the tensors to the same device as the model. This step ensures that both the model and the input tensors reside on the same device, either the GPU or CPU.

After that, we store the length of the input tokens. This value helps us determine which tokens belong to the model’s generated response, rather than the original prompt.

Next, we perform inference using the model. We use torch.no_grad() to disable gradient computations. Since we are only performing inference, this reduces memory usage and improves performance.

Inside this block, we generate the model’s output. The generate() function performs autoregressive text generation. The parameter max_new_tokens limits the length of the generated response. We also set do_sample=False, which ensures deterministic outputs instead of random sampling.

Next, we extract only the tokens generated by the model. This removes the original prompt tokens, leaving only the newly generated tokens.

Finally, we convert the generated tokens into readable text. The decode() method converts token IDs back into text. We also remove special tokens and strip unnecessary whitespace.

At this point, the function returns the final response generated by the Vision-Language Model.

This function will serve as the core interface between our agentic system and the Vision-Language Model. In the next sections, we will use it to extract segmentation prompts and evaluate the outputs produced by the segmentation model.

Implementing the SAM 3 Text-Prompted Segmentation Function

Now, we implement a helper function that runs segmentation using the SAM3 model. This function will take an input image and optional prompts, run the SAM3 model, and return the segmentation results.

In our agentic pipeline, this function serves as the tool used by the agent to perform segmentation.

Specifically, it returns three important outputs:

• segmentation masks
• bounding boxes
• confidence scores

def call_sam(
   image: Image.Image,
   text_prompt: str   = None,
   input_boxes        = None,   # list of [x1,y1,x2,y2]
   input_boxes_labels = None,   # list of 0/1 labels per box
   threshold: float   = 0.5,
) -> dict:
   """
   Mirrors: call_sam_service()
   Returns dict with keys: masks, boxes, scores (all as numpy arrays).
   """
   kwargs = dict(images=image, return_tensors="pt")
   if text_prompt:
       kwargs["text"] = text_prompt
   if input_boxes is not None:
       kwargs["input_boxes"] = [input_boxes]
       kwargs["input_boxes_labels"] = [input_boxes_labels or [1] * len(input_boxes)]

   inputs = sam_processor(**kwargs).to(device)

   with torch.no_grad():
       outputs = sam_model(**inputs)

   results = sam_processor.post_process_instance_segmentation(
       outputs,
       threshold=threshold,
       mask_threshold=0.5,
       target_sizes=inputs.get("original_sizes").tolist(),
   )[0]

   return {
       "masks":  results["masks"].cpu().numpy(),                          # [N, H, W] bool
       "boxes":  results["boxes"].cpu().to(torch.float32).numpy(),        # [N, 4]    xyxy
       "scores": results["scores"].cpu().to(torch.float32).numpy(),       # [N]
   }

First, we define the function call_sam. This function accepts several inputs:

• The image parameter is the input image that we want to segment.
• The text_prompt parameter allows us to perform concept-based segmentation. SAM3 can segment objects using natural language prompts such as "bag" or "leftmost bag".
• The input_boxes parameter allows us to guide the segmentation model using bounding boxes. Each box is defined by four coordinates: [x1, y1, x2, y2]
• Similarly, input_boxes_labels specifies whether each box corresponds to a positive or negative prompt.
• Finally, the threshold parameter determines the confidence threshold used when filtering segmentation results.

Next, we prepare the inputs required by the SAM3 processor.

Here, we create a dictionary containing the image input. The return_tensors="pt" argument ensures that the processed outputs are returned as PyTorch tensors.

If a text prompt is provided, we include it in the input dictionary. This allows SAM3 to perform text-guided segmentation.

Next, we check whether bounding boxes are provided. If bounding boxes exist, we pass them to the processor along with their labels. If no labels are specified, we automatically assign positive labels (1) to all boxes.

Next, we preprocess the inputs using the SAM3 processor. The processor converts the image, prompts, and bounding boxes into tensors that the model can understand. We also move these tensors to the selected device (GPU or CPU).

Now we perform inference using SAM3. We wrap the inference step inside torch.no_grad() to disable gradient calculations. Since we are performing inference only, this improves performance and reduces memory usage. The model returns raw segmentation outputs.

Next, we convert the raw model outputs into usable segmentation results. The post_process_instance_segmentation() function performs several important tasks:

• filters predictions using the confidence threshold
• converts predicted masks to the correct image resolution
• extracts bounding boxes and scores

The [0] index retrieves the results corresponding to the input image.

Finally, we return the segmentation results. The function returns a dictionary containing three elements.

• The masks array contains the segmentation masks with shape: [N, H, W] where N represents the number of detected objects.
• The boxes array contains the bounding box coordinates in the format: [x1, y1, x2, y2]
• Finally, the scores array contains the confidence score for each detected object.

We also move the tensors to the CPU and convert them into NumPy arrays. This makes them easier to process and visualize in later steps.

At this point, the call_sam() function provides a simple interface for running SAM3 segmentation within our agentic vision pipeline.

Implementing the Agentic AI Segmentation Pipeline with Iterative Refinement

Now we implement the core function of our system. This function orchestrates the entire agentic workflow by combining the Vision-Language Model and the segmentation model.

Instead of running segmentation only once, the system follows an agentic loop where the Vision-Language Model interprets the user request, runs segmentation, verifies the result, and refines the prompt if needed.

def run_single_image_inference(
   image_path: str,
   user_prompt: str,
   max_agent_rounds: int = 3,
   seg_threshold: float  = 0.5,
   output_dir: str       = "agent_output",
   debug: bool           = True,
) -> str | None:
   """
   Mirrors: run_single_image_inference() from sam3.agent.inference

   Agentic loop:
     Round 1 — VLM reads image + user prompt → produces a concise SAM3 concept phrase
     Round 2 — SAM3 segments with that phrase → VLM verifies / refines if needed
     Round N — repeat until VLM is satisfied or max_agent_rounds reached
   Returns path to the saved output image (or None on failure).
   """
   os.makedirs(output_dir, exist_ok=True)
   image = Image.open(image_path).convert("RGB")

   # ── Round 1: VLM extracts a clean SAM3 text prompt ──────────────────────
   extraction_messages = [
       {
           "role": "system",
           "content": (
               "You are a precise vision assistant. "
               "Your job is to convert a user's free-form description into a SHORT, "
               "clean object concept phrase suitable for an open-vocabulary segmentation model. "
               "Reply with ONLY a JSON object: {\"sam_prompt\": \"<phrase>\"}. "
               "No explanation, no markdown, just the JSON."
           ),
       },
       {
           "role": "user",
           "content": [
               {"type": "image", "image": image},
               {"type": "text",  "text": f"User description: \"{user_prompt}\""},
           ],
       },
   ]

The run_single_image_inference function serves as the main entry point of our agentic vision system. It accepts several inputs:

• image_path: the path to the image we want to analyze
• user_prompt: the natural language description of the object to segment
• max_agent_rounds: the maximum number of refinement iterations
• seg_threshold: the confidence threshold for segmentation
• output_dir: the directory where the output image will be saved
• debug: a flag that enables detailed logging

The function returns the path of the saved output image or None if segmentation fails.

First, we create the output directory and load the image. The os.makedirs() function ensures that the output directory exists. If the directory already exists, the exist_ok=True argument prevents an error. Next, we open the input image using Pillow and convert it to RGB format.

Here, we define a system message that instructs the Vision-Language Model to convert the user description into a short concept phrase. The SAM3 model performs better with short noun-style prompts such as:

• leftmost bag
• red apple
• wooden chair

rather than long sentences.

We also include the user input. This message contains both the image and the user instruction.

if debug:
       print(f"\n[Agent] Round 1 — extracting SAM3 prompt from: '{user_prompt}'")

   vlm_reply = vlm_generate(image, extraction_messages)
   if debug:
       print(f"[Agent] VLM raw reply: {vlm_reply}")

   # Parse the JSON; fall back to raw reply if needed
   try:
       clean = vlm_reply.strip().lstrip("```json").rstrip("```").strip()
       sam_prompt = json.loads(clean)["sam_prompt"]
   except Exception:
       sam_prompt = user_prompt  # graceful fallback
   if debug:
       print(f"[Agent] SAM3 prompt → '{sam_prompt}'")

Next, we call the VLM inference function. The Vision-Language Model analyzes the image and generates a clean segmentation prompt.

For example:

User prompt: "the bag on the leftmost side"
Model output: {"sam_prompt": "leftmost bag"}

Next, we extract the segmentation prompt from the JSON response. This step removes formatting artifacts and converts the JSON string into a Python dictionary.

If the response cannot be parsed, we fall back to the original user prompt.

# ── Agentic segmentation loop ────────────────────────────────────────────
   sam_result = None
   final_prompt = sam_prompt

   for round_idx in range(max_agent_rounds):
       if debug:
           print(f"\n[Agent] Round {round_idx + 2} — calling SAM3 with '{final_prompt}'")

       sam_result = call_sam(image, text_prompt=final_prompt, threshold=seg_threshold)
       n_masks = len(sam_result["masks"])
       if debug:
           print(f"[Agent] SAM3 found {n_masks} instance(s)")

Now we begin the agentic segmentation loop. Here, we initialize two variables:

• sam_result: stores the segmentation output
• final_prompt: stores the prompt used for segmentation

Next, we enter the iterative loop. This loop allows the system to refine segmentation prompts up to a maximum number of rounds.

Inside the loop, we call the SAM3 segmentation function. This function returns segmentation results including masks, bounding boxes, and confidence scores.

Next, we count the number of detected objects. This value helps determine whether the segmentation succeeded.

       # ── Verification: ask VLM if the result looks right ─────────────────
       if n_masks == 0:
           # No masks found — ask VLM to rephrase
           refine_messages = [
               {
                   "role": "system",
                   "content": (
                       "You are a vision assistant helping refine segmentation prompts. "
                       "The segmentation model found NO objects. "
                       "Suggest a simpler or broader alternative concept phrase. "
                       "Reply ONLY with JSON: {\"sam_prompt\": \"<phrase>\"}."
                   ),
               },
               {
                   "role": "user",
                   "content": [
                       {"type": "image", "image": image},
                       {"type": "text",  "text": (
                           f"Original user intent: \"{user_prompt}\". "
                           f"Failed prompt: \"{final_prompt}\". "
                           "Suggest a better phrase."
                       )},
                   ],
               },
           ]
           vlm_reply = vlm_generate(image, refine_messages)
           if debug:
               print(f"[Agent] VLM refine reply: {vlm_reply}")
           try:
               clean = vlm_reply.strip().lstrip("```json").rstrip("```").strip()
               final_prompt = json.loads(clean)["sam_prompt"]
           except Exception:
               break  # give up if we can't parse
       

If SAM3 fails to detect any objects, we ask the Vision-Language Model to refine the segmentation prompt. We construct a new prompt asking the model to generate a simpler or broader concept phrase.

For example:

Original prompt: "leftmost brown grocery bag"
Suggested prompt: "bag"

The VLM then generates a new segmentation prompt, and the loop repeats.

else:
           # We have masks — ask VLM to verify they match the user intent
           verify_messages = [
               {
                   "role": "system",
                   "content": (
                       "You are a vision QA assistant. "
                       "Given the original user intent and the segmentation result metadata, "
                       "decide if the segmentation is correct. "
                       "Reply ONLY with JSON: {\"ok\": true/false, \"reason\": \"...\", \"sam_prompt\": \"<refined phrase if not ok>\"}."
                   ),
               },
               {
                   "role": "user",
                   "content": [
                       {"type": "image", "image": image},
                       {"type": "text",  "text": (
                           f"User intent: \"{user_prompt}\".\n"
                           f"SAM3 was given prompt: \"{final_prompt}\".\n"
                           f"Result: {n_masks} mask(s) found, "
                           f"scores: {sam_result['scores'].tolist()}, "
                           f"boxes: {sam_result['boxes'].tolist()}.\n"
                           "Is this correct? If yes, ok=true. If not, provide a better sam_prompt."
                       )},
                   ],
               },
           ]
           vlm_reply = vlm_generate(image, verify_messages, max_new_tokens=256)
           if debug:
               print(f"[Agent] VLM verify reply: {vlm_reply}")
           try:
               clean = vlm_reply.strip().lstrip("```json").rstrip("```").strip()
               verdict = json.loads(clean)
               if verdict.get("ok", True):
                   if debug:
                       print("[Agent] VLM verified result ✓ — stopping.")
                   break
               else:
                   final_prompt = verdict.get("sam_prompt", final_prompt)
                   if debug:
                       print(f"[Agent] VLM says not ok → retrying with '{final_prompt}'")
           except Exception:
               break  # can't parse verdict, accept current result

If SAM3 successfully detects objects, we verify whether the result matches the user intent.

In this step, we ask the Vision-Language Model to evaluate the segmentation results.

The model receives:

• the original user instruction
• the segmentation prompt used
• the number of detected masks
• the confidence scores
• the bounding boxes

Based on this information, the model decides whether the segmentation result is correct.

The model returns a JSON response such as:

{
"ok": true,
"reason": "correct object detected"
}

or

{
"ok": false,
"sam_prompt": "bag"
}

If the segmentation is incorrect, the system updates the segmentation prompt. The loop then repeats using the new prompt. If the segmentation result is correct, the loop stops. This verification step allows the system to self-correct its segmentation decisions.

   # ── Render and save output ───────────────────────────────────────────────
   if sam_result is None or len(sam_result["masks"]) == 0:
       print("[Agent] No masks produced — check your prompt or image.")
       return None

   output_path = os.path.join(
       output_dir,
       os.path.splitext(os.path.basename(image_path))[0] + "_segmented.png"
   )
   _save_overlay(image, sam_result, output_path, title=f'"{user_prompt}"')
   print(f"\n[Agent] Output saved → {output_path}")
   return output_path

After the agentic loop finishes, we check whether segmentation succeeded. If no objects were detected, the function returns None. Otherwise, we generate the output image path.

Finally, we visualize the segmentation results. This function creates an image containing the segmentation masks and bounding boxes. The result is saved to disk.

This function implements the agentic reasoning loop that makes our system powerful.

Instead of relying on a single segmentation attempt, the system:

• interprets the user request
• generates a segmentation prompt
• runs segmentation
• evaluates the results
• refines the prompt if necessary

This iterative process allows the system to produce more accurate results and demonstrates how multiple AI models can collaborate within an agentic vision pipeline.

Visualizing and Saving the Segmentation Results

After running the agentic segmentation pipeline, we want to visualize the results in a clear and interpretable way. For this purpose, we implement a helper function that overlays the segmentation masks and bounding boxes on top of the original image.

This function generates a side-by-side visualization showing both the detected bounding boxes and the segmentation masks.

def _save_overlay(image: Image.Image, sam_result: dict, output_path: str, title: str = ""):
   masks  = sam_result["masks"]
   boxes  = sam_result["boxes"]
   scores = sam_result["scores"]

   fig, axes = plt.subplots(1, 2, figsize=(16, 8))

   # Left: original + boxes
   axes[0].imshow(image)
   axes[0].set_title(f"Detected boxes  |  {title}", fontsize=11)
   axes[0].axis("off")
   cmap = matplotlib.colormaps.get_cmap("rainbow").resampled(max(len(masks), 1))
   for i, (box, score) in enumerate(zip(boxes, scores)):
       x1, y1, x2, y2 = box
       color = cmap(i)[:3]
       rect = plt.Rectangle(
           (x1, y1), x2 - x1, y2 - y1,
           linewidth=2, edgecolor=color, facecolor="none"
       )
       axes[0].add_patch(rect)
       axes[0].text(x1, y1 - 4, f"{score:.2f}", color=color, fontsize=9, fontweight="bold")

   # Right: mask overlay
   composite = image.convert("RGBA")
   for i, mask in enumerate(masks):
       color = tuple(int(c * 255) for c in cmap(i)[:3])
       mask_img = Image.fromarray((mask * 255).astype(np.uint8))
       overlay  = Image.new("RGBA", composite.size, color + (0,))
       overlay.putalpha(mask_img.point(lambda v: int(v * 0.5)))
       composite = Image.alpha_composite(composite, overlay)

   axes[1].imshow(composite)
   axes[1].set_title(f"SAM3 masks  ({len(masks)} instance(s))", fontsize=11)
   axes[1].axis("off")

   plt.tight_layout()
   plt.savefig(output_path, dpi=150, bbox_inches="tight")
   plt.close()

We begin by defining the _save_overlay function, which takes the original image, the segmentation output from SAM3, the output path, and an optional title. From the segmentation results, we extract the masks, bounding boxes, and confidence scores. The masks represent pixel-level regions for each detected object, the boxes define object boundaries, and the scores indicate how confident the model is for each detection.

To visualize these results, we create a figure with two side-by-side panels. The left panel displays the original image along with bounding boxes, while the right panel shows the segmentation masks overlaid on the image.

The process starts by rendering the original image and assigning a distinct color to each detected object using a colormap. For every detection, we draw a rectangle corresponding to its bounding box and place the confidence score near it. This provides a quick overview of what the model has detected and how reliable those detections are.

For the mask visualization, the image is first converted to RGBA format so that transparent overlays can be applied. Each segmentation mask is then assigned a color, converted into an image, and used to create a semi-transparent overlay. These overlays are composited onto the original image, allowing the segmented regions to stand out while still preserving the underlying content.

The final composite is displayed in the second panel, along with the number of detected instances. The visualization is then saved to disk using a resolution of 150 DPI for clarity, with tight_layout() ensuring proper spacing and bbox_inches="tight" removing unnecessary margins. The figure is closed afterward to free up memory.

This results in a clean and intuitive visualization that combines bounding boxes, confidence scores, and segmentation masks, making it easy to verify the model’s predictions.

Running the Agentic AI Vision System on Real Images

Now that we have implemented all the components of our pipeline, we can run the complete agentic vision system on an example image.

In this step, we provide an image along with a natural language instruction and let the system handle the rest.

output_image_path = run_single_image_inference(
   image_path  = "/content/groceries.jpg",
   user_prompt = "the bag on the leftmost side",
   max_agent_rounds = 3,
   seg_threshold    = 0.5,
   output_dir       = "agent_output",
   debug            = True,
)

if output_image_path:
   img = Image.open(output_image_path)
   img.show()

We begin by calling the run_single_image_inference() function, which executes the complete agentic pipeline. The input image is provided through the image_path parameter, and in this example, we use groceries.jpg. Along with the image, we pass a natural language instruction — “the bag on the leftmost side”. This instruction is intentionally written in free-form language to demonstrate how the system can interpret human-like queries.

The pipeline is configured to allow up to three refinement iterations using max_agent_rounds=3. A confidence threshold of 0.5 is used to filter segmentation results, and the final output is saved to the agent_output directory. Debugging is enabled to log intermediate steps such as prompt generation, segmentation outputs, and verification decisions.

Once the pipeline runs, it returns the path to the output image if segmentation is successful. We then load this image using Pillow and display it. The final visualization includes bounding boxes around detected objects, segmentation masks overlaid on the image, and confidence scores for each detection.

Under the hood, the system follows an iterative process. The Vision-Language Model first analyzes the image and converts the user’s instruction into a concise segmentation prompt. This prompt is passed to SAM3, which generates segmentation masks. The result is then evaluated by the Vision-Language Model to determine whether it matches the user’s intent. If the output is not satisfactory, the prompt is refined and the process repeats. Once the result is verified, the system produces the final visualization and saves it to disk.

Agentic Segmentation Output: Iterative Prompt Refinement in Action

The input image (Figure 1) shows multiple grocery bags placed inside the trunk of a car.

We provide the following natural language instruction:

"the bag on the leftmost side"

This instruction is not a fixed label. Instead, it includes spatial reasoning, which makes the task more challenging for standard segmentation models.

Now let’s examine how the system processes this instruction.

[Agent] Round 1 — extracting SAM3 prompt from: 'the bag on the leftmost side'
[Agent] VLM raw reply: {"sam_prompt": "leftmost paper bag"}

First, the Vision-Language Model interprets the instruction and generates an initial segmentation prompt:

[Agent] SAM3 prompt -> 'leftmost paper bag'

[Agent] Round 2 — calling SAM3 with 'leftmost paper bag'
[Agent] SAM3 found 0 instance(s)

Next, SAM3 attempts segmentation using this prompt.

However, no objects are detected.

This shows an important limitation: SAM3 is sensitive to how the prompt is phrased.

[Agent] VLM refine reply: {"sam_prompt": "leftmost brown paper bag"}

The system does not stop here.

Instead, the Vision-Language Model refines the prompt by adding more descriptive information.

[Agent] Round 3 — calling SAM3 with 'leftmost brown paper bag'
[Agent] SAM3 found 0 instance(s)

Again, SAM3 fails to detect any objects.

At this point, we observe something important: More detailed prompts do not always improve segmentation.

[Agent] VLM refine reply: {"sam_prompt": "leftmost bag"}

Now, the model simplifies the prompt.

This step is critical. Instead of making the prompt more complex, the system makes it more general.

[Agent] Round 4 — calling SAM3 with 'leftmost bag'
[Agent] SAM3 found 1 instance(s)

This time, SAM3 successfully detects the object.

[Agent] VLM verify reply: {
 "ok": true,
 "reason": "The segmentation correctly identifies the leftmost bag as per the user's intent."
 "sam_prompt": ""
}

Finally, the Vision-Language Model verifies the result and confirms that the segmentation is correct.

The agentic loop stops here, and the system saves the final output image with a bounding box and segmentation mask overlaid on the input image.

The output image (Figure 3) shows:

• the detected bounding box around the leftmost bag
• the segmentation mask highlighted in color
• the correct object selected based on the user’s instruction

What's next? We recommend PyImageSearch University.

Course information:
86+ total classes • 115+ hours hours of on-demand code walkthrough videos • Last updated: April 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:

• ✓ 86+ courses on essential computer vision, deep learning, and OpenCV topics
• ✓ 86 Certificates of Completion
• ✓ 115+ hours 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, we built an agentic AI vision system that combines a Vision-Language Model with a segmentation model to solve a real-world problem.

Instead of relying on a single model, we designed a pipeline where multiple components work together in a loop. This allows the system to not only perform segmentation, but also understand instructions, evaluate results, and improve itself automatically.

First, we used a Vision-Language Model to interpret the user’s natural language query and convert it into a clean segmentation prompt.

Next, we used SAM3 to perform open-vocabulary segmentation using that prompt.

Then, we introduced an agentic loop where the Vision-Language Model verifies the segmentation output and refines the prompt if necessary.

Finally, we visualized the results by overlaying bounding boxes and segmentation masks on the original image.

This approach highlights an important shift in computer vision. Instead of building static pipelines, we are now moving toward interactive and self-correcting systems that can adapt to user intent.

Such systems can be extended to a wide range of applications, including:

• interactive image editing
• robotics and autonomous perception
• visual assistants
• multimodal search systems

In the future, we can further improve this system by:

• adding support for multiple images or video inputs
• integrating more tools into the agent loop
• introducing memory for long-term reasoning
• optimizing inference for real-time applications

By combining Vision-Language Models with powerful segmentation models, we take a step closer to building intelligent visual systems that can understand and act on human instructions.

This represents the foundation of next-generation AI systems.

Citation Information

Thakur, P. “Agentic AI Vision System: Object Segmentation with SAM 3 and Qwen,” PyImageSearch, S. Huot, G. Kudriavtsev, and A. Sharma, eds., 2026, https://pyimg.co/ohlwd

@incollection{Thakur_2026_building-an-agentic-ai-vision-system-with-sam-3-and-qwen,
  author = {Piyush Thakur},
  title = {{Agentic AI Vision System: Object Segmentation with SAM 3 and Qwen}},
  booktitle = {PyImageSearch},
  editor = {Susan Huot and Georgii Kudriavtsev and Aditya Sharma},
  year = {2026},
  url = {https://pyimg.co/ohlwd},
}

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 Agentic AI Vision System: Object Segmentation with SAM 3 and Qwen appeared first on PyImageSearch.

Autoregressive Model Limits and Multi-Token Prediction in DeepSeek-V3

In the first three parts of this series, we built the foundation of DeepSeek-V3 by implementing its configuration and Rotary Positional Embeddings (RoPE), exploring the efficiency gains of Multi-Head Latent Attention (MLA), and scaling capacity through the Mixture of Experts (MoE). Each of these components adds a crucial piece to the puzzle, progressively shaping a model that balances performance, scalability, and efficiency. With these building blocks in place, we are now ready to tackle another defining innovation: Multi-Token Prediction (MTP).

Unlike traditional autoregressive models that predict one token at a time, MTP enables DeepSeek-V3 to forecast multiple tokens simultaneously, significantly accelerating training and inference. This approach not only reduces computational overhead but also improves the model’s ability to capture richer contextual patterns across sequences.

In this lesson, we will explore the theory behind MTP, examine why it represents a leap forward in language modeling, and implement it step by step. As with the earlier lessons, this installment continues our broader mission to reconstruct DeepSeek-V3 from scratch, showing how innovations including RoPE, MLA, MoE, and now MTP fit together into a cohesive architecture that will culminate in the assembly and training of the full model.

This lesson is the 4th in a 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 (this tutorial)
5. Lesson 5
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

Why Next-Token Prediction Limits DeepSeek-V3

Traditional language models are trained with a simple objective: given tokens , predict the next token . Mathematically, we maximize:

.

This autoregressive factorization is elegant and has proven remarkably effective. However, it has a fundamental limitation: the model only receives a training signal for immediate next-token prediction. It never explicitly learns to plan multiple steps ahead.

Consider generating the sentence: “The cat sat on the mat because it was comfortable.” When predicting “because,” the model should already be considering how the sentence will complete — including the subordinate clause, the pronoun reference, and the conclusion. But with next-token prediction alone, there’s no explicit gradient signal encouraging this forward planning. The model might learn it implicitly through exposure to many examples, but we’re not directly optimizing for it.

This limitation becomes especially apparent in tasks requiring long-term coherence (e.g., story generation, multi-paragraph reasoning, or code generation), where later statements must be consistent with earlier declarations. The model can easily generate locally fluent text that globally contradicts itself because its training objective only looks one token ahead.

Multi-Token Prediction in DeepSeek-V3: Predicting Multiple Tokens Ahead

Multi-Token Prediction (Figure 1) addresses this by adding auxiliary prediction heads that forecast multiple tokens into the future. Alongside the standard prediction , we also predict:

and so on for tokens ahead. Critically, these predictions are computed in parallel during training (not autoregressively) — we know all ground truth tokens, so we can supervise all predictions simultaneously.

The complete training objective becomes:

,

where is the number of future tokens we predict, are weighting coefficients (typically decreasing with distance: ), and we’ve explicitly shown that predictions at depth condition on both the context up to position and the intermediate tokens up to .

DeepSeek-V3 Architecture: Multi-Token Prediction Heads Explained

Implementing MTP requires architectural additions. We can’t just reuse the main language modeling head for future predictions — we need to condition on the intermediate tokens. DeepSeek-V3 implements this through a hierarchy of prediction heads, each specialized for a particular future depth.

Head Architecture: For predicting tokens ahead, we have a head that combines:

• The hidden representation from the Transformer at position :
• The embedding of the token at position :

The combination follows:

This combined representation is then processed through a mini-Transformer (lightweight attention and feedforward layers) before projecting to the vocabulary:

The intuition is powerful: to predict token , we start with the representation at position (encoding all context), incorporate the embedding of token (telling us what word we’ve just generated), process through a small Transformer (allowing the model to refine this combination), and project to vocabulary (producing logits over the vocabulary). This architecture naturally encourages forward planning — the model must learn representations at position that are useful for predictions multiple steps ahead.

Gradient Insights for Multi-Token Prediction in DeepSeek-V3

From an optimization perspective, MTP provides richer gradient signals. In standard training, only the hidden representation receives gradients from predicting . With MTP, also receives gradients from predicting . These additional gradients encourage to encode information relevant not just for the immediate next token, but for multiple future tokens.

Moreover, the gradients from future predictions flow through different pathways — through the MTP heads’ mini-Transformers. This creates a form of multi-task learning in which different prediction depths impose distinct consistency constraints on the learned representations. A representation that works well for predicting 1 token ahead might not be good for predicting 5 tokens ahead; MTP encourages learning representations that support both.

We can think of this as adding an implicit regularizer. The additional prediction objectives constrain the learned representations to be more structured, more forward-looking, and more globally coherent. It’s similar in spirit to multi-task learning, where auxiliary tasks improve representation quality even if we care primarily about one main task.

DeepSeek-V3 Training vs. Inference: How MTP Changes Both

During Training: We compute all predictions in parallel. For a sequence of length , we predict:

• Main head: positions 1 through predict positions 2 through
• Depth-1 head: positions 1 through predict positions 3 through
• Depth-2 head: positions 1 through predict positions 4 through

Each prediction uses the ground truth intermediate tokens (available during training), so there’s no error accumulation. The losses are computed independently and summed with appropriate weights.

During Inference: Interestingly, MTP heads are typically not used during autoregressive generation. Once training is complete, we generate text using only the main prediction head in the standard autoregressive manner. The MTP heads have served their purpose by improving the learned representations; we don’t need their multi-step predictions at inference time.

This is computationally appealing: we get the benefits of MTP (better representations, improved coherence) during training, but inference remains as efficient as a standard language model. There’s no additional computational cost at deployment.

Multi-Token Prediction Loss Weighting and Decay for DeepSeek-V3

The weighting coefficients are important hyperparameters. Intuitively, predictions further in the future are harder and less reliable, so we should weight them less heavily. A common scheme is exponential decay:

where . For example, with :

• Depth 1 (predicting from ): weight 1.0
• Depth 2 (predicting from ): weight 0.5
• Depth 3 (predicting from ): weight 0.25

In our implementation, we use a simpler approach: uniform weighting of 0.3 for all MTP losses relative to the main loss. This is less sophisticated but easier to tune and still provides the core benefits.

Step-by-Step Implementation of Multi-Token Prediction Heads in DeepSeek-V3

Let’s implement the complete MTP system:

class MultiTokenPredictionHead(nn.Module):
    """
    Multi-Token Prediction Head

    Each head predicts a token at a specific future position.
    Combines previous hidden state with future token embedding.
    """
    def __init__(self, config: DeepSeekConfig, depth: int):
        super().__init__()
        self.depth = depth
        self.n_embd = config.n_embd

        # Combine previous hidden state with future token embedding
        self.combine_proj = nn.Linear(2 * config.n_embd, config.n_embd, bias=config.bias)

        # Normalization
        self.norm1 = RMSNorm(config.n_embd)
        self.norm2 = RMSNorm(config.n_embd)

        # Transformer components (mini-transformer for each head)
        self.attn = MultiheadLatentAttention(config)
        self.mlp = MixtureOfExperts(config)
        self.attn_norm = RMSNorm(config.n_embd)
        self.mlp_norm = RMSNorm(config.n_embd)

Lines 1-24: Prediction Head Structure. Each MultiTokenPredictionHead is specialized for a particular depth — head 1 predicts 1 token ahead, head 2 predicts 2 tokens ahead, etc. We store the depth for potential depth-conditional processing (though we don’t use it in this simple implementation).

The architecture has 3 main components: a combination projection that merges the hidden state and future token embeddings, normalization layers for stabilization, and a mini-Transformer consisting of an attention module and an MoE. This mini-Transformer is complete but lightweight — it has the same architecture as our main model blocks but serves a specialized purpose.

    def forward(self, prev_hidden, future_token_embed):
        """
        Args:
            prev_hidden: [B, T, D] - Hidden states from previous layer
            future_token_embed: [B, T, D] - Embeddings of future tokens

        Returns:
            hidden: [B, T, D] - Processed hidden states
        """
        # Normalize inputs
        prev_norm = self.norm1(prev_hidden)
        future_norm = self.norm2(future_token_embed)

        # Combine representations
        combined = torch.cat([prev_norm, future_norm], dim=-1)
        hidden = self.combine_proj(combined)

        # Process through mini-transformer
        hidden = hidden + self.attn(self.attn_norm(hidden))
        moe_out, _ = self.mlp(self.mlp_norm(hidden))
        hidden = hidden + moe_out

        return hidden

Lines 26-41: The Combination Strategy. The forward method takes two inputs: prev_hidden (the hidden representation at position , encoding all context up to that point) and future_token_embed (the embedding of the token at position , providing information about what’s been generated). We normalize both inputs independently — this prevents scale mismatches between the hidden representations (which may have grown or shrunk through many Transformer layers) and the embeddings (which come fresh from the embedding layer). We concatenate along the feature dimension, doubling the dimensionality, then project back to n_embd dimensions. This projection learns how to merge content from these two different sources.

Lines 44-46: Mini-Transformer Processing. The combined representation flows through a lightweight Transformer. First, attention with a residual connection: the model can attend across the sequence, allowing position to gather information from other positions when predicting . This is crucial because the prediction might depend on context earlier in the sequence. Then, MoE with a residual connection: the expert networks can apply non-linear transformations, refining the combined representation. The use of the same MLA attention and MoE that we’ve already implemented is elegant — we’re reusing well-tested components. The pre-norm architecture (normalizing before attention and MoE rather than after) has become standard in modern Transformers for training stability.

Line 48: Returning Refined Hidden State. The output hidden state has the same dimensionality as the input (), so it can be projected through the vocabulary matrix to get logits for predicting . This hidden state has been enriched with information from both the context (via prev_hidden) and the intermediate token (via future_token_embed), and has been refined through attention and expert processing. It represents the model’s best understanding of what should come next-next, not just next.

Integrating Multi-Token Prediction with DeepSeek-V3’s Core Transformer

The MTP heads integrate into the main model during training. After computing the final hidden states from the main Transformer, we apply the following operations:

• Main prediction: Project to vocabulary to predict , compute cross-entropy loss
• Depth-1 prediction: For each position , get embedding of (ground truth), combine with through head 1, project to vocabulary to predict , compute cross-entropy loss
• Depth-2 prediction: For each position , get embedding of (ground truth), combine with head-1 output, project to vocabulary to predict , compute cross-entropy loss

The key insight is that we chain the heads: head 2’s input includes head 1’s output. This creates a hierarchical structure in which each head builds on the previous one, progressively looking further into the future.

Theoretical Foundations: MTP, Curriculum Learning, and Auxiliary Tasks

MTP has interesting theoretical connections to other areas of machine learning:

Temporal Difference Learning: In reinforcement learning, temporal difference learning propagates value information backward from future states. MTP does something analogous — it propagates gradient information backward from future predictions, encouraging current representations to encode future-relevant information.

Auxiliary Tasks: MTP can be viewed as an auxiliary task framework in which the auxiliary tasks are future token predictions. Research in multi-task learning shows that auxiliary tasks improve representation quality when they are related but distinct from the main task. Future token prediction is perfectly related (it is the same task at different time steps) but distinct (it requires different information).

Curriculum Learning: The depth-weighted loss structure implements a form of curriculum — we emphasize near-future predictions (easier, more reliable) more than far-future predictions (harder, noisier). This gradually increasing difficulty may help training by first learning short-term dependencies before tackling long-term structure.

Multi-Token Prediction Benefits: Coherence, Planning, and Faster Convergence

Research on Multi-Token Prediction shows several empirical benefits:

• Improved Coherence: Models trained with MTP generate more globally coherent text, with fewer contradictions or topic drift over long generations
• Better Planning: For tasks like story writing or code generation, where early decisions constrain later possibilities, MTP helps the model make forward-compatible choices
• Faster Convergence: The additional training signals can accelerate learning, reaching target performance with fewer training steps
• Regularization: MTP acts as a regularizer, preventing overfitting by encouraging representations that support multiple related objectives

However, MTP also has costs. Training becomes more complex — we must manage multiple prediction heads and carefully weight their losses. Training is slower — computing multiple predictions per position increases computation by a factor of roughly for future tokens (the factor is not linear because not all positions can predict tokens ahead). Memory usage increases due to the additional heads’ parameters.

The tradeoff is typically favorable for larger models and longer-form generation tasks. For small models or short-sequence tasks, the overhead may outweigh the benefits. In our children’s story generation task, MTP should help with maintaining narrative consistency across a story.

What's next? We recommend PyImageSearch University.

Course information:
86+ total classes • 115+ hours hours of on-demand code walkthrough videos • Last updated: April 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:

• ✓ 86+ courses on essential computer vision, deep learning, and OpenCV topics
• ✓ 86 Certificates of Completion
• ✓ 115+ hours 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 the first three lessons of this series, we progressively assembled the foundations of DeepSeek-V3: starting with its configuration and Rotary Positional Embeddings (RoPE), then advancing to the efficiency of Multi-Head Latent Attention (MLA), and scaling capacity through the Mixture of Experts (MoE). Each of these innovations has added a crucial piece to the architecture, balancing efficiency, scalability, and representational power. With those components in place, we turn to another breakthrough that redefines how language models learn and generate text: Multi-Token Prediction (MTP).

Traditional autoregressive models rely on next-token prediction, a strategy that, while effective, can be shortsighted — focusing only on immediate context rather than broader sequence-level patterns. MTP addresses this limitation by enabling the model to predict multiple tokens ahead, accelerating training and inference while enriching contextual understanding. In this lesson, we explore the shortcomings of next-token prediction, introduce the architecture of specialized prediction heads, and examine why MTP works from a gradient perspective.

We then dive into practical considerations (e.g., weighted loss, decay strategies, and implementation details), before integrating MTP into the main model. By the end, we see how this innovation not only improves efficiency but also strengthens the theoretical and empirical foundations of DeepSeek-V3, bringing us closer to assembling the complete architecture.

Citation Information

Mangla, P. “Autoregressive Model Limits and Multi-Token Prediction in DeepSeek-V3,” PyImageSearch, S. Huot, A. Sharma, and P. Thakur, eds., 2026, https://pyimg.co/alrep

@incollection{Mangla_2026_autoregressive-model-limits-and-mTP-in-deepseek-v3,
  author = {Puneet Mangla},
  title = {{Autoregressive Model Limits and Multi-Token Prediction in DeepSeek-V3}},
  booktitle = {PyImageSearch},
  editor = {Susan Huot and Aditya Sharma and Piyush Thakur},
  year = {2026},
  url = {https://pyimg.co/alrep},
}

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 Autoregressive Model Limits and Multi-Token Prediction in DeepSeek-V3 appeared first on PyImageSearch.

DeepSeek-V3 from Scratch: Mixture of Experts (MoE)

In the first two parts of this series, we established the foundations of DeepSeek-V3 by implementing its core configuration and positional encoding, followed by a deep dive into Multi-Head Latent Attention (MLA). Together, these components set the stage for a model that is both efficient and capable of handling long-range dependencies. With those building blocks in place, we now explore another key innovation in DeepSeek-V3: the Mixture of Experts (MoE).

MoE introduces a dynamic way of scaling model capacity without proportionally increasing computational cost. Instead of activating every parameter for every input, the model selectively routes tokens through specialized “expert” networks, allowing it to expand representational power while keeping inference efficient. In this lesson, we’ll unpack the theory behind MoE, explain how expert routing works, and then implement it step by step. This installment continues our broader goal of reconstructing DeepSeek-V3 from scratch — showing how each innovation, from RoPE to MLA to MoE, fits together into a cohesive architecture that balances scale, efficiency, and performance.

This lesson is the 3rd in a 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) (this tutorial)
4. Lesson 4
5. Lesson 5
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 Scaling Challenge in Neural Networks

As we scale neural networks, we face a fundamental tradeoff: larger models have greater capacity to learn complex patterns, but they’re more expensive to train and deploy. A standard Transformer feedforward layer applies the same computation to every token:

,

where and are weight matrices, typically with . For our model with , this means , giving us approximately 256K parameters per FFN (FeedForward Network) per layer.

To increase model capacity, we could simply make larger — say, instead of . This doubles the FFN parameters and theoretically doubles capacity. But it also doubles the computation for every token, even if most don’t need that extra capacity.

Mixture of Experts (Figure 1) offers a more efficient scaling paradigm: instead of a single large FFN, we create multiple smaller expert FFNs and route each token to a subset of these experts. This gives us the capacity of a much larger model while maintaining computational efficiency.

Mixture of Experts (MoE): Mathematical Foundation and Routing Mechanism

Consider expert networks, each with the same architecture as a standard FFN:

for . Instead of using all experts for every token, we select the top-k experts. The selection is determined by a learned routing function:

where is the router weight matrix and is a learnable bias vector. This gives us a probability distribution over experts for each token.

Top-k Routing: We select the top-k experts based on router probabilities:

The final output combines the selected experts, weighted by their normalized routing probabilities:

The renormalization ensures the selected experts’ weights sum to 1.

Capacity and Computation: With experts and (our configuration), each token activates 2 out of 4 experts. If each expert has the same size as a standard FFN, we have the parameters but only the computation per token. This is the MoE efficiency advantage: parameter count scales with , but computation scales with .

SwiGLU Activation in DeepSeek-V3: Improving MoE Non-Linearity

DeepSeek uses SwiGLU (Swish-Gated Linear Unit) instead of the traditional GELU (Gaussian Error Linear Units) activation. SwiGLU is a gated activation function that has shown superior performance in language models:

where:

• : projects input to hidden dimension
• : is another projection to hidden dimension
• : is the Swish activation (smooth version of ReLU)
• : denotes element-wise multiplication
• The result is then projected back:

The gating mechanism allows the network to control information flow more precisely than simple activation functions. The activation provides smooth gradients everywhere, improving training dynamics compared to ReLU’s hard threshold.

Shared Expert in DeepSeek-V3: Universal Processing in MoE Layers

DeepSeek introduces a shared expert that processes all tokens in addition to the routed experts. This design addresses a key limitation of pure MoE: some computations are beneficial for all tokens regardless of their content.

The shared expert has a larger hidden dimension (768 in our configuration vs 512 for individual experts) and processes every token. This ensures that:

• Common patterns are efficiently handled by dedicated capacity
• Specialized experts can focus on token-specific features
• Training is more stable with guaranteed gradient flow

The shared expert serves as a “base” computation that’s always present, while routed experts add specialized processing on top of it.

Auxiliary-Loss-Free Load Balancing in DeepSeek-V3 MoE

A critical challenge in MoE is load balancing. If the router learns to always send tokens to the same one or two experts, we lose the benefits of having multiple experts — the unused experts contribute nothing, and the overused ones become bottlenecks.

Traditional MoE models use an auxiliary loss that penalizes uneven expert usage:

where is the number of tokens routed to expert , is batch size, and is a coefficient. However, auxiliary losses add complexity and require careful tuning.

DeepSeek’s Innovation: Auxiliary-loss-free load balancing through dynamic bias updates. Instead of penalizing imbalance during training, we adjust the router biases to encourage balanced usage:

During training, we monitor how many tokens are routed to each expert. This gives us an expert_usage vector, where each entry counts the number of tokens assigned to a particular expert. We then compute the average usage across all experts.

To maintain a balanced load, we adjust the router biases: if an expert is used more than the average, its bias is decreased to make it less likely to be chosen in the future; if it is used less than the average, its bias is increased to make it more likely to be selected. This dynamic bias update encourages fair distribution of tokens across experts without requiring an explicit auxiliary loss.

Let denote the usage (number of tokens) of expert , and let

be the average usage across all experts. The router bias for expert , denoted , is updated as:

,

where is the learning rate controlling the magnitude of the bias adjustment.

This approach:

• Eliminates the need for auxiliary loss hyperparameter tuning
• Provides smoother load balancing over time
• Doesn’t interfere with the primary task loss
• Automatically adapts to data distribution changes

The bias updates are performed with a small learning rate (0.001 in our implementation) to ensure gradual adjustment without disrupting training.

Sequence-Wise Load Balancing for Mixture of Experts Models

For even better load balancing, DeepSeek can use a complementary sequence-wise auxiliary loss. This encourages different sequences in a batch to use different experts:

,

where is the expert usage vector for sequence (i.e., which experts were used), and measures similarity. By minimizing this loss, we encourage sequences to be complementary — if sequence A uses experts 1 and 2 heavily, sequence B should use experts 3 and 4.

Expert Specialization in MoE: Emergent Behavior in DeepSeek-V3

A fascinating property of MoE is expert specialization. Even though we don’t explicitly tell experts what to specialize in, they often learn to handle different types of patterns. In language models, researchers have observed:

• Syntactic experts: Handle grammatical structures, verb conjugations
• Semantic experts: Process meaning, synonyms, and conceptual relationships
• Domain experts: Specialize in specific topics (e.g., scientific text, dialogue)
• Numerical experts: Handle arithmetic, dates, quantities

This specialization emerges naturally as the routing function learns which experts are most effective for different inputs. Gradient flow during training reinforces this — when an expert performs well on certain patterns, the router learns to send similar patterns to that expert.

Mathematically, we can think of each expert as learning a local model that’s particularly good in some region of the input space. The router function implicitly partitions the input space, assigning different regions to different experts. This is similar to a mixture of experts in classical machine learning, but learned end-to-end through backpropagation.

Implementation: Building the DeepSeek-V3 MoE Layer from Scratch

Let’s implement the complete MoE layer with expert networks, routing, and load balancing:

class SwiGLU(nn.Module):
    """SwiGLU activation function used in DeepSeek experts"""
   
    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int, bias: bool = True):
        super().__init__()
        self.gate_proj = nn.Linear(input_dim, hidden_dim, bias=bias)
        self.up_proj = nn.Linear(input_dim, hidden_dim, bias=bias)
        self.down_proj = nn.Linear(hidden_dim, output_dim, bias=bias)
       
    def forward(self, x: torch.Tensor):
        gate = F.silu(self.gate_proj(x))  # SiLU activation
        up = self.up_proj(x)
        return self.down_proj(gate * up)

Lines 1-13: SwiGLU Activation: The SwiGLU class implements a gated activation mechanism. We have 3 linear projections:

• gate_proj: for the gating signal
• up_proj: for the value branch
• down_proj: for the output projection

The forward pass applies SiLU (Sigmoid Linear Unit) to the gate projection, multiplies it element-wise with the up-projection, and projects back down. This creates a more expressive activation than simple GELU, with the gating mechanism allowing fine-grained control over information flow.

class MoEExpert(nn.Module):
    """Expert network for Mixture of Experts using SwiGLU"""

    def __init__(self, config: DeepSeekConfig):
        super().__init__()
        self.expert_mlp = SwiGLU(
            config.n_embd,
            config.expert_intermediate_size,
            config.n_embd,
            config.bias
        )

    def forward(self, x: torch.Tensor):
        return self.expert_mlp(x)

Lines 14-27: Expert with SwiGLU: Each MoEExpert is now a SwiGLU network instead of a simple FFN. The intermediate size (expert_intermediate_size) controls capacity — we use 512 in our configuration, which is smaller than the shared expert’s 768. This asymmetry reflects the fact that routed experts handle specialized patterns, while the shared expert handles common operations.

class MixtureOfExperts(nn.Module):
    """
    DeepSeek MoE layer with shared expert and auxiliary-loss-free load balancing
   
    Key features:
    - Shared expert that processes all tokens
    - Auxiliary-loss-free load balancing via bias updates
    - Top-k routing to selected experts
    """

    def __init__(self, config: DeepSeekConfig):
        super().__init__()
        self.config = config
        self.n_experts = config.n_experts
        self.top_k = config.n_experts_per_token
        self.n_embd = config.n_embd

        # Router: learns which experts to use for each token
        self.router = nn.Linear(config.n_embd, config.n_experts, bias=False)

        # Expert networks
        self.experts = nn.ModuleList([
            MoEExpert(config) for _ in range(config.n_experts)
        ])

        # Shared expert (processes all tokens)
        if config.use_shared_expert:
            self.shared_expert = SwiGLU(
                config.n_embd,
                config.shared_expert_intermediate_size,
                config.n_embd,
                config.bias
            )
        else:
            self.shared_expert = None

        # Auxiliary-loss-free load balancing
        self.register_buffer('expert_bias', torch.zeros(config.n_experts))
        self.bias_update_rate = 0.001

        self.dropout = nn.Dropout(config.dropout)

Lines 28-68: MoE Layer Structure: The MixtureOfExperts class orchestrates routing and expert execution. The 3 key additions:

• shared_expert: full-capacity expert that processes all tokens
• expert_bias: buffer for auxiliary-loss-free balancing
• bias_update_rate: controls how quickly biases adapt

The dropout provides regularization across the entire MoE output.

    def forward(self, x: torch.Tensor):
        batch_size, seq_len, hidden_dim = x.shape
        x_flat = x.view(-1, hidden_dim)

        # Routing phase with bias for load balancing
        router_logits = self.router(x_flat) + self.expert_bias

        # Top-k routing
        top_k_logits, top_k_indices = torch.topk(router_logits, self.top_k, dim=-1)
        routing_weights = torch.zeros_like(router_logits)
        routing_weights.scatter_(-1, top_k_indices, F.softmax(top_k_logits, dim=-1))

        # Expert computation
        output = torch.zeros_like(x_flat)
        expert_usage = torch.zeros(self.n_experts, device=x.device)

Lines 70-84: Routing with Learnable Bias. The forward pass begins by flattening the input for efficient processing. We compute router logits and add the expert bias — this is the key to auxiliary-loss-free balancing. Overused experts have negative bias (making them less likely to be selected), while underused experts have positive bias (encouraging them to be selected). We then perform top-k selection and softmax normalization across the selected experts.

        # Process through selected experts
        for expert_idx in range(self.n_experts):
            expert_mask = (top_k_indices == expert_idx).any(dim=-1)
            expert_usage[expert_idx] = expert_mask.sum().float()

            if expert_mask.any():
                expert_input = x_flat[expert_mask]
                expert_output = self.experts[expert_idx](expert_input)

                # Weight by routing probability
                weights = routing_weights[expert_mask, expert_idx].unsqueeze(-1)
                output[expert_mask] += expert_output * weights

        # Add shared expert output (processes all tokens)
        if self.shared_expert is not None:
            shared_output = self.shared_expert(x_flat)
            output += shared_output

        # Auxiliary-loss-free load balancing (update biases during training)
        if self.training:
            with torch.no_grad():
                avg_usage = expert_usage.mean()
                for i in range(self.n_experts):
                    if expert_usage[i] > avg_usage:
                        self.expert_bias[i] -= self.bias_update_rate
                    else:
                        self.expert_bias[i] += self.bias_update_rate

        output = self.dropout(output)
        return output.view(batch_size, seq_len, hidden_dim), router_logits.view(batch_size, seq_len, -1)

Lines 86-97: Expert Processing. We iterate over all experts, identifying which tokens route to each one via the expert_mask. For each expert with assigned tokens, we extract those tokens, process them through the expert network, weight them by routing probability, and accumulate them into the output. This selective execution is what makes MoE efficient — we don’t compute all experts for all tokens.

Lines 100-102: Shared Expert. The shared expert processes all tokens unconditionally and adds its output to the routed experts’ output. This ensures every token receives some baseline processing, improving training stability and providing capacity for universal patterns. The shared expert’s larger hidden dimension (768 vs 512) reflects its broader responsibility.

Lines 105-112: Auxiliary-Loss-Free Balancing. During training, we update expert biases based on usage. We compute average usage across experts, then adjust biases: overused experts receive negative adjustments (discouraging future selection), while underused experts receive positive adjustments (encouraging future selection). Using the torch.no_grad() context ensures these bias updates don’t interfere with gradient computation. The small update rate (0.001) provides smooth, stable balancing over time.

Lines 114-115: Output and Return. We apply dropout to the combined output (routed + shared experts) and reshape back to the original dimensions. We return both the output and router logits — the latter can be used for optional auxiliary loss computation.

    def _complementary_sequence_aux_loss(self, router_logits, seq_mask=None):
      """
      router_logits: [batch_size, seq_len, num_experts]
          Raw logits from the router before softmax.
      seq_mask: optional mask for padding tokens.
      """

      # Convert to probabilities
      probs = F.softmax(router_logits, dim=-1)  # [B, T, E]

      # Aggregate per-sequence expert usage
      if seq_mask is not None:
          probs = probs * seq_mask.unsqueeze(-1)  # mask padding
      seq_usage = probs.sum(dim=1)  # [B, E]

      # Normalize per sequence
      seq_usage = seq_usage / seq_usage.sum(dim=-1, keepdim=True)

      # Compute pairwise similarity between sequences
      sim_matrix = torch.matmul(seq_usage, seq_usage.transpose(0, 1))  # [B, B]

      # Encourage complementarity: minimize similarity off-diagonal
      batch_size = seq_usage.size(0)
      off_diag = sim_matrix - torch.eye(batch_size, device=sim_matrix.device)
      loss = off_diag.mean()

      return loss

Lines 117-143: Complementary Sequence-Wise Loss. This method implements an alternative load-balancing approach. It converts router logits to probabilities, aggregates expert usage for each sequence, and computes pairwise similarity between sequences’ expert usage patterns. By minimizing off-diagonal similarity, we encourage different sequences to use different experts, promoting diversity in expert utilization. This can be added to the training loss with a small weight (e.g., 0.01).

MoE Design Decisions in DeepSeek-V3: SwiGLU, Shared Experts, and Routing

Several implementation choices merit discussion:

SwiGLU vs GELU: We use SwiGLU instead of traditional GELU because empirical research shows it consistently outperforms GELU in language models. The gating mechanism provides more expressive power, and SiLU’s smoothness improves gradient flow. The computational cost is slightly higher (three projections instead of two), but the quality improvement justifies it.

Shared Expert Design: The shared expert is a DeepSeek innovation that addresses a key limitation of pure MoE: some computations benefit all tokens. By providing dedicated capacity for universal processing, we free routed experts to specialize more aggressively. The larger hidden dimension (768 vs 512) for the shared expert reflects empirical findings that shared capacity requires more parameters than individual experts.

Auxiliary-Loss-Free Balancing: Traditional MoE uses auxiliary losses, such as:

where is the fraction of tokens routed to expert and is the average routing probability. This requires tuning (typically 0.01-0.1). Our bias-based approach eliminates the need for this hyperparameter, simplifying training. The tradeoff is that bias updates are less direct than gradient-based learning, but in practice, the smoother adaptation works well.

Complementary Sequence-Wise Loss: This alternative balancing approach is useful when batch diversity is high. By encouraging different sequences to use different experts, we naturally achieve balance. However, if the batch contains very similar sequences (e.g., all from the same domain), this loss may not be effective. It’s best used in combination with bias-based balancing or as an optional auxiliary objective.

Expert Capacity: Production MoE systems often implement expert capacity constraints — if too many tokens route to one expert, excess tokens are dropped or routed to a second choice. We don’t implement this in our educational model, but the formula would be:

where factor is typically 1.25-1.5. Tokens beyond this capacity are handled via overflow strategies.

MoE Computational and Memory Analysis in DeepSeek-V3

Let’s analyze the computational cost. For a standard FFN with hidden dimension :

For our MoE with routed experts (each with ), selected, and shared expert ():

The SwiGLU computation involves three projections:

For our configuration:

• Routing: (negligible)
• Routed experts:
• Shared expert:
• Total: 2.75M FLOPs per token

Compare to a standard FFN with : FLOPs. Our MoE uses 2.6× more computation but has much higher capacity (4 experts × 512 + 1 shared × 768 = 2,816 vs 1,024). We get 2.7× capacity for 2.6× computation — roughly linear scaling, which is the goal.

Memory usage during the forward pass stores activations for active experts only. During backpropagation, we need gradients for all experts (since routing is differentiable), yet the memory remains manageable. The bias vector is tiny (4 floats for 4 experts).

MoE Expert Specialization in Practice: Real-World Behavior

While we can’t demonstrate this in our small toy model, in larger-scale MoE models, expert specialization is observable through analysis of routing patterns. Researchers have visualized which experts activate for different types of inputs, revealing clear specialization. For example:

• Multilingual models: Different experts handle different languages
• Code models: Some experts handle syntax, others semantics, others API patterns
• Reasoning models: Numerical experts for math, logical experts for inference, retrieval experts for factual recall

This specialization isn’t programmed — it emerges from optimization. The routing function learns to partition the input space, and experts learn to excel in their assigned partitions. It’s a beautiful example of how end-to-end learning can discover structured solutions.

Training Dynamics of MoE: Load Balancing and Expert Utilization

In practice, MoE training exhibits interesting dynamics:

Early Training: Routing is initially random or near-uniform. All experts receive a similar load. The shared expert learns basic patterns that benefit all tokens.

Mid Training: Routing starts specializing. Some experts become preferred for certain patterns. Load imbalance can emerge without careful management. Bias-based balancing begins correcting the imbalance.

Late Training: Experts are clearly specialized. Routing is confident (high softmax probabilities for selected experts). Load is balanced through continuous bias adjustment. The shared expert handles universal operations while routed experts focus on specialized patterns.

Monitoring expert usage during training is valuable. We can log:

• Per-expert selection frequency
• Routing entropy (higher means more uniform)
• Expert bias magnitudes (large values indicate strong correction needed)

Mixture of Experts vs Related Techniques: Switch Transformers and Sparse Models

MoE shares ideas with several other architectural patterns:

Switch Transformers: Use top-1 routing (only one expert per token) for maximum efficiency. Simpler but less expressive than top-k.

Expert Choice: Instead of tokens choosing experts, experts choose tokens. Helps with load balancing but changes the computational pattern.

Sparse Attention: Like MoE, selectively activates parts of the network. Can be combined with MoE for extreme efficiency.

Dynamic Networks: Adapt network structure based on input. MoE is a specific form of dynamic computation.

With our MoE implementation complete, we’ve added efficient scaling to our model — the capacity grows superlinearly with computation cost. Combined with MLA’s memory efficiency and the upcoming MTP’s improved training signal, we’re building a model that’s efficient in training, efficient in inference, and capable of strong performance. Next, we’ll tackle Multi-Token Prediction, which improves the training signal itself by having the model look further ahead.

What's next? We recommend PyImageSearch University.

Course information:
86+ total classes • 115+ hours hours of on-demand code walkthrough videos • Last updated: April 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:

• ✓ 86+ courses on essential computer vision, deep learning, and OpenCV topics
• ✓ 86 Certificates of Completion
• ✓ 115+ hours 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 the third installment of our DeepSeek-V3 from Scratch series, we turn our attention to the Mixture of Experts (MoE) framework, a powerful approach to scaling neural networks efficiently. We begin by unpacking the scaling challenge in modern architectures and how MoE addresses it through selective expert activation. From its mathematical foundation to the introduction of SwiGLU activation, we explore how enhanced non-linearity and universal shared experts contribute to more flexible and expressive models.

We then examine the mechanics of load balancing, highlighting innovations (e.g., auxiliary-loss-free balancing and complementary sequence-wise strategies). These techniques ensure that experts are used effectively without introducing unnecessary complexity. We also explore how expert specialization emerges naturally during training, leading to diverse behaviors across experts that improve overall performance. This emergent specialization is not just theoretical — it becomes visible in practice, shaping how the model processes different types of input.

Finally, we walk through the implementation of MoE, discussing design decisions, computational trade-offs, and memory analysis. We connect these insights to related techniques, showing how MoE integrates into the broader landscape of efficient deep learning. By the end, we not only understand the theory but also gain practical knowledge of how to implement and optimize MoE within DeepSeek-V3. This part of the series equips us with the tools to harness expert specialization while keeping training dynamics balanced and efficient.

Citation Information

Mangla, P. “DeepSeek-V3 from Scratch: Mixture of Experts (MoE),” PyImageSearch, S. Huot, A. Sharma, and P. Thakur, eds., 2026, https://pyimg.co/a1w0g

@incollection{Mangla_2026_deepseek-v3-from-scratch-moe,
  author = {Puneet Mangla},
  title = {{DeepSeek-V3 from Scratch: Mixture of Experts (MoE)}},
  booktitle = {PyImageSearch},
  editor = {Susan Huot and Aditya Sharma and Piyush Thakur},
  year = {2026},
  url = {https://pyimg.co/a1w0g},
}

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: Mixture of Experts (MoE) appeared first on PyImageSearch.

Build DeepSeek-V3: Multi-Head Latent Attention (MLA) Architecture

In the first part of this series, we laid the foundation by exploring the theoretical underpinnings of DeepSeek-V3 and implementing key configuration elements such as Rotary Positional Embeddings (RoPE). That tutorial established how DeepSeek-V3 manages long-range dependencies and sets up its architecture for efficient scaling. By grounding theory in working code, we ensured that readers not only understood the concepts but also saw how they translate into practical implementation.

With that groundwork in place, we now turn to one of DeepSeek-V3’s most distinctive innovations: Multi-Head Latent Attention (MLA). While traditional attention mechanisms have proven remarkably effective, they often come with steep computational and memory costs. MLA reimagines this core operation by introducing a latent representation space that dramatically reduces overhead while preserving the model’s ability to capture rich contextual relationships.

In this lesson, we’ll break down the theory behind MLA, explore why it matters, and then implement it step by step. This installment continues our hands-on approach — moving beyond abstract concepts to practical code — while advancing the broader goal of the series: to reconstruct DeepSeek-V3 from scratch, piece by piece, until we assemble and train the full architecture.

This lesson is the 2nd of 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 (this tutorial)
3. Lesson 3
4. Lesson 4
5. Lesson 5
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 KV Cache Memory Problem in DeepSeek-V3

To understand why MLA is revolutionary, we must first understand the memory bottleneck in Transformer inference. Standard multi-head attention computes:

,

where are query, key, and value matrices for sequence length . In autoregressive generation (producing one token at a time), we cannot recompute attention over all previous tokens from scratch at each step — that would be computation per token generated.

Instead, we cache the key and value matrices. When generating token , we only compute (the query for the new token), then compute attention using and the cached . This reduces computation from to per generated token — a dramatic speedup.

However, this cache comes at a steep memory cost. For a model with layers, attention heads, and head dimension , the KV cache requires:

.

For a model like GPT-3 with 96 layers, 96 heads, 128-head dimensions, and 2048 sequence length, this is:

.

This means you can only serve a handful of users concurrently on even high-end GPUs. The memory bottleneck is often the limiting factor in deployment, not computation.

Multi-Head Latent Attention (MLA): KV Cache Compression with Low-Rank Projections

MLA (Figure 1) solves this through a compress-decompress strategy inspired by Low-Rank Adaptation (LoRA). The key insight: we do not need to store full -dimensional representations. We can compress them into a lower-dimensional latent space for storage, then decompress when needed for computation.

Step 1. Key-Value Compression: Instead of storing directly, we project them through a low-rank bottleneck:

,

where is the input, is the down-projection, and is the low-rank dimension. We only cache rather than the full and .

Step 2. Key-Value Decompression: When we need the actual key and value matrices for attention computation, we decompress:

,

where are up-projection matrices. This decomposition approximates the full key and value matrices through a low-rank factorization: and .

Memory Savings: Instead of caching , we cache . The reduction factor is . For our configuration with and , this is a 4× reduction. For larger models with and , it’s a 16× reduction — transformative for deployment.

Query Compression and Rotary Positional Embeddings (RoPE) Integration

MLA extends compression to queries, though less aggressively since queries are not cached:

,

where can be different from . In our configuration, versus — we give queries slightly more capacity.

Now comes the clever part: integrating RoPE. We split both queries and keys into content and positional components:

,

where denotes concatenation. The content components come from the compression-decompression process described above. The positional components are separate projections that we apply RoPE to:

,

where denotes applying rotary embedding at position . This separation is crucial: content and position are independently represented and combined only in the attention scores.

Attention Computation with Multi-Head Latent Attention (MLA)

The complete attention computation becomes:

.

Then standard multi-head attention:

,

where are per-head projections. The attention scores naturally incorporate both content similarity (through ) and positional information (through ).

Causal Masking: For autoregressive language modeling, we must prevent tokens from attending to future positions. We apply a causal mask:

.

This ensures position can only attend to positions , maintaining the autoregressive property.

Attention Weights and Output: After computing scores with the causal mask applied:

,

where is the effective key dimension (content plus RoPE dimensions). We apply attention to values:

,

where is the output projection. Finally, dropout is applied for regularization, and the result is added to the residual connection.

Implementation: Multi-Head Latent Attention (MLA)

Here is the complete implementation of MLA:

class MultiheadLatentAttention(nn.Module):
    """
    Multihead Latent Attention (MLA) - DeepSeek's efficient attention mechanism

    Key innovations:
    - Compression/decompression of queries and key-values
    - LoRA-style low-rank projections for efficiency
    - RoPE with separate content and positional components
    """

    def __init__(self, config: DeepSeekConfig):
        super().__init__()
        self.config = config
        self.n_embd = config.n_embd
        self.n_head = config.n_head
        self.head_dim = config.n_embd // config.n_head

        # Compression dimensions
        self.kv_lora_rank = config.kv_lora_rank
        self.q_lora_rank = config.q_lora_rank
        self.rope_dim = config.rope_dim

Lines 11-21: Configuration and Dimensions. We extract key parameters from the configuration object, computing the head dimension as . We store compression ranks (kv_lora_rank and q_lora_rank) and the RoPE dimension. These define the memory-accuracy tradeoff — lower ranks mean more compression but potentially lower quality. Our choices balance efficiency with model capacity.

        # KV decompression
        self.k_decompress = nn.Linear(self.kv_lora_rank, self.n_head * self.head_dim, bias=False)
        self.v_decompress = nn.Linear(self.kv_lora_rank, self.n_head * self.head_dim, bias=False)

        # Query compression
        self.q_proj = nn.Linear(self.n_embd, self.q_lora_rank, bias=False)
        self.q_decompress = nn.Linear(self.q_lora_rank, self.n_head * self.head_dim, bias=False)

        # RoPE projections
        self.k_rope_proj = nn.Linear(self.n_embd, self.n_head * self.rope_dim, bias=False)
        self.q_rope_proj = nn.Linear(self.q_lora_rank, self.n_head * self.rope_dim, bias=False)

        # Output projection
        self.o_proj = nn.Linear(self.n_head * self.head_dim, self.n_embd, bias=config.bias)

        # Dropout
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)

        # RoPE
        self.rope = RotaryEmbedding(self.rope_dim, config.block_size)

        # Causal mask
        self.register_buffer(
            "causal_mask",
            torch.tril(torch.ones(config.block_size, config.block_size)).view(
                1, 1, config.block_size, config.block_size
            )
        )

Lines 23-29: KV Compression Pipeline. The compression-decompression architecture follows the low-rank factorization principle. The kv_proj layer performs the down-projection from to , cutting the dimensionality in half. We apply RMSNorm to the compressed representation for stability — this normalization helps prevent the compressed representation from drifting to extreme values during training. The decompression layers k_decompress and v_decompress then expand back to dimensions. Note that we use bias=False for these projections — empirical research shows that biases in attention projections do not significantly help and add unnecessary parameters.

Lines 31-33: Query Processing and RoPE Projections. Query handling follows a similar compression pattern but with a slightly higher rank (). The asymmetry makes sense: we do not cache queries, so memory pressure is lower, and we can afford more capacity. The RoPE projections are separate pathways — k_rope_proj projects directly from the input , while q_rope_proj projects from the compressed query representation. Both target the RoPE dimension of 64. This separation of content and position is architecturally elegant: the model learns different transformations for “what” (content) versus “where” (position).

Lines 36-51: Infrastructure Components. The output projection o_proj combines multi-head outputs back to the model dimension. We include 2 dropout layers:

• attn_dropout: applied to attention weights (reducing overfitting on attention patterns)
• resid_dropout: applied to the final output (regularizing the residual connection)

The RoPE module is instantiated with our chosen dimension and maximum sequence length. Finally, we create and register a causal mask as a buffer — by using register_buffer, this tensor moves with the model to GPU/CPU and is included in the state dict, but is not treated as a learnable parameter.

    def forward(self, x: torch.Tensor, attention_mask: Optional[torch.Tensor] = None):
        B, T, C = x.size()

        # Compression phase
        kv_compressed = self.kv_norm(self.kv_proj(x))
        q_compressed = self.q_proj(x)

        # Decompression phase
        k_content = self.k_decompress(kv_compressed)
        v = self.v_decompress(kv_compressed)
        q_content = self.q_decompress(q_compressed)

        # RoPE components
        k_rope = self.k_rope_proj(x)
        q_rope = self.q_rope_proj(q_compressed)

        # Reshape [B, H, T, d_head] for multi-head attention
        k_content = k_content.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        q_content = q_content.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        k_rope = k_rope.view(B, T, self.n_head, self.rope_dim).transpose(1, 2)
        q_rope = q_rope.view(B, T, self.n_head, self.rope_dim).transpose(1, 2)

        # Apply RoPE
        cos, sin = self.rope(x, T)
        q_rope = apply_rope(q_rope, cos, sin)
        k_rope = apply_rope(k_rope, cos, sin)

        # Concatenate content and rope parts
        q = torch.cat([q_content, q_rope], dim=-1)
        k = torch.cat([k_content, k_rope], dim=-1)

Lines 52-57: Compression Phase. The forward pass begins by compressing the input. We project onto the KV latent space, apply normalization, and project back onto the query latent space. These operations are lightweight — just matrix multiplications. The compressed representations are what we would cache during inference. Notice that kv_compressed has shape versus the original — we’ve already halved the memory footprint.

Lines 60-73: Decompression and RoPE. We decompress to get content components and compute separate RoPE projections. Then comes a crucial reshaping step: we convert from to , moving the head dimension before the sequence dimension. This layout is required for multi-head attention — each head operates independently, and we want to batch those operations. The .transpose(1, 2) operation efficiently swaps dimensions without copying data.

Lines 76-82: RoPE Application and Concatenation. We fetch cosine and sine tensors from our RoPE module and apply the rotation to both queries and keys. Critically, we only rotate the RoPE components, not the content components. This maintains the separation between “what” and “where” information. We then concatenate along the feature dimension, creating final query and key tensors of shape . The attention scores will capture both content similarity and relative position.

        # Attention computation
        scale = 1.0 / math.sqrt(q.size(-1))
        scores = torch.matmul(q, k.transpose(-2, -1)) * scale

        # Apply causal mask
        scores = scores.masked_fill(self.causal_mask[:, :, :T, :T] == 0, float('-inf'))

        # Apply padding mask if provided
        if attention_mask is not None:
            padding_mask_additive = (1 - attention_mask).unsqueeze(1).unsqueeze(2) * float('-inf')
            scores = scores + padding_mask_additive

        # Softmax and dropout
        attn_weights = F.softmax(scores, dim=-1)
        attn_weights = self.attn_dropout(attn_weights)

        # Apply attention to values
        out = torch.matmul(attn_weights, v)

        # Reshape and project
        out = out.transpose(1, 2).contiguous().view(B, T, self.n_head * self.head_dim)
        out = self.resid_dropout(self.o_proj(out))

        return out

Lines 84-94: Attention Score Computation and Masking. We compute scaled dot-product attention: . The scaling factor is critical for training stability — without it, attention logits would grow large as dimensions increase, leading to vanishing gradients in the softmax. We apply the causal mask using masked_fill, setting future positions to negative infinity so they contribute zero probability after softmax. If an attention mask is provided (for handling padding), we convert it to an additive mask and add it to scores. This handles variable-length sequences in a batch.

Lines 97-107: Attention Weights and Output. We apply softmax to convert scores to probabilities, ensuring they sum to 1 over the sequence dimension. Dropout is applied to attention weights — this has been shown to help with generalization, perhaps by preventing the model from becoming overly dependent on specific attention patterns. We multiply attention weights by values to get our output. The final transpose and reshape convert from the multi-head layout back to , concatenating all heads. The output projection and residual dropout complete the attention module.

Multi-Head Latent Attention and KV Cache Optimization

Multi-Head Latent Attention (MLA) is one approach to KV cache optimization — compression through low-rank projections. Other approaches include the following:

• Multi-Query Attention (MQA), where all heads share a single key and value
• Grouped-Query Attention (GQA), where heads are grouped to share KV pairs
• KV Cache Quantization, which stores keys and values at lower precision (INT8 or INT4)
• Cache Eviction Strategies, which discard less important past tokens

Each approach has the following trade-offs:

• MQA and GQA reduce quality more than MLA but are simpler
• Quantization can degrade accuracy
• Cache eviction strategies discard historical context

DeepSeek-V3’s MLA offers an appealing middle ground — significant memory savings with minimal quality loss through a principled compression approach.

For readers interested in diving deeper into KV cache optimization, we recommend exploring the “KV Cache Optimization” series, which covers these techniques in detail, including implementation strategies, benchmarking results, and guidance on choosing the right approach for a given use case.

With MLA implemented, we have addressed one of the primary memory bottlenecks in Transformer inference — the KV cache. Our attention mechanism can now serve longer contexts and more concurrent users within the same hardware budget. In the next lesson, we will address another critical challenge: scaling model capacity efficiently through Mixture of Experts (MoE).

What's next? We recommend PyImageSearch University.

Course information:
86+ total classes • 115+ hours hours of on-demand code walkthrough videos • Last updated: April 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:

• ✓ 86+ courses on essential computer vision, deep learning, and OpenCV topics
• ✓ 86 Certificates of Completion
• ✓ 115+ hours 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 2nd lesson of our DeepSeek-V3 from Scratch series, we dive into the mechanics of Multi-Head Latent Attention (MLA) and why it is a crucial innovation for scaling large language models.

We begin by introducing MLA and framing it against the KV cache memory problem, a common bottleneck in Transformer architectures. By understanding this challenge, we set the stage for how MLA provides a more efficient solution through compression and smarter attention computation.

We then explore how low-rank projections enable MLA to compress key-value representations without losing essential information. This compression is paired with query compression and RoPE integration, ensuring that positional encoding remains geometrically consistent while reducing computational overhead.

Together, these techniques rethink the attention mechanism, balancing efficiency and accuracy and making MLA a powerful tool for modern architectures.

Finally, we walk through the implementation of MLA, showing how it connects directly to KV cache optimization.

By the end of this lesson, we not only understand the theory but also gain hands-on experience implementing MLA and integrating it into DeepSeek-V3. This practical approach shows how MLA reshapes attention computation, paving the way for more memory-efficient and scalable models.

Citation Information

Mangla, P. “Build DeepSeek-V3: Multi-Head Latent Attention (MLA) Architecture,” PyImageSearch, S. Huot, A. Sharma, and P. Thakur, eds., 2026, https://pyimg.co/scgjl

@incollection{Mangla_2026_build-deepseek-v3-mla-architecture,
  author = {Puneet Mangla},
  title = {{Build DeepSeek-V3: Multi-Head Latent Attention (MLA) Architecture}},
  booktitle = {PyImageSearch},
  editor = {Susan Huot and Aditya Sharma and Piyush Thakur},
  year = {2026},
  url = {https://pyimg.co/scgjl},
}

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 Build DeepSeek-V3: Multi-Head Latent Attention (MLA) Architecture appeared first on PyImageSearch.

Introduction to the DeepSeek-V3 Model

The landscape of large language models has been rapidly evolving, with innovations in architecture, training efficiency, and inference optimization pushing the boundaries of what is possible in natural language processing. The DeepSeek-V3 model represents a significant milestone in this evolution, introducing a suite of cutting-edge techniques that address some of the most pressing challenges in modern language model development:

• memory efficiency during inference
• computational cost during training
• effective capture of long-range dependencies

In this comprehensive lesson, we embark on an ambitious journey to build DeepSeek-V3 from scratch, implementing every component from first principles. This isn’t just another theoretical overview. We will write actual, working code that you can run, modify, and experiment with. By the end of this series, you will have a deep understanding of 4 revolutionary architectural innovations and how they synergistically combine to create a powerful language model.

This lesson is the 1st in a 6-part series on Building DeepSeek-V3 from Scratch:

1. DeepSeek-V3 Model: Theory, Config, and Rotary Positional Embeddings (this tutorial)
2. Lessons 2
3. Lesson 3
4. Lesson 4
5. Lesson 5
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 Four Pillars of DeepSeek-V3

Multihead Latent Attention (MLA): Traditional Transformer models face a critical bottleneck during inference: the key-value (KV) cache grows linearly with sequence length, consuming massive amounts of memory. For a model with 32 attention heads and a hidden dimension of 4096, storing keys and values for a single sequence of 2048 tokens requires over 1GB of memory. DeepSeek’s MLA addresses this by introducing a clever compression-decompression mechanism inspired by Low-Rank Adaptation (LoRA). Instead of storing full key and value matrices, MLA compresses them into a low-rank latent space, achieving up to a 75% reduction in KV cache memory while maintaining model quality. This isn’t just a theoretical improvement; it translates directly to the ability to serve more concurrent users or process longer contexts with the same hardware (Figure 1).

Mixture of Experts (MoE): The challenge in scaling language models is balancing capacity with computational cost. Simply making models wider and deeper becomes prohibitively expensive. MoE offers an elegant solution: instead of every token passing through the same feedforward network, we create multiple “expert” networks and route each token to only a subset of them. DeepSeek-V3 implements this with a learned routing mechanism that dynamically selects the most relevant experts for each token. With 4 experts and top-2 routing, we effectively quadruple the model’s capacity while only doubling the computation per token. The routing function learns to specialize different experts for different types of patterns — perhaps one expert becomes good at handling numerical reasoning, another at processing dialogue, and so on.

Multi-Token Prediction (MTP): Traditional language models predict one token at a time, receiving a training signal only for the immediate next token. This is somewhat myopic — humans don’t just think about the very next word; we plan ahead, considering how sentences and paragraphs will unfold. MTP addresses this by training the model to predict multiple future tokens simultaneously. If we are at position in the sequence, standard training predicts token . MTP adds auxiliary prediction heads that predict tokens , , and so on. This provides a richer training signal, encouraging the model to learn better long-range planning and coherence. It is particularly valuable for tasks requiring forward-looking reasoning.

Rotary Positional Embeddings (RoPE): Transformers don’t inherently understand position — they need explicit positional information. Early approaches used absolute position embeddings, but these struggle with sequences longer than those seen during training. RoPE takes a geometric approach: it rotates query and key vectors in a high-dimensional space, with the rotation angle proportional to the position. This naturally encodes relative position information and exhibits remarkable extrapolation properties. A model trained on 512-token sequences can often handle 2048-token sequences at inference time without degradation.

The combination of these 4 techniques is more than the sum of its parts. MLA reduces memory pressure, allowing us to handle longer contexts or larger batch sizes. MoE increases model capacity without proportional compute increases, making training more efficient. MTP provides richer gradients, accelerating learning and improving model quality. RoPE enables better position understanding and length generalization. Together, they create a model that is efficient to train, efficient to serve, and capable of producing high-quality outputs.

What You Will Build

By the end of this series, you will have implemented a working DeepSeek-V3 model trained on the TinyStories dataset — a curated collection of simple children’s stories. The dataset is ideal for demonstrating core language modeling concepts without requiring massive computational resources. Your model will be able to generate coherent, creative stories in the style of children’s literature. More importantly, you will understand every line of code, every architectural decision, and every mathematical principle behind the model.

The DeepSeek-V3 model we build uses carefully chosen hyperparameters for educational purposes:

• 6 Transformer layers
• 256-dimensional token embeddings
• 8 attention heads
• 4 MoE experts with top-2 routing
• 2-token-ahead prediction training objective (MTP)

These choices balance pedagogical clarity with practical performance: the model is small enough to train on a single GPU in a reasonable time, yet large enough to generate meaningful outputs and demonstrate the key architectural innovations.

Prerequisites and Setup for Building the DeepSeek-V3 Model

Before we dive in, ensure you have a working Python environment with PyTorch 2.0+, the transformers library, and standard scientific computing packages (e.g., numpy, datasets). A GPU is highly recommended but not required — you can train on a CPU, though it will be slower. The complete code is available as a Jupyter notebook, allowing you to experiment interactively.

# Install required packages
!pip install -q transformers datasets torch accelerate tensorboard

# Import core libraries
import os
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass
from typing import Optional, Tuple, List, Dict
import logging
import json

print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
print(f"Device: {torch.device('cuda' if torch.cuda.is_available() else 'cpu')}")

Implementing DeepSeek-V3 Model Configuration and RoPE

DeepSeek-V3 Model Parameters and Configuration

Before we can build any neural network, we need a systematic way to manage its hyperparameters — the architectural decisions that define the model. In modern deep learning, the configuration pattern has become essential: we encapsulate all hyperparameters in a single, serializable object that can be saved, loaded, and modified independently of the model code. This is not just good software engineering — it is crucial for reproducibility, experimentation, and deployment.

DeepSeek-V3’s configuration must capture parameters across multiple dimensions. First, there are the standard Transformer parameters:

• vocabulary size
• number of Transformer layers
• hidden dimension
• number of attention heads
• maximum context length

These follow from the canonical Transformer architecture, where the model transforms input sequences through layers of self-attention and feedforward processing.

Beyond these basics, we need parameters specific to the DeepSeek-V3 innovations. For MLA, we require the LoRA ranks for key-value compression () and query compression (), as well as the RoPE dimension (). For MoE, we specify the number of experts (), how many to activate per token (), and coefficients for auxiliary losses. For MTP, we define how many tokens ahead to predict ().

The mathematical relationship between these parameters determines the model’s computational and memory characteristics. The standard Transformer attention complexity scales as for sequence length . With MLA’s compression, we reduce the KV cache from to approximately , where . For our chosen parameters with and , this represents approximately a 50% reduction in KV cache size.

Rotary Positional Embeddings: Geometric Position Encoding

RoPE (Figure 2) represents one of the most elegant ideas in modern Transformer research. To understand it, we must first examine why position matters and where earlier approaches had limitations.

The Position Problem: Self-attention mechanisms are permutation-invariant — if we shuffle the input tokens, we get the same output (modulo the shuffling). But language is sequential; “The cat chased the mouse” means something very different from “The mouse chased the cat.” We need to inject positional information.

Absolute Positional Embeddings: The original Transformer used sinusoidal positional embeddings: and . These are added to input embeddings. Learned absolute positional embeddings are another option. But both struggle with extrapolation — a model trained on sequences up to length 512 often fails when applied to sequences of length 1024.

Relative Position Approaches: Some models (e.g., Transformer-XL) use relative positional encodings, explicitly modeling the distance between tokens. This helps with extrapolation but adds computational overhead.

RoPE’s Geometric Insight: RoPE takes a different approach, encoding position through rotation in complex space. Consider the attention score between query at position and key at position :

RoPE modifies this by rotating both and by angles proportional to their positions:

where is the rotation matrix corresponding to position . The key insight: rotation matrices satisfy , so the attention score naturally depends on the relative position rather than absolute positions.

In practice, we implement this in 2D rotation pairs. For a -dimensional vector, we split it into pairs and rotate each pair:

where follows the same frequency pattern as sinusoidal embeddings. This gives us multiple rotation frequencies, allowing the model to capture both fine-grained and coarse-grained positional relationships.

Why RoPE Extrapolates Well: The rotation formulation naturally extends to positions beyond training data. If the model learns that a relative position of +5 corresponds to a certain rotation angle, it can apply the same principle to positions beyond its training range. The continuous nature of trigonometric functions means there are no discrete position embeddings that “run out.”

RMSNorm: A Modern Normalization Choice: Before diving into code, we should mention RMSNorm (Root Mean Square Normalization), which DeepSeek uses instead of LayerNorm. While LayerNorm computes:

RMSNorm simplifies by removing the mean-centering and bias:

This is computationally cheaper and empirically performs just as well for language models. The key insight is that the mean-centering term in LayerNorm may not be necessary for Transformers, where the activations are already roughly centered.

Implementation: Configuration and Rotary Positional Embeddings

Now let’s implement these concepts. We’ll start with the configuration class:

import json


@dataclass
class DeepSeekConfig:
    """Configuration for DeepSeek model optimized for children's stories"""
    vocab_size: int = 50259  # GPT-2 vocabulary size + <|story|> + </|story|> tokens
    n_layer: int = 6         # Number of transformer blocks
    n_head: int = 8          # Number of attention heads
    n_embd: int = 256        # Embedding dimension
    block_size: int = 1024   # Maximum context window
    dropout: float = 0.1     # Dropout rate
    bias: bool = True        # Use bias in linear layers

    # MLA (Multihead Latent Attention) config
    kv_lora_rank: int = 128  # LoRA rank for key-value projection
    q_lora_rank: int = 192   # LoRA rank for query projection
    rope_dim: int = 64       # RoPE dimension

    # MoE (Mixture of Experts) config
    n_experts: int = 4       # Number of experts
    n_experts_per_token: int = 2  # Number of experts per token (top-k)
    expert_intermediate_size: int = 512  # Expert hidden size
    shared_expert_intermediate_size: int = 768  # Shared expert hidden size
    use_shared_expert: bool = True  # Enable shared expert
    aux_loss_weight: float = 0.0  # Auxiliary loss weight (0.0 for aux-free)

    # Multi-token prediction
    multi_token_predict: int = 2  # Predict next 2 tokens

Lines 1-5: Configuration Class Structure: We use Python’s @dataclass decorator to define our DeepSeekConfig class, which automatically generates initialization and representation methods. This is more than syntactic sugar — it ensures type hints are respected and provides built-in equality comparisons. The configuration serves as a single source of truth for model hyperparameters, making it easy to experiment with different architectures by simply modifying this object.

Lines 7-13: Standard Transformer Parameters: We define the core Transformer dimensions. The vocabulary size of 50,259 comes from the GPT-2 tokenizer, with two additional custom tokens for story boundaries. We choose 6 layers and a 256-dimensional embedding size as a balance between model capacity and computational cost — this is small enough to train on a single consumer GPU but large enough to demonstrate the key DeepSeek innovations. The block size of 1024 determines the model’s maximum context length, sufficient for coherent short stories. The dropout rate of 0.1 provides regularization without being overly aggressive.

Lines 16-18: MLA Configuration: These parameters control our Multihead Latent Attention mechanism. The kv_lora_rank of 128 means we compress key-value representations from 256 dimensions down to 128 — a 50% reduction that translates directly to KV cache memory savings. The q_lora_rank of 192 provides slightly more capacity for query compression since queries don’t need to be cached during inference. The rope_dim of 64 specifies how many dimensions use RoPE — we don’t apply RoPE to all dimensions, only to a subset, allowing some dimensions to focus purely on content rather than position.

Lines 21-29: MoE and MTP Configuration: We configure 4 expert networks with top-2 routing, meaning each token will be processed by exactly 2 out of 4 experts. This gives us 2× more parameters than a standard feedforward layer while maintaining the same computational cost. The aux_loss_weight of 0.01 determines how strongly we penalize uneven expert usage — this is crucial for preventing all tokens from routing to just one or two experts. The multi_token_predict parameter determines how many future tokens the model is trained to predict at each step.

    def __post_init__(self):
        """Initialize special tokens after dataclass initialization"""
        self.special_tokens = {
            "story_start": "<|story|>",
            "story_end": "</|story|>",
        }

    def to_dict(self):
        """Convert configuration to dictionary"""
        return {
            'vocab_size': self.vocab_size,
            'n_layer': self.n_layer,
            'n_head': self.n_head,
            'n_embd': self.n_embd,
            'block_size': self.block_size,
            'dropout': self.dropout,
            'bias': self.bias,
            'kv_lora_rank': self.kv_lora_rank,
            'q_lora_rank': self.q_lora_rank,
            'rope_dim': self.rope_dim,
            'n_experts': self.n_experts,
            'n_experts_per_token': self.n_experts_per_token,
            'expert_intermediate_size': self.expert_intermediate_size,
            'shared_expert_intermediate_size': self.shared_expert_intermediate_size,
            'use_shared_expert': self.use_shared_expert,
            'aux_loss_weight': self.aux_loss_weight,
            'multi_token_predict': self.multi_token_predict,
            'special_tokens': self.special_tokens,
        }

    def to_json_string(self, indent=2):
        """Convert configuration to JSON string"""
        return json.dumps(self.to_dict(), indent=indent)

    @classmethod
    def from_dict(cls, config_dict):
        """Create configuration from dictionary"""
        # Remove special_tokens from dict as it's set in __post_init__
        config_dict = {k: v for k, v in config_dict.items() if k != 'special_tokens'}
        return cls(**config_dict)

    @classmethod
    def from_json_string(cls, json_string):
        """Create configuration from JSON string"""
        return cls.from_dict(json.loads(json_string))

Lines 31-75: Special Methods for Serialization: We implement __post_init__ to add special tokens after initialization, ensuring they’re always present but not required in the constructor. The to_dict and to_json_string methods enable easy serialization for saving configurations alongside trained models. The class methods from_dict and from_json_string provide deserialization, creating a complete round-trip for configuration management. This pattern is essential for reproducibility — we can save a configuration with our trained model and later reconstruct the exact architecture.

Next, we implement the RoPE module.

class RMSNorm(nn.Module):
    """Root Mean Square Layer Normalization"""
    def __init__(self, ndim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(ndim))

    def forward(self, x):
        norm = x.norm(dim=-1, keepdim=True) * (x.size(-1) ** -0.5)
        return self.weight * x / (norm + self.eps)

RMSNorm Implementation (Lines 1-10): Our RMSNorm class is remarkably simple. In the constructor, we create a learnable weight parameter (the in our equations) initialized to ones. In the forward pass, we compute the L2 norm of the input along the feature dimension, multiply by to get the RMS, and then scale the input by the inverse of this norm (plus epsilon for numerical stability) and multiply by the learned weight parameter. This normalization ensures our activations have unit RMS, helping with training stability and gradient flow.

class RotaryEmbedding(nn.Module):
    """Rotary Positional Embedding (RoPE) for better position understanding"""
    def __init__(self, dim, max_seq_len=2048):
        super().__init__()
        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer('inv_freq', inv_freq)
        self.max_seq_len = max_seq_len

    def forward(self, x, seq_len=None):
        if seq_len is None:
            seq_len = x.shape[-2]

        t = torch.arange(seq_len, device=x.device).type_as(self.inv_freq)
        freqs = torch.outer(t, self.inv_freq)
        cos, sin = freqs.cos(), freqs.sin()
        return cos, sin

def apply_rope(x, cos, sin):
    """Apply rotary position embedding"""
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)

The RotaryEmbedding Class (Lines 12-27): The constructor creates the inverse frequency vector inv_freq following the same frequency schedule used in sinusoidal positional embeddings, where each pair of dimensions is assigned a frequency following the schedule . We use register_buffer rather than a parameter because these frequencies shouldn’t be learned — they’re fixed by our positional encoding design. In the forward pass, we create position indices from 0 to seq_len, compute the outer product with inverse frequencies (giving us a matrix where entry is , and compute the cosine and sine values. These will be broadcast and applied to query and key vectors. The resulting cosine and sine tensors broadcast across the batch, head, and sequence dimensions during attention computation.

The apply_rope Function (Lines 29-32): This elegant function applies the 2D rotation. We chunk the input into pairs of dimensions (effectively treating each pair of dimensions as the real and imaginary components of a complex number). We then apply the rotation formula:

The chunking operation splits along the last dimension. We compute each rotated component and then concatenate them back together. This vectorized implementation is far more efficient than iterating over dimension pairs in Python.

Design Choices and Tradeoffs: Several decisions merit discussion. We chose partial RoPE (rope_dim=64 rather than full n_embd=256) because empirical research shows that applying RoPE to all dimensions can sometimes hurt performance — some dimensions benefit from remaining content-focused rather than encoding position. Our LoRA ranks are fairly high (128 and 192) relative to the 256-dimensional embeddings; in larger models, the compression ratio would be more aggressive. The special tokens pattern (story_start and story_end) provides explicit boundaries that help the model learn story structure — it knows when a generation should terminate.

What's next? We recommend PyImageSearch University.

Course information:
86+ total classes • 115+ hours hours of on-demand code walkthrough videos • Last updated: April 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:

• ✓ 86+ courses on essential computer vision, deep learning, and OpenCV topics
• ✓ 86 Certificates of Completion
• ✓ 115+ hours 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 blog, we walk through the foundations of DeepSeek-V3, starting with its theoretical underpinnings and the four pillars that shape its architecture. We explore why these pillars matter, how they guide the design of the model, and what we aim to build by the end of the lesson. By laying out the prerequisites and setup, we ensure that we’re equipped with the right tools and mindset before diving into the implementation details.

Next, we focus on the model configuration, where we break down the essential parameters that define DeepSeek-V3’s behavior. We discuss how these configurations influence performance, scalability, and adaptability, and why they are critical for building a robust model. Alongside this, we introduce Rotary Positional Embeddings (RoPE), a geometric approach to positional encoding that enhances the model’s ability to capture sequential information with precision.

Finally, we bring theory into practice by implementing both the configuration and RoPE step by step. We highlight how these components integrate seamlessly, forming the backbone of DeepSeek-V3. By the end, we not only understand the theoretical aspects but also gain hands-on experience in building and customizing the model. Together, these steps demystify the process and set the stage for deeper experimentation with advanced Transformer architectures.

Citation Information

Mangla, P. “DeepSeek-V3 Model: Theory, Config, and Rotary Positional Embeddings,” PyImageSearch, S. Huot, A. Sharma, and P. Thakur, eds., 2026, https://pyimg.co/1atre

@incollection{Mangla_2026_deepseek-v3-model-theory-config-and-rotary-positional-embeddings,
  author = {Puneet Mangla},
  title = {{DeepSeek-V3 Model: Theory, Config, and Rotary Positional Embeddings}},
  booktitle = {PyImageSearch},
  editor = {Susan Huot and Aditya Sharma and Piyush Thakur},
  year = {2026},
  url = {https://pyimg.co/1atre},
}

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 Model: Theory, Config, and Rotary Positional Embeddings appeared first on PyImageSearch.

SAM 3 for Video: Concept-Aware Segmentation and Object Tracking

In Part 1 of this series, we introduced Segment Anything Model 3 (SAM 3) and saw how it moves beyond geometric prompts to concept-based visual understanding. We learned how the model can segment all instances of a concept using natural language and visual examples.

In Part 2, we went one step further. We explored multi-modal prompting and interactive workflows. We combined text, bounding boxes, and point clicks to build precise and controllable segmentation pipelines on images.

So far, however, everything we have done lives in a static world.

Images do not move. Objects do not disappear. There is no notion of time.

Video changes everything.

In videos, segmentation alone is not enough. We also need temporal consistency. If a person appears in frame 1 and walks across the scene, the model must not only segment that person — it must also remember that it is the same person in frame 200.

This is where SAM3 becomes fundamentally different from previous systems.

SAM3 does not treat video as a bag of independent images. Instead, it maintains a streaming memory and a tracking state that allows it to propagate object identities across frames. Detection, segmentation, and tracking are no longer separate steps. They are part of a single, unified pipeline.

In other words, SAM3 does not just answer: “Where is the object in this frame?”

It answers: “Where is this concept over time?”

In this tutorial, we will focus entirely on using SAM3 with video. We will build several practical pipelines that combine detection, segmentation, and tracking into one coherent workflow.

Specifically, we will implement 4 tasks:

• Video Detection, Segmentation, and Tracking Using a Text Prompt
  Here, we use a text prompt such as "person" and let SAM3 detect, segment, and track all instances of that concept throughout the video.
• Real-Time Detection, Segmentation, and Tracking Using a Text Prompt via Webcam
  This is the same idea, but running in real time on a live camera stream.
• Detection, Segmentation, and Tracking Using a Single Click on an Object
  Here, we do not use text. We simply click on an object in the first frame and let SAM3 track it.
• Detection, Segmentation, and Tracking Using Multiple Clicks on Objects
  In this case, we select multiple objects interactively and track all of them at once.

Across these examples, we will see the same core idea again and again:

• First, SAM3 recognizes what to track.
• Then, it segments it.
• Finally, it remembers and propagates it through time.

This lesson is the 3rd of a 4-part series on SAM 3:

1. SAM 3: Concept-Based Visual Understanding and Segmentation
2. Advanced SAM 3: Multi-Modal Prompting and Interactive Segmentation
3. SAM 3 for Video: Concept-Aware Segmentation and Object Tracking (this tutorial)
4. Lesson 4

To learn how to build SAM3-powered video segmentation and tracking pipelines — using text prompts, real-time webcam streams, and interactive multi-object tracking inside a dynamic Gradio interface, just keep reading.

Looking for the source code to this post?

Jump Right To The Downloads Section

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 have the following libraries installed on your system.

!pip install --q git+https://github.com/huggingface/transformers av gradio

First, we install the latest development version of transformers directly from GitHub.

We do this because SAM3 is very new. Its video processors, tracker models, and streaming APIs are not yet available in older stable releases. Installing from GitHub ensures we get access to:

• Sam3VideoProcessor and Sam3VideoModel for concept-based segmentation
• Sam3TrackerVideoProcessor and Sam3TrackerVideoModel for video tracking
• The video session and streaming inference utilities

In short, this gives us all SAM3 image and video capabilities in one place.

Next, we install av. This is a Python binding for FFmpeg. We use it to decode video files, read frames efficiently, and handle video streams coming from disk or webcam. Without av, working with video frames would be slow and unreliable.

Finally, we install gradio. We use Gradio to build interactive demos, run webcam-based real-time segmentation, create simple UI components for clicking on objects and visualizing results. This allows us to turn SAM3 into a live, interactive video application, not just a notebook script.

We also pass the --q flag to keep the installation output quiet. This keeps our notebook clean and easy to read.

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!

Setup and Imports

Once the dependencies are installed, we can import the libraries we need for video processing, model execution, and interactive demos.

import cv2
import torch
import numpy as np
import gradio as gr

from PIL import Image
from accelerate import Accelerator
from transformers.video_utils import load_video
from transformers import Sam3VideoModel, Sam3VideoProcessor
from transformers import Sam3TrackerVideoModel, Sam3TrackerVideoProcessor

First, we import OpenCV (cv2). We use OpenCV to read frames from videos and webcams, convert between different image formats, and perform basic image and video I/O operations. This forms the low-level video input layer of our system.

Next, we import PyTorch (torch). PyTorch is the backend that runs SAM3. We use it to move models and tensors to GPU or CPU, run inference efficiently, and control precision and memory usage. All heavy computation in this tutorial happens inside PyTorch.

We also import NumPy (numpy). NumPy is used for simple array manipulations, converting between OpenCV images, PIL images, and tensors, and handling masks and frame buffers in a lightweight way.

Then, we import Gradio (gradio). We use Gradio to build interactive demos, create webcam-based real-time segmentation apps, and handle mouse clicks for point-based object selection. This turns our video pipeline into a live, interactive application instead of a static script.

Next, we import PIL’s Image. This is used to convert frames into PIL format when required by the processor, handle RGB image conversions cleanly, and bridge between OpenCV and the Transformers preprocessing pipeline.

Then, we import Accelerator from Hugging Face Accelerate. The Accelerator class helps us automatically place the model on CPU or GPU, write device-agnostic code, and scale to different hardware setups without changing logic. This keeps the code clean and portable.

Next, we import load_video from transformers.video_utils. This utility function loads a video file from disk, decodes it into a list (or generator) of frames, and handles resizing and format conversion in a consistent way. We use this for file-based video experiments.

Finally, we import the SAM3 video models and processors. These are the core components of this tutorial.

• Sam3VideoModel and Sam3VideoProcessor are used for:
  • Text-prompted video segmentation
  • Concept detection on video frames
• Sam3TrackerVideoModel and Sam3TrackerVideoProcessor are used for:
  • Point-based prompting
  • Interactive object selection
  • Multi-object tracking with memory across frames

Together, these 2 model families allow us to cover all 4 workflows:

• Text-prompted tracking
• Webcam-based tracking
• Single-click object tracking
• Multi-click object tracking

Text-Prompt Video Tracking

In this section, we use a natural language prompt such as "person" or "car" to detect, segment, and track objects across an entire video. SAM3 maintains temporal memory, ensuring that object identities remain consistent from the first frame to the last. The result is a fully annotated video with masks, bounding boxes, and tracking IDs propagated over time.

Load the SAM3 Video Model

Before we process any video, we need to load the SAM3 model and its processor.

Since these models are large, we load them once and reuse them across all videos and frames.

# ------------------------------------------------
# Model setup (loaded once)
# ------------------------------------------------
accelerator = Accelerator()
device = accelerator.device

model = Sam3VideoModel.from_pretrained(
   "facebook/sam3"
).to(device, dtype=torch.bfloat16)

processor = Sam3VideoProcessor.from_pretrained("facebook/sam3")

First, we create an Accelerator object. The Accelerator automatically:

• Detects whether a GPU is available
• Chooses the best device (CPU or GPU)
• Handles device placement in a clean and consistent way

By using this, we avoid hardcoding "cuda" or "cpu" anywhere in our code. The device variable now represents where the model will run.

Next, we load the SAM3 Video model. This does 3 important things:

• It downloads the pretrained SAM3 weights from Hugging Face
• It constructs the full video-capable segmentation model
• It moves the model to the selected device (GPU or CPU)

​​We also explicitly set dtype=torch.bfloat16. This tells PyTorch to run the model in bfloat16 precision. Using bfloat16 reduces memory usage significantly, speeds up inference on modern GPUs, and has almost no impact on segmentation quality. This is especially important because SAM3 is a very large model.

Then, we load the processor. The processor is responsible for preprocessing video frames (resizing, normalization, padding), encoding text prompts, formatting inputs for the model, and post-processing outputs (resizing masks back to the original resolution).

At this point, we have:

• A fully loaded SAM3 video model on the correct device
• A video processor that knows how to prepare inputs and decode outputs

Helper Function: Visualizing Video Segmentation Masks, Bounding Boxes, and Tracking IDs

Before we start running video inference, we define a small helper function to visualize segmentation and tracking results.

This function overlays:

• Segmentation masks
• Bounding boxes
• Object IDs
• Confidence scores

directly on top of each video frame.

# ------------------------------------------------
# Visualization helper
# ------------------------------------------------
def overlay_masks_boxes(image, masks, boxes, scores, object_ids, alpha=0.5):
   image = image.copy()
   h, w = image.shape[:2]

   for i, mask in enumerate(masks):
       color = np.random.randint(0, 255, (3,), dtype=np.uint8)

       if mask.shape[-2:] != (h, w):
           mask = cv2.resize(mask.astype(np.uint8), (w, h)) > 0

       colored = np.zeros_like(image)
       colored[mask] = color
       image = cv2.addWeighted(image, 1.0, colored, alpha, 0)

       x1, y1, x2, y2 = boxes[i].astype(int)
       cv2.rectangle(image, (x1, y1), (x2, y2), color.tolist(), 2)

       label = f"ID {int(object_ids[i])} | {scores[i]:.2f}"
       cv2.putText(
           image,
           label,
           (x1, max(y1 - 5, 15)),
           cv2.FONT_HERSHEY_SIMPLEX,
           0.5,
           color.tolist(),
           1,
           cv2.LINE_AA,
       )

   return image

First, we create a copy of the input image and read its spatial resolution. We do this to avoid modifying the original frame and ensure all masks are resized to the same height and width.

Next, we loop over each detected object. Each iteration corresponds to one tracked instance in the frame. Inside the loop, we generate a random color for that object. This makes each object visually distinct in the overlay.

Then, we ensure the mask matches the image resolution. Because masks are produced at a lower resolution, we resize them to the full frame size and convert them to a Boolean mask. This ensures perfect alignment with the video frame.

Next, we create a colored overlay and blend it with the image. This:

• Paints the object region with the selected color
• Blends it with transparency (alpha)
• Keeps the original image visible underneath

Then, we draw the bounding box for the same object. This helps us visually confirm the detection region and the spatial extent of each tracked object.

Next, we prepare a label string. This shows the tracking ID assigned by SAM3 and the confidence score of the detection.

We then render this text near the bounding box. We slightly shift the text upward to avoid overlapping with the box. Finally, after all objects are drawn, we return the annotated frame.

Main Pipeline: Running the Full Video Segmentation and Tracking Workflow

Now we build the main function that runs SAM3 on a video using a text prompt and produces an annotated output video.

This function does 4 things:

• Loads the video frames
• Initializes a SAM3 video session with memory
• Propagates segmentation and tracking across frames
• Writes an annotated video to disk

Let us walk through the full pipeline step by step.

# ------------------------------------------------
# Main pipeline
# ------------------------------------------------
def run_sam3_video(video_path, text_prompt):
   # Load frames for SAM3
   video_frames, _ = load_video(video_path)

   # Reliable FPS extraction (OpenCV)
   cap = cv2.VideoCapture(video_path)
   fps = cap.get(cv2.CAP_PROP_FPS)
   cap.release()

   if fps is None or fps <= 0:
       fps = 25.0
   fps = float(fps)

   # Init SAM3 session
   inference_session = processor.init_video_session(
       video=video_frames,
       inference_device=device,
       processing_device="cpu",
       video_storage_device="cpu",
       dtype=torch.bfloat16,
   )

   inference_session = processor.add_text_prompt(
       inference_session=inference_session,
       text=text_prompt,
   )

   outputs_per_frame = {}

   for model_outputs in model.propagate_in_video_iterator(
       inference_session=inference_session,
       max_frame_num_to_track=len(video_frames),
   ):
       processed = processor.postprocess_outputs(
           inference_session,
           model_outputs,
       )
       outputs_per_frame[model_outputs.frame_idx] = processed

   # Prepare output video
   h, w = video_frames[0].shape[:2]
   out_path = "sam3_annotated.mp4"

   writer = cv2.VideoWriter(
       out_path,
       cv2.VideoWriter_fourcc(*"mp4v"),
       fps,
       (w, h),
   )

   for idx, frame in enumerate(video_frames):
       outputs = outputs_per_frame.get(idx)

       if outputs and len(outputs["object_ids"]) > 0:
           frame = overlay_masks_boxes(
               frame,
               outputs["masks"].cpu().numpy(),
               outputs["boxes"].cpu().numpy(),
               outputs["scores"].cpu().numpy(),
               outputs["object_ids"].cpu().numpy(),
           )

       # OpenCV expects BGR
       writer.write(frame[:, :, ::-1])

   writer.release()

   return out_path, out_path

First, we load the video frames using the Transformers utility. This returns:

• A list of RGB frames as NumPy arrays
• (Optionally) audio, which we ignore here

Next, we extract the frame rate using OpenCV. We do this because:

• We want the output video to play at the same speed as the input
• Some video loaders do not always return reliable FPS metadata

If FPS is missing or invalid, we fall back to a default value.

Instead of processing each frame independently, SAM3 uses a video session that maintains the following:

• Temporal memory
• Object identities
• Propagation state

We initialize the session using processor.init_video_session(...). We pass:

• The full list of frames
• The device for model inference (GPU or CPU)
• CPU for processing and storage (to save GPU memory)
• dtype=torch.bfloat16 for efficient computation

This session object now holds the entire video, memory, and tracking state.

Next, we tell SAM3 what concept we want to track using processor.add_text_prompt(...).

For example:

• "person"
• "car"
• "player in red jersey"

From this point on, the session is configured to: detect, segment, and track all instances of this concept across the video.

The real work happens when we call model.propagate_in_video_iterator(...).

This function:

• Processes frames sequentially
• Uses memory to propagate masks and identities
• Emits results for each frame

For each frame:

• We post-process raw model outputs
• Convert them into usable masks, boxes, scores, and IDs
• Store them in a dictionary indexed by frame number

At the end of this loop, we have a full timeline of segmentation and tracking results.

We now prepare an OpenCV video writer.

• Same resolution as input
• Same FPS
• MP4 format

We now loop over each frame. If SAM3 produced results for that frame: we overlay masks, boxes, IDs, and scores using our helper function. Then we write the frame to the output video. We convert RGB → BGR because OpenCV expects BGR format.

We close the video writer and return the path to the annotated video. At this point, we have a complete pipeline that: takes a video + text prompt → returns a fully segmented and tracked video.

Launch the Gradio Application

Now that our video pipeline is ready, we wrap everything into a simple Gradio web interface.

This allows us to:

• Upload a video
• Enter a text prompt (e.g., "person")
• Run SAM3 segmentation and tracking
• Preview the result
• Download the annotated video

# ------------------------------------------------
# Gradio UI
# ------------------------------------------------
with gr.Blocks() as demo:
   gr.Markdown("# 🎥 SAM3 Video Segmentation & Tracking")

   with gr.Row():
       video_input = gr.Video(label="Upload Video")
       prompt = gr.Textbox(
           label="Text Prompt",
           placeholder="e.g. person, chair, bed",
           value="person",
       )

   run_btn = gr.Button("Run Segmentation")

   video_out = gr.Video(label="Annotated Output")
   download = gr.File(label="Download Video")

   run_btn.click(
       fn=run_sam3_video,
       inputs=[video_input, prompt],
       outputs=[video_out, download],
   )

demo.launch(debug=True)

gr.Blocks() lets us build a custom UI layout instead of a simple one-function demo. Everything inside this block becomes part of our web interface.

This is just a header that appears at the top of the page.

Here, we place 2 widgets side by side:

• A video uploader
• A text box for the concept prompt

The default value is "person", so the app works immediately without typing anything.

This button will trigger the SAM3 pipeline.

We create 2 outputs:

• One to preview the annotated video
• One to download the result file

It tells Gradio: When the button is clicked, call run_sam3_video(video, prompt) and display its outputs.

Recall that our function returns 2 values.

So:

• The first output goes to the video preview
• The second output goes to the download widget

This starts a local web server and opens the interface in the browser.

The debug=True flag helps:

• Show errors in the console
• Make debugging easier during development

At this point, we have a complete application:

• upload a video
• type a concept
• get a fully segmented and tracked output video

This completes the text-prompt video segmentation and tracking pipeline.

Output: Text-Prompt Video Segmentation and Tracking Results

Real-Time Text-Prompt Tracking (Webcam)

Here, we extend text-prompt tracking to a live webcam stream. Instead of processing a preloaded video, frames arrive continuously, and SAM3 updates its tracking state in real time. This enables live, concept-aware segmentation with stable object identities across streaming frames.

Helper Function: Stable Color Overlays for Real-Time Video Tracking

In the previous pipeline, we used a visualization helper to draw masks and bounding boxes on each frame. For streaming and webcam scenarios, we slightly modify this helper so that each tracked object keeps a stable color across frames.

This small change makes tracking much easier to follow visually, especially when objects move across the scene.

Here is the updated helper:

# ------------------------------------------------
# Visualization helper
# ------------------------------------------------
def overlay_masks_boxes(image, masks, boxes, scores, object_ids, alpha=0.5):
   image = image.copy()
   h, w = image.shape[:2]

   for i, mask in enumerate(masks):
       # Stable color per object id
       rng = np.random.default_rng(int(object_ids[i]))
       color = rng.integers(0, 255, size=3, dtype=np.uint8)

       if mask.shape[-2:] != (h, w):
           mask = cv2.resize(mask.astype(np.uint8), (w, h)) > 0

       colored = np.zeros_like(image)
       colored[mask] = color
       image = cv2.addWeighted(image, 1.0, colored, alpha, 0)

       x1, y1, x2, y2 = boxes[i].astype(int)
       cv2.rectangle(image, (x1, y1), (x2, y2), color.tolist(), 2)

       label = f"ID {int(object_ids[i])} | {scores[i]:.2f}"
       cv2.putText(
           image,
           label,
           (x1, max(y1 - 5, 15)),
           cv2.FONT_HERSHEY_SIMPLEX,
           0.5,
           color.tolist(),
           1,
           cv2.LINE_AA,
       )

   return image

Let us walk through what changed and why it matters.

In streaming inference, objects appear across many frames. If colors change every frame, it becomes hard to follow which object is which.

To solve this, we generate colors based on the object ID:

Here, on Lines 10 and 11:

• The object ID is used as the random seed.
• The same object always produces the same color.
• Tracking becomes visually consistent across frames.

So if object ID 3 appears in 200 frames, it will always use the same color.

Sometimes masks are produced at a lower resolution, so we resize them to match the frame. Lines 13 and 14 ensure masks perfectly align with the streaming frame.

Next, we create a colored mask and blend it with the original frame on Lines 16-18. The transparency factor alpha controls how strongly the mask appears.

We then draw bounding boxes for each tracked object. Line 21 helps confirm detection regions visually.

Finally, we render the object ID and confidence score and draw it above the box. This allows us to verify identity consistency and inspect tracking confidence per frame (Lines 23-33).

Streaming Inference Function: Maintaining Temporal Memory Across Live Video Frames

So far, we processed videos by loading all frames first and then propagating segmentation across the entire clip.

For webcam or live-stream scenarios, this approach does not work because frames arrive one at a time. Instead, we need a streaming pipeline that:

• Maintains tracking memory across frames
• Processes each incoming frame independently
• Updates segmentation and tracking state continuously

To achieve this, we maintain a persistent session and reuse it across frames.

# ------------------------------------------------
# Global streaming state (kept across frames)
# ------------------------------------------------
STREAM_STATE = {
   "session": None,
   "prompt": None,
}

First, we create a small global structure that keeps track of the active inference session.

This dictionary stores:

• The active SAM3 inference session
• The currently used text prompt

This allows us to reuse the same session across frames instead of recreating it repeatedly.

# ------------------------------------------------
# Streaming inference function
# ------------------------------------------------
def process_webcam_frame(frame, text_prompt):
   global STREAM_STATE

   if frame is None:
       return None

   # Initialize session if needed or if prompt changed
   if (
       STREAM_STATE["session"] is None
       or STREAM_STATE["prompt"] != text_prompt
   ):
       session = processor.init_video_session(
           inference_device=device,
           processing_device="cpu",
           video_storage_device="cpu",
           dtype=torch.bfloat16,
       )
       session = processor.add_text_prompt(
           inference_session=session,
           text=text_prompt,
       )
       STREAM_STATE["session"] = session
       STREAM_STATE["prompt"] = text_prompt

   session = STREAM_STATE["session"]

   # Preprocess frame
   inputs = processor(images=frame, device=device, return_tensors="pt")

   # Streaming forward pass
   with torch.no_grad():
       model_outputs = model(
           inference_session=session,
           frame=inputs.pixel_values[0],
           reverse=False,
       )

   # Postprocess to original resolution
   outputs = processor.postprocess_outputs(
       session,
       model_outputs,
       original_sizes=inputs.original_sizes,
   )

   # Visualize
   if outputs and len(outputs["object_ids"]) > 0:
       frame = overlay_masks_boxes(
           frame,
           outputs["masks"].cpu().numpy(),
           outputs["boxes"].cpu().numpy(),
           outputs["scores"].cpu().numpy(),
           outputs["object_ids"].cpu().numpy(),
       )

   return frame

Now we define the function that processes each incoming webcam frame. We initialize the global stream state defined earlier. If no frame is available, we simply return None (Lines 14 and 15).

Next, we check whether a session already exists or whether the user has changed the prompt. If either condition is true, we create a new session. Unlike offline video processing, we do not pass frames during initialization because frames arrive one at a time. We then attach the prompt and store the session globally (Lines 18-33).

This session (Line 35) now contains:

• Tracking memory
• Object identity history
• Propagation state

across previous frames.

To convert each frame into tensors before inference, we initialize the processor, which handles resizing, normalization, tensor formatting, and device placement (Line 38).

We now run inference for the current frame (Lines 41-46). Key points:

• torch.no_grad() disables gradients, improving speed.
• The frame is processed using the existing session.
• SAM3 updates tracking memory internally.

So segmentation and identities propagate automatically.

Model outputs are resized back to original resolution. This produces masks, bounding boxes, scores, and object IDs for the current frame (Lines 49-53).

If objects are detected, we overlay results on the frame. This produces the annotated frame. The processed frame is returned to the UI or video stream (Lines 56-65).

Launch the Gradio Application

Now we connect our streaming inference pipeline to a live webcam interface using Gradio.

This allows us to:

• Capture frames directly from a webcam
• Run SAM3 segmentation and tracking in real time
• Visualize masks and tracked objects continuously

# ------------------------------------------------
# Gradio UI
# ------------------------------------------------
with gr.Blocks() as demo:
   gr.Markdown("# 📷 SAM3 Live Webcam Segmentation & Tracking")

   with gr.Row():
       webcam = gr.Image(
           sources=["webcam"],
           streaming=True,
           label="Webcam",
           type="numpy",
       )

       output = gr.Image(
           label="Live Segmentation",
           type="numpy",
       )

   prompt = gr.Textbox(
       label="Text Prompt",
       value="person",
       placeholder="e.g. person, face, chair, bottle",
   )

   webcam.stream(
       fn=process_webcam_frame,
       inputs=[webcam, prompt],
       outputs=output,
   )

demo.launch(debug=True)

Line 1 creates the Gradio app layout where we can combine multiple UI components. Line 2 displays a header explaining what the demo does.

We place input and output side by side. Inside gr.Row, we define 2 components:

First component:

• Captures frames from the webcam
• Streams frames continuously
• Sends frames as NumPy arrays to our function

The streaming=True flag enables real-time frame delivery.

The second component displays the processed frame returned by our streaming pipeline.

Next, we create a textbox for specifying the concept to track. Users can change this dynamically while the webcam runs. If the prompt changes, our pipeline automatically resets the session.

The key connection happens on Lines 26-30. This tells Gradio:

• Send every webcam frame to process_webcam_frame
• Pass along the current text prompt
• Display the returned frame in the output panel

This loop runs continuously while the webcam is active.

Finally, we launch the interface. This starts a local server and opens the demo in a browser. The debug=True flag helps diagnose errors during development.

Output: Real-Time Webcam Video Segmentation Results

Single-Click Object Tracking

In this workflow, we remove text prompts and select an object by clicking on it in the first frame. SAM3 segments the clicked object and propagates its mask throughout the video using its tracking memory. With just one foreground point, we obtain consistent object tracking across the full sequence.

Load the SAM3 Tracker Video Model

So far, we used text prompts to detect and track concepts across videos and live streams.

Now we move to a different workflow: interactive tracking, where we manually select objects and let SAM3 track them across frames.

To enable this, we switch from the text-prompt video model to the tracker-specific video model.

Here is how we load it:

# Initialize model
device = Accelerator().device
model = Sam3TrackerVideoModel.from_pretrained("facebook/sam3").to(device, dtype=torch.bfloat16)
processor = Sam3TrackerVideoProcessor.from_pretrained("facebook/sam3")

The Accelerator().device automatically selects the available hardware:

• GPU if available
• CPU otherwise

This keeps our code portable across machines.

We load the SAM3 tracker model, which is optimized for:

• Point-based object prompting
• Interactive tracking
• Multi-object identity propagation
• Frame-to-frame tracking consistency

Unlike the previous model, this one does not require text prompts. Instead, it expects clicks or point annotations.

We also move the model to the selected device and run it in bfloat16 precision, reducing memory usage and speeding up inference.

We load the processor which prepares inputs and postprocesses outputs specifically for tracking workflows. It handles frame preprocessing, prompt encoding (clicks or points), mask decoding, and identity propagation formatting.

Extract First Frame: Preparing the Initial Frame for Object Selection

Before we start tracking objects in a video, we need a way to select them.

In our workflow, we select objects by clicking on them in the first frame. So, the first step is to extract that frame from the video.

Here is a small helper function that does exactly that:

def extract_first_frame(video_path):
   """Extract first frame from video for point selection"""
   cap = cv2.VideoCapture(video_path)
   ret, frame = cap.read()
   cap.release()
   if ret:
       return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
   return None

First, OpenCV opens the video file and prepares it for frame reading. Then it attempts to read the next frame from the video.

• ret indicates whether reading succeeded.
• frame contains the actual image data.

Since we just opened the video, this call returns the first frame.

We release the video resource immediately after reading the frame. This is important because:

• It frees system resources
• It prevents file locks
• It keeps later video processing clean

OpenCV loads images in BGR format, but most visualization and processing pipelines expect RGB. So we convert BGR to RGB before returning the frame.

If the frame cannot be read, the function safely returns None.

Tracking Object Function: Propagating a Single Object Mask Across Video Frames

We now build the core function that allows us to track an object through an entire video using a single click.

The workflow is simple:

• The user uploads a video.
• The user clicks on the object in the first frame.
• SAM3 segments that object.
• The tracker propagates the object mask across all frames.
• A new annotated video is generated.

Let us walk through the implementation step by step:

def track_object(video_path, point_coords):
   """Track object through video based on clicked point"""
   if video_path is None:
       return None, "Please upload a video first"

   if point_coords is None:
       return None, "Please click on the first frame to select an object"

   try:
       # Load video
       video_frames, _ = load_video(video_path)

       # Get click coordinates
       x, y = int(point_coords[0]), int(point_coords[1])

       # Initialize session
       inference_session = processor.init_video_session(
           video=video_frames,
           inference_device=device,
           dtype=torch.bfloat16,
       )

       # Add point annotation
       points = [[[[x, y]]]]
       labels = [[[1]]]  # 1 for foreground point

       processor.add_inputs_to_inference_session(
           inference_session=inference_session,
           frame_idx=0,
           obj_ids=1,
           input_points=points,
           input_labels=labels,
       )

       # First, segment the object on the first frame
       outputs = model(
           inference_session=inference_session,
           frame_idx=0,
       )
       first_frame_masks = processor.post_process_masks(
           [outputs.pred_masks],
           original_sizes=[[inference_session.video_height, inference_session.video_width]],
           binarize=False
       )[0]

       # Propagate through video
       video_segments = {0: first_frame_masks}
       for sam3_tracker_video_output in model.propagate_in_video_iterator(inference_session):
           video_res_masks = processor.post_process_masks(
               [sam3_tracker_video_output.pred_masks],
               original_sizes=[[inference_session.video_height, inference_session.video_width]],
               binarize=False
           )[0]
           video_segments[sam3_tracker_video_output.frame_idx] = video_res_masks

       # Create output video with masks
       output_path = "/tmp/output_tracked.mp4"
       fourcc = cv2.VideoWriter_fourcc(*'mp4v')
       height, width = video_frames[0].shape[:2]
       out = cv2.VideoWriter(output_path, fourcc, 30.0, (width, height))

       for idx in range(len(video_frames)):
           frame = video_frames[idx].copy().astype(np.uint8)
           if idx in video_segments:
               masks = video_segments[idx]
               # Convert mask to float32 first, then to boolean
               mask = masks[0, 0].float().cpu().numpy() > 0.0
               # Overlay red mask
               overlay = frame.copy()
               overlay[mask] = [255, 0, 0]
               frame = cv2.addWeighted(frame.astype(np.float32), 0.6, overlay.astype(np.float32), 0.4, 0).astype(np.uint8)
           out.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))

       out.release()

       status = f"✅ Successfully tracked object through {len(video_segments)} frames at point ({x}, {y})"
       return output_path, status

   except Exception as e:
       return None, f"❌ Error: {str(e)}"

The track_object() function accepts:

• video_path: Path to the uploaded video file.
• point_coords: (x, y) coordinates of the user’s click on the first frame.

The goal is simple: “Given a video and one clicked point, track that object through the entire video.”

If no video is uploaded, tracking cannot begin. The function returns:

• None: No output video
• A message explaining the issue

Likewise, tracking requires a foreground prompt. SAM3 needs at least one positive point to know: which object to segment, and where that object exists in the first frame. Without it, tracking is undefined.

Inside a try block, the load_video() function reads the video file, extracts all frames into memory, and returns them as a list of NumPy arrays.

Why load all frames?

Because SAM3 tracking requires:

• Access to the entire temporal sequence
• Mask propagation across frames
• Internal memory consistency

Each frame shape is typically: (height, width, 3). The UI provides floating-point coordinates. We convert them to integers because:

• Pixel indices must be integers
• Mask indexing requires integer positions

This (x, y) now represents a foreground location inside frame 0.

Next, we initialize a video session which creates an internal tracking session, stores all video frames, model memory state, and object tracking buffers. We also set computation device (either CPU or GPU) and use bfloat16 for faster inference and lower memory usage. This prepares SAM3’s brain to process a full video.

Then, we prepare a foreground prompt. Since SAM3 expects inputs in batch format, points = [[[[x, y]]]] means:

• Batch size = 1
• Object ID = 1
• One point
• Coordinates (x, y)

and in labels = [[[1]]]

1 means:

• Foreground point

If it were 0, it would mean:

• Background point

So this tells SAM3: “This pixel belongs to the object.”

Then processor.add_inputs_to_inference_session() injects the prompt into the tracking session.

• frame_idx=0: Object exists in the first frame
• obj_ids=1: This is object ID 1
• input_points: Where the object is
• input_labels: Foreground signal

At this point, the model knows that object 1 is located at (x, y) in frame 0.

We explicitly run segmentation on frame 0. This produces:

• Raw mask logits
• Low-resolution mask predictions

Then processor.post_process_masks():

• Resizes masks to original resolution
• Converts internal representation into full-size masks
• Keeps them as float probabilities (not binarized)

We now have a full-resolution mask for frame 0.

On Line 47, we store frame 0 results first.​​ Then model.propagate_in_video_iterator() runs SAM3’s tracking mechanism.

What happens internally:

• It uses memory from frame 0
• Matches object appearance across frames
• Predicts masks for each new frame

For each frame processor.post_process_masks(...), we resize masks and store them in a dictionary.

Final structure:

video_segments = {
  0: mask0,
  1: mask1,
  2: mask2,
  ...
}

Now we have segmentation results for the full video.

Next we define the output video path, and define the codec format to mp4v. We get the resolution which ensures output video matches input resolution. We also initialize the OpenCV writer using cv2.VideoWriter() and pass the output video path, codec format, 30 FPS, and height and width.

We then iterate over each frame and create a copy to avoid modifying the original frames. We move the output tensor to the CPU, convert it to a NumPy array, compute probabilities, and threshold the result to obtain a Boolean mask. The resulting mask is True where the object exists and False elsewhere.

We then overlay a red mask and blend it using cv2.addWeighted(), producing a frame with 60% original content and 40% red overlay for smooth visualization. Because OpenCV expects BGR format, we convert the frame using cv2.cvtColor() with the cv2.COLOR_RGB2BGR flag.

out.release() finalizes the file, writes any remaining buffers, and closes the video properly. Without this call, the output file may become corrupted.

Finally, we return the path to the saved video and a success message. If an error occurs:

• The function safely returns an error message
• The application does not crash

Launch the Gradio Application

We now connect our tracking pipeline to an interactive Gradio interface. This interface allows users to upload a video, click on an object in the first frame, and automatically track that object across the entire clip.

Here is the full interface code:

# Create Gradio interface with blocks for better control
with gr.Blocks(title="SAM3 Video Tracker") as demo:
   gr.Markdown("# 🎯 SAM3 Video Object Tracker")
   gr.Markdown("Upload a video and click on an object in the first frame to track it throughout the video")

   with gr.Row():
       with gr.Column():
           video_input = gr.Video(label="Upload Video")
           first_frame = gr.Image(label="Click on object to track", type="numpy")
           point_display = gr.Textbox(label="Selected Point", interactive=False)
           track_btn = gr.Button("Track Object", variant="primary")

       with gr.Column():
           video_output = gr.Video(label="Tracked Video")
           status_output = gr.Textbox(label="Status")

   # Store clicked point
   clicked_point = gr.State(None)

   # Extract first frame when video is uploaded
   def on_video_upload(video):
       if video:
           frame = extract_first_frame(video)
           return frame, None, "Upload complete. Click on the object you want to track."
       return None, None, ""

   video_input.change(
       on_video_upload,
       inputs=[video_input],
       outputs=[first_frame, clicked_point, status_output]
   )

   # Handle click on first frame
   def on_click(img, evt: gr.SelectData):
       x, y = evt.index[0], evt.index[1]
       # Draw a circle on the clicked point
       img_copy = img.copy()
       cv2.circle(img_copy, (x, y), 5, (255, 0, 0), -1)
       return img_copy, (x, y), f"Point selected: ({x}, {y})"

   first_frame.select(
       on_click,
       inputs=[first_frame],
       outputs=[first_frame, clicked_point, point_display]
   )

   # Track button
   track_btn.click(
       track_object,
       inputs=[video_input, clicked_point],
       outputs=[video_output, status_output]
   )

# Launch
demo.launch(debug=True)

The Gradio interface is built using gr.Blocks, which gives full control over layout, components, and event handling. The goal is simple: “Allow a user to upload a video, click on a single object in the first frame, and track that object throughout the entire video.”

At the top of the interface, we display two gr.Markdown() sections. The first acts as the main heading so users immediately understand what the application does. The second provides short instructions explaining the workflow: upload a video, click on an object in the first frame, and then track it.

Next, we structure the layout using a gr.Row(). Inside that row, we create two columns. The left column contains all inputs and interactions. The right column displays outputs.

In the left column, we first add a video upload component. This allows the user to upload a video file from their system. Once uploaded, the backend receives the file path. That file path is later used to extract frames and run tracking.

Below the video upload, we place an image component. This image will display the first frame of the uploaded video. We set its type to NumPy so that the backend receives the frame as a NumPy array. This is important because we draw visual markers on the frame using OpenCV when the user clicks.

Below the image, we add a textbox labeled "Selected Point". This textbox is non-interactive, meaning the user cannot manually edit it. It simply displays the coordinates of the selected point so the user can confirm their click.

Under that, we add a "Track Object" button. This button is styled as primary so it visually stands out as the main action. When clicked, it triggers the tracking pipeline.

In the right column, we create a video output component. This will display the processed video after tracking is complete. Below it, we add a status textbox. This displays messages such as upload confirmation, tracking success, or error details.

To maintain interaction state, we use a Gradio State variable called clicked_point. This variable stores the coordinates of the selected object. Initially, it is set to None. State is important because the click event and the tracking event happen at different times, and we need a way to remember which point the user selected.

When a video is uploaded, a function on_video_upload() is triggered. This function checks whether a valid video exists. If it does, we extract the first frame using a helper function. That first frame is returned to the image component so the user can see it. We also reset the stored clicked point to None, ensuring that any previous selection is cleared. Finally, we return a status message informing the user that the upload is complete and they should click on an object.

If no video is uploaded, the function returns empty values, keeping the interface clean.

When the user clicks on the first frame image, another function on_click() handles the event. The click event provides the pixel coordinates of the selected location through evt.index. We extract the x and y coordinates from that event data.

Next, we create a copy of the displayed image. This is important because we do not want to modify the original image directly. On this copy, we draw a small filled circle at the clicked location using OpenCV. The circle visually marks the selected object so the user knows exactly where they clicked.

After drawing the marker, we return 3 things:

• The updated image with the circle drawn
• The tuple (x, y) stored in the state variable
• A formatted string such as "Point selected: (x, y)" displayed in the textbox

This ensures the UI updates immediately and the selected point is stored for later use.

The Track Object button is connected to the backend tracking function. When pressed, it sends 2 inputs:

• The uploaded video path
• The stored clicked point

The tracking function then performs segmentation and mask propagation across the entire video. Once processing is complete, it returns:

• The path to the output tracked video
• A status message indicating success or failure

These outputs are displayed in the video output component and the status textbox, respectively.

Finally, the application is launched with debug mode enabled. Debug mode prints detailed logs in case of errors, which is helpful during development and testing.

The complete flow works as follows:

• The user uploads a video.
• The first frame is extracted and displayed.
• The user clicks on an object.
• A visual marker appears and the coordinates are stored.
• The user presses Track Object.
• The backend processes the video and returns the tracked result.
• The output video and status message are displayed.

This design keeps the interface simple, intuitive, and focused on a single-click tracking workflow while properly managing state and user interaction.

Output: Single-Click Video Object Tracking Results

Multi-Click Object Tracking

In this final setup, we select multiple objects by clicking different locations in the first frame. Each click initializes a unique object ID, and SAM3 tracks all selected objects simultaneously. The output video shows multiple masks with distinct colors, preserving identity consistency across frames.

Initialize Few Colors: Defining a Color Palette for Multi-Object Tracking Visualization

When tracking multiple objects at the same time, visualization becomes very important. If all objects share the same mask color, it becomes difficult to understand which mask corresponds to which object.

To solve this, we assign different colors to different tracked objects.

Here is a small color palette we use:

# Different colors for different objects
COLORS = [
   [255, 0, 0],    # Red
   [0, 255, 0],    # Green
   [0, 0, 255],    # Blue
   [255, 255, 0],  # Yellow
   [255, 0, 255],  # Magenta
   [0, 255, 255],  # Cyan
   [255, 128, 0],  # Orange
   [128, 0, 255],  # Purple
]

Each entry in this list represents an RGB color used to render masks and overlays for a tracked object.

For example:

• Object 1 may appear in red
• Object 2 in green
• Object 3 in blue
• and so on.

During visualization, we typically assign colors based on object index or object ID, cycling through the list if the number of objects exceeds the available colors.

Extract First Frame: Preparing the First Frame for Multi-Object Selection

For multi-object tracking, we again begin by extracting the first frame of the video. This frame is used as the interaction surface where users click on multiple objects they want to track.

The helper function below reads the first frame from the uploaded video.

def extract_first_frame(video_path):
   """Extract first frame from video for point selection"""
   cap = cv2.VideoCapture(video_path)
   ret, frame = cap.read()
   cap.release()
   if ret:
       return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
   return None

First, we open the video using OpenCV’s VideoCapture. Immediately after opening, we read a single frame. Since the video has just been opened, this corresponds to the very first frame.

Next, we release the video handle to free system resources and avoid locking the file for later processing steps.

OpenCV loads images in BGR format, but our visualization and model pipelines expect RGB images. Therefore, we convert the frame from BGR to RGB before returning it.

If frame extraction fails, the function safely returns None, allowing the application to handle the error gracefully.

This function now allows users to click on multiple objects in the first frame, which we will use as prompts for tracking several objects simultaneously in the next step.

Tracking Object Function: Tracking Multiple Objects with Unique IDs Across Video Frames

def track_objects(video_path, points_list):
   """Track multiple objects through video based on clicked points"""
   if video_path is None:
       return None, "Please upload a video first"

   if not points_list or len(points_list) == 0:
       return None, "Please click on at least one object in the first frame"

   try:
       # Load video
       video_frames, _ = load_video(video_path)

       # Initialize session
       inference_session = processor.init_video_session(
           video=video_frames,
           inference_device=device,
           dtype=torch.bfloat16,
       )

       # Prepare points for all objects
       obj_ids = list(range(1, len(points_list) + 1))
       input_points = [[[[int(x), int(y)]] for x, y in points_list]]
       input_labels = [[[1] for _ in points_list]]  # All are foreground points

       # Add all objects to inference session
       processor.add_inputs_to_inference_session(
           inference_session=inference_session,
           frame_idx=0,
           obj_ids=obj_ids,
           input_points=input_points,
           input_labels=input_labels,
       )

       # First, segment objects on the first frame
       outputs = model(
           inference_session=inference_session,
           frame_idx=0,
       )
       first_frame_masks = processor.post_process_masks(
           [outputs.pred_masks],
           original_sizes=[[inference_session.video_height, inference_session.video_width]],
           binarize=False
       )[0]

       # Initialize video segments with first frame
       video_segments = {0: {
           obj_id: first_frame_masks[i]
           for i, obj_id in enumerate(inference_session.obj_ids)
       }}

       # Propagate through video
       for sam3_tracker_video_output in model.propagate_in_video_iterator(inference_session):
           video_res_masks = processor.post_process_masks(
               [sam3_tracker_video_output.pred_masks],
               original_sizes=[[inference_session.video_height, inference_session.video_width]],
               binarize=False
           )[0]
           video_segments[sam3_tracker_video_output.frame_idx] = {
               obj_id: video_res_masks[i]
               for i, obj_id in enumerate(inference_session.obj_ids)
           }

       # Create output video with masks
       output_path = "/tmp/output_tracked.mp4"
       fourcc = cv2.VideoWriter_fourcc(*'mp4v')
       height, width = video_frames[0].shape[:2]
       out = cv2.VideoWriter(output_path, fourcc, 30.0, (width, height))

       for idx in range(len(video_frames)):
           frame = video_frames[idx].copy().astype(np.uint8)

           if idx in video_segments:
               # Create overlay for all objects
               overlay = frame.copy().astype(np.float32)

               for obj_idx, (obj_id, masks) in enumerate(video_segments[idx].items()):
                   # Convert mask to float32 first, then to boolean
                   mask = masks[0].float().cpu().numpy() > 0.0

                   # Use different color for each object
                   color = COLORS[obj_idx % len(COLORS)]
                   overlay[mask] = np.array(color, dtype=np.float32)

               # Blend overlay with original frame
               frame = cv2.addWeighted(frame.astype(np.float32), 0.6, overlay, 0.4, 0).astype(np.uint8)

           out.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))

       out.release()

       status = f"✅ Successfully tracked {len(points_list)} object(s) through {len(video_segments)} frames"
       return output_path, status

   except Exception as e:
       import traceback
       return None, f"❌ Error: {str(e)}\n{traceback.format_exc()}"

The track_objects() function accepts:

• video_path: Path to the uploaded video file.
• points_list: A list of (x, y) coordinates where the user clicked on the first frame (one click per object).

The goal is simple: “Given a video and multiple clicked points, track all selected objects through the entire video.”

If no video is uploaded, tracking cannot begin. The function returns:

• None: No output video
• A message explaining the issue

Likewise, tracking requires at least one foreground prompt. SAM3 needs one or more positive points to know:

• Which objects to segment
• Where those objects exist in the first frame

If points_list is empty, tracking is undefined.

Inside a try block, the load_video() function reads the video file, extracts all frames into memory, and returns them as a list of NumPy arrays.

Why load all frames?

Because SAM3 tracking requires:

• Access to the entire temporal sequence
• Mask propagation across frames
• Internal memory consistency

Each frame shape is typically: (height, width, 3). Next, we initialize a video session using processor.init_video_session(). This creates an internal tracking session that:

• Stores all video frames
• Maintains model memory state
• Manages object tracking buffers

We also:

• Set the computation device (CPU or GPU)
• Use bfloat16 for faster inference and lower memory usage

This step prepares SAM3’s internal tracking mechanism to process the full video.

Now we prepare inputs for multiple objects. If the user clicked 3 points, this becomes:

[1, 2, 3]

Each clicked point is treated as a separate object with its own ID.

Then we structure the coordinates:

input_points = [[[[int(x), int(y)]] for x, y in points_list]]

SAM3 expects batch format:

[batch][object][points][coordinates]

So this structure means:

• Batch size = 1
• Multiple objects
• One point per object
• Each point has (x, y)

We convert coordinates to integers because:

• Pixel indices must be integers
• Mask indexing requires integer positions

Each (x, y) now represents a foreground location for a different object in frame 0.

Next, we define labels:

input_labels = [[[1] for _ in points_list]]

1 means:

• Foreground point

If it were 0, it would mean:

• Background point

So this tells SAM3: “Each of these clicked pixels belongs to a separate object.”

Then we inject everything into the inference session using processor.add_inputs_to_inference_session():

Here:

• frame_idx=0: Objects exist in the first frame
• obj_ids: Multiple object identifiers
• input_points: Click locations
• input_labels: Foreground signals

At this point, the model knows that multiple objects are present at the selected coordinates in frame 0.

Next, we explicitly run segmentation on frame 0. This produces:

• Raw mask logits
• Low-resolution mask predictions for all objects

Then we post-process the masks using processor.post_process_masks(...). This step:

• Resizes masks to original resolution
• Converts internal representation into full-size masks
• Keeps them as float probabilities (not binarized)

We now have: Full-resolution masks for all selected objects in frame 0.

Now we initialize storage on Lines 46-49:

This means:

• Frame 0 results are stored first
• Each object ID maps to its own mask

Structure becomes:

video_segments = {
   0: {
       1: mask_for_object_1,
       2: mask_for_object_2,
       ...
   }
}

Next, we propagate through the video using model.propagate_in_video_iterator(). This runs SAM3’s tracking mechanism.

What happens internally:

• It uses memory from frame 0
• Matches object appearance across frames
• Maintains identity consistency for each object
• Predicts masks for each new frame

For every frame:

• We resize masks
• Store them in a dictionary per object

Final structure:

video_segments = {
 0: {1: mask0_1, 2: mask0_2},
 1: {1: mask1_1, 2: mask1_2},
 2: {1: mask2_1, 2: mask2_2},
 ...
}

Now we have segmentation results for all objects across the full video.

Next, we define:

• Output video path
• Codec format (mp4v)

We get the resolution from the first frame to ensure the output video matches the input resolution. Then we initialize cv2.VideoWriter() which sets:

• Output path
• Codec
• 30 FPS
• Frame dimensions

Now we iterate through each frame. We copy each frame to avoid modifying the original.

If masks exist for that frame:

• We create an overlay
• For each object:
  • Convert tensor: CPU
  • Convert to NumPy
  • Convert probabilities: Boolean mask

Now:

• True: Object pixels
• False: Background

Unlike the single-object version, here we:

• Assign a different color for each object
• Use COLORS[obj_idx % len(COLORS)]
• Overlay masks for multiple objects

Then we blend using cv2.addWeighted(). This results in:

• 60% original frame
• 40% colored overlay
• Smooth visualization

OpenCV expects BGR format, so we convert using cv2.COLOR_RGB2BGR.

Finally, out.release() finalizes the file, writes any remaining buffers, and properly closes the video. Without this step, the output video may become corrupted.

At the end, we return:

• Path to the saved video
• A success message indicating how many objects were tracked

If anything fails:

• The function safely returns the error
• The application does not crash
• The traceback is included for debugging

Launch the Gradio Application

# Create Gradio interface with blocks for better control
with gr.Blocks(title="SAM3 Multi-Object Video Tracker") as demo:
   gr.Markdown("# 🎯 SAM3 Multi-Object Video Tracker")
   gr.Markdown("Upload a video and click on multiple objects in the first frame to track them. Each object gets a different color!")

   with gr.Row():
       with gr.Column():
           video_input = gr.Video(label="Upload Video")
           first_frame = gr.Image(label="Click on objects to track (multiple clicks supported)", type="numpy")

           with gr.Row():
               clear_points_btn = gr.Button("Clear Points", variant="secondary")
               track_btn = gr.Button("Track Objects", variant="primary")

           points_display = gr.Textbox(label="Selected Points", interactive=False, lines=5)

       with gr.Column():
           video_output = gr.Video(label="Tracked Video")
           status_output = gr.Textbox(label="Status")
           gr.Markdown("""
           ### Color Legend:
           - 🔴 Red - Object 1
           - 🟢 Green - Object 2
           - 🔵 Blue - Object 3
           - 🟡 Yellow - Object 4
           - 🩷 Magenta - Object 5
           - 🩵 Cyan - Object 6
           - 🟠 Orange - Object 7
           - 🟣 Purple - Object 8
           """)

   # Store clicked points and original frame
   clicked_points = gr.State([])
   original_frame = gr.State(None)

   # Extract first frame when video is uploaded
   def on_video_upload(video):
       if video:
           frame = extract_first_frame(video)
           return frame, frame, [], "Upload complete. Click on objects you want to track."
       return None, None, [], ""

   video_input.change(
       on_video_upload,
       inputs=[video_input],
       outputs=[first_frame, original_frame, clicked_points, status_output]
   )

   # Handle click on first frame
   def on_click(img, orig_frame, points, evt: gr.SelectData):
       if orig_frame is None:
           return img, points, "Please upload a video first"

       x, y = evt.index[0], evt.index[1]

       # Add point to list
       points.append((x, y))

       # Draw all points on the image
       img_copy = orig_frame.copy()
       for i, (px, py) in enumerate(points):
           color = COLORS[i % len(COLORS)]
           cv2.circle(img_copy, (px, py), 8, tuple(color), -1)
           cv2.circle(img_copy, (px, py), 10, (255, 255, 255), 2)
           # Add number label
           cv2.putText(img_copy, str(i+1), (px+15, py+5),
                      cv2.FONT_HERSHEY_SIMPLEX, 0.6, tuple(color), 2)

       points_text = "\n".join([f"Object {i+1}: ({x}, {y})" for i, (x, y) in enumerate(points)])

       return img_copy, points, points_text

   first_frame.select(
       on_click,
       inputs=[first_frame, original_frame, clicked_points],
       outputs=[first_frame, clicked_points, points_display]
   )

   # Clear points button
   def clear_points(orig_frame):
       return orig_frame, [], ""

   clear_points_btn.click(
       clear_points,
       inputs=[original_frame],
       outputs=[first_frame, clicked_points, points_display]
   )

   # Track button
   track_btn.click(
       track_objects,
       inputs=[video_input, clicked_points],
       outputs=[video_output, status_output]
   )

# Launch
demo.launch(debug=True)

The Gradio interface is built using gr.Blocks, which allows us to design a structured, interactive layout with full control over components and events. The goal is simple: “Create an interactive UI where a user uploads a video, clicks on multiple objects in the first frame, and then tracks them across the entire video.”

At the top of the interface, we display a title and short instructions using gr.Markdown(). This helps users immediately understand what the application does and what steps they need to follow. Clear instructions reduce confusion and improve usability.

Next, we organize the layout using a gr.Row(). Inside that row, we create 2 columns using gr.Columns(). The left column handles inputs and interactions. The right column displays outputs and tracking results. This separation keeps the workflow intuitive and clean.

In the left column, we first create a video upload component. This allows the user to upload a video file from their system. Once a video is uploaded, the backend receives the file path, which is later used to extract frames and perform tracking.

Below the video upload, we place an image component that displays the first frame of the uploaded video. This image is interactive and supports click events. We set its type to NumPy so the backend receives the image as a NumPy array. This is important because we draw circles and labels on the frame using OpenCV.

Under the image, we add 2 buttons side by side. The first button clears selected points. The second button triggers the tracking pipeline. The track button is styled as the primary button so it visually stands out as the main action.

Below the buttons, we add a textbox that displays selected points. This textbox is non-interactive, meaning users cannot edit it manually. It simply shows a formatted list of selected objects and their coordinates. This helps users confirm that they clicked the correct locations before starting tracking.

In the right column, we create a video output component. This will display the processed video returned by the tracking function. Below that, we add a status textbox to show messages such as upload confirmation, tracking success, or error details. Finally, we display a color legend using Markdown so users understand which color corresponds to which object during visualization.

To manage interaction data across events, we use Gradio’s State component. One state variable stores the list of clicked points. This list grows as the user clicks on multiple objects. Another state variable stores the original first frame. This is important because every time a new click occurs, we redraw all points on a clean copy of the original frame instead of repeatedly drawing over an already modified image. Without this, the markers would stack incorrectly and distort the visualization.

When a video is uploaded, a function on_video_upload() is triggered. This function extracts the first frame from the video and returns 4 values:

• The extracted first frame for display
• The same frame stored as the original clean frame
• An empty list of clicked points
• A status message confirming upload completion

This ensures that each new upload resets the application state properly.

When the user clicks on the first frame image, another function on_click() handles the event. The click event provides the pixel coordinates of the selected location. First, we check whether a video has been uploaded. If not, we return a message asking the user to upload one.

If a frame exists, we extract the x and y coordinates from the click event. We then append this coordinate pair to the stored list of points. After updating the list, we redraw the image. We copy the original clean frame and loop over all stored points. For each point:

• We select a color from the predefined COLORS list
• We draw a filled circle at the clicked location
• We draw a white border around the circle for better visibility
• We add a numeric label next to the point indicating Object 1, Object 2, and so on

This ensures each selected object is visually distinct and clearly labeled.

We also generate formatted text listing all selected objects and their coordinates. This text is displayed in the textbox so the user can verify selections.

The Clear Points button is connected to a function that resets the interface. It restores the original clean frame, empties the clicked points list, and clears the points display textbox. This allows the user to start fresh without reloading the video.

The Track Objects button is connected to the tracking function. When clicked, it sends:

• The uploaded video
• The stored list of clicked points

to the backend tracking pipeline. The tracking function processes the video, segments and propagates masks for all selected objects, and returns:

• The path to the processed output video
• A status message

These are then displayed in the output video component and the status textbox.

Finally, the application is launched with debug mode enabled. Debug mode provides detailed logs in case errors occur during development, making troubleshooting easier.

Overall, the interface follows this flow:

• The user uploads a video.
• The first frame is extracted and displayed.
• The user clicks multiple objects.
• Points are stored and visualized with colors and labels.
• The user presses Track Objects.
• The backend processes the video and returns the tracked result.
• The output video and status message are displayed.

State management ensures smooth interaction across multiple events, and the 2-column layout keeps inputs and outputs clearly separated.

Output: Multi-Object Video Segmentation and Tracking Results

What's next? We recommend PyImageSearch University.

Course information:
86+ total classes • 115+ hours hours of on-demand code walkthrough videos • Last updated: April 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:

• ✓ 86+ courses on essential computer vision, deep learning, and OpenCV topics
• ✓ 86 Certificates of Completion
• ✓ 115+ hours 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 tutorial, we extended SAM3 from image-based segmentation workflows to full video understanding and tracking. We first built pipelines that detect, segment, and track concepts across videos using simple text prompts, enabling automatic tracking of objects such as people or vehicles without manual annotation.

Next, we moved to streaming inference, where SAM 3 processes frames continuously from a webcam while maintaining tracking memory across time. This allowed us to build real-time segmentation and tracking systems that operate on live video streams.

We then explored interactive tracking workflows, where users select objects directly using click prompts. Starting from single-object tracking, we progressed to multi-object tracking, enabling several objects to be tracked simultaneously with consistent identity and color-coded visualization.

By the end of this tutorial, we developed complete end-to-end systems that combine detection, segmentation, tracking, and interactive workflows into practical applications using Gradio interfaces. Together with the previous parts of this series, we now have a full understanding of how SAM 3 enables concept-aware segmentation and tracking across both images and videos, opening the door to intelligent video editing, annotation, and analysis workflows.

Citation Information

Thakur, P. “SAM 3 for Video: Concept-Aware Segmentation and Object Tracking,” PyImageSearch, S. Huot, G. Kudriavtsev, and A. Sharma, eds., 2026, https://pyimg.co/luxfd

@incollection{Thakur_2026_sam-3-sam3-for-video-concept-aware-segmentation-and-object-tracking,
  author = {Piyush Thakur},
  title = {{SAM 3 for Video: Concept-Aware Segmentation and Object Tracking}},
  booktitle = {PyImageSearch},
  editor = {Susan Huot and Georgii Kudriavtsev and Aditya Sharma},
  year = {2026},
  url = {https://pyimg.co/luxfd},
}

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 SAM 3 for Video: Concept-Aware Segmentation and Object Tracking appeared first on PyImageSearch.

Vector Search Using Ollama for Retrieval-Augmented Generation (RAG)

In the previous lessons, you learned how to generate text embeddings, store them efficiently, and perform fast vector search using FAISS. Now, it’s time to put that search power to use — by connecting it with a language model to build a complete Retrieval-Augmented Generation (RAG) pipeline.

RAG is the bridge between retrieval and reasoning — it lets your LLM (large language model) access facts it hasn’t memorized. Instead of relying solely on pre-training, the model fetches relevant context from your own data before answering, ensuring responses that are accurate, up-to-date, and grounded in evidence.

Think of it as asking a well-trained assistant a question: they don’t guess — they quickly look up the right pages in your company wiki, then answer with confidence.

This lesson is the last of a 3-part series on Retrieval-Augmented Generation (RAG):

1. TF-IDF vs. Embeddings: From Keywords to Semantic Search
2. Vector Search with FAISS: Approximate Nearest Neighbor (ANN) Explained
3. Vector Search Using Ollama for Retrieval-Augmented Generation (RAG) (this tutorial)

To learn how to make your LLM do the same, just keep reading.

Looking for the source code to this post?

Jump Right To The Downloads Section

How Vector Search Powers Retrieval-Augmented Generation (RAG)

Before we start wiring our first Retrieval-Augmented Generation (RAG) pipeline, let’s pause to understand how far we’ve come — and why this next step is a natural progression.

In Lesson 1, we learned how to translate language into geometry.

Each sentence became a vector — a point in high-dimensional space — where semantic closeness means directional similarity. Instead of matching exact words, embeddings capture meaning.

In Lesson 2, we tackled the scale problem: when millions of such vectors exist, finding the nearest ones efficiently demands specialized data structures such as FAISS indexes — Flat, HNSW, and IVF.

These indexes allow us to perform lightning-fast approximate nearest neighbor (ANN) searches with only a small trade-off in precision.

Now, in Lesson 3, we finally connect this retrieval ability to an LLM.

Think of the FAISS index as a semantic memory vault — it remembers every sentence you’ve embedded.

RAG acts as the retrieval layer that fetches the most relevant facts when you ask a question, passing those snippets to the model before it generates an answer.

From Search to Context

Traditional vector search stops at retrieval:

You enter a query, it finds semantically similar passages, and displays them as search results.

RAG goes one step further — it feeds those retrieved passages into the language model’s input prompt.

Instead of reading raw similarity scores, the model sees sentences such as:

Context:
1. Vector databases store and search embeddings efficiently using ANN.
2. FAISS supports multiple indexing strategies including Flat, HNSW, and IVF.

User Question:
What’s the advantage of using HNSW over Flat indexes?

Now the model doesn’t have to “guess” — it answers with contextually grounded reasoning.

That is what transforms search into retrieval-based reasoning (Figure 1).

The Flow of Meaning

Let’s connect all the components (Table 1).

This is the essence of RAG — combining the recall strength of retrieval with the reasoning power of generation.

Putting It All Together

Imagine browsing through a giant photo album of your entire text corpus.

Vector search helps you instantly find pictures with similar colors and patterns — that’s embeddings at work.

But RAG doesn’t stop there. It shows those pictures to a storyteller (the LLM), who uses them to narrate a coherent story about what’s happening across them.

Embeddings give you semantic lookup.

RAG gives you semantic understanding (Figure 2).

If this flow made sense, you’re ready for the real action — understanding how Retrieval Augmented Generation actually works under the hood.

Next, we’ll break down the architecture, components, and the 2-stage process that powers modern RAG pipelines.

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.

What Is Retrieval-Augmented Generation (RAG)?

Large Language Models (LLMs) have changed how we interact with information.

But they come with two fundamental weaknesses: they can’t access external data and they forget easily.

Even the most powerful LLMs (e.g., GPT-4 or Mistral) rely entirely on patterns learned during training.

They don’t know about the latest company reports, your private PDFs, or a proprietary codebase unless explicitly retrained — which is expensive, slow, and often impossible for organizations working with sensitive data.

This is exactly where Retrieval-Augmented Generation (RAG) steps in.

RAG acts as a bridge between frozen LLM knowledge and fresh, external information.

Instead of forcing the model to memorize everything, we give it a retrieval memory system — a searchable knowledge store filled with your domain data.

Imagine giving your LLM a library card — and access to an intelligent librarian.

Whenever a question arrives, the LLM doesn’t rely on its memory alone — it sends the librarian to fetch relevant documents, reads them carefully, and then generates a grounded, evidence-based response.

The Retrieve-Read-Generate Architecture Explained

RAG systems follow a predictable 3-step pipeline that connects information retrieval with text generation:

Retrieve

The user’s question is first converted into a numerical vector (embedding).

This vector represents the semantic meaning of the query and is matched against stored document vectors in a vector index (e.g., FAISS, Pinecone, or Milvus).

The top-k closest matches — meaning the most semantically similar chunks — are returned as potential context.

Read

These retrieved chunks are merged into a short context window — effectively a mini-knowledge pack relevant to the user’s query.

This step is vital: instead of dumping the entire corpus into the model, we pass only the most useful and concise context.

Generate

The LLM (e.g., one running locally through Ollama or remotely via an API) takes both the query and retrieved context, then composes an answer that blends natural language fluency with factual grounding.

If well-designed, the model avoids hallucinating and gracefully responds “I don’t know” when information is missing.

Figure 3 displays a high-level visual summary of this process.

Why Retrieval-Augmented Generation (RAG) Improves LLM Accuracy

At first glance, RAG may appear to be “just another way to query a model,” but it represents a fundamental shift in how LLMs reason.

Traditional LLMs store knowledge in their parameters — they memorize facts.

RAG decouples knowledge from parameters and instead retrieves it on demand.

This means you can keep your model small, fast, and efficient, while still answering domain-specific queries with accuracy.

Let’s unpack this with a few concrete advantages, as reported in Table 2.

The result?

A modular intelligence system — where the retriever evolves with your data, and the generator focuses purely on language reasoning.

The Broader Picture: A Hybrid of Search and Generation

You can think of RAG as the perfect fusion of information retrieval and natural language generation.

Traditional search engines stop at retrieval — they return ranked documents.

LLMs go further — they interpret and explain.

RAG combines both: find relevant context, then generate insights from it.

It’s the same principle behind how humans answer questions:

• We first recall or look up what we know.
• Then we synthesize an answer in our own words.

RAG gives LLMs the same skill — combining retrieval precision with generative fluency.

Key Takeaway

RAG doesn’t replace fine-tuning — it complements it.

It’s the fastest, cheapest, and most reliable way to make LLMs domain-aware without touching their weights.

Once you set up your retriever (built from the FAISS indexes we created in Lesson 2) and connect it to a generator (which we’ll later run via Ollama), you’ll have a self-contained intelligent assistant — one that can reason over your data and answer complex questions in natural language.

How to Build a RAG Pipeline with FAISS and Ollama (Local LLM)

Now that you understand what Retrieval Augmented Generation is and why it matters, let’s break down how to actually build one — conceptually first, before we dive into the code.

A RAG pipeline may sound complicated, but in practice it’s a clean, modular system made of 3 major parts: the retriever, the reader, and the generator.

Each part does one job well, and together they form the backbone of every production-grade RAG system — whether you’re querying a few PDFs or an entire knowledge base.

Step 1: Implementing HNSW Vector Search with FAISS for RAG

The retriever’s job is to search your document corpus and return the chunks most relevant to a user query.

It’s powered by the vector indexes you built in Lesson 2, which enable efficient approximate nearest-neighbor (ANN) search.

When a user asks a question, here’s what happens:

• The query text is embedded using the same Sentence Transformer model used during indexing.
• That query embedding is compared with your stored document embeddings via a FAISS index.
• The retriever returns the top-k results (typically 3-5 chunks) ranked by cosine similarity.

Think of it as Google Search for your private data — except instead of matching keywords, it matches meaning (Figure 4).

Step 2: Prompt Engineering for Retrieval-Augmented Generation (RAG)

Once the relevant chunks are retrieved, we can’t just throw them at the LLM.

They must be assembled and formatted into a coherent, bounded prompt.

This is the job of the reader — a lightweight logic layer that:

• Ranks and filters retrieved chunks by similarity score or metadata (e.g., document name or section).
• Merges them into a context block that stays within the LLM’s context-window limit (say, 4K-8K tokens).
• Wraps them inside a consistent prompt template.

In our code, this will be handled using utilities from config.py — notably build_prompt(), which combines system prompts, retrieved text, and user queries into a final message ready for the model (Figure 5).

Step 3: Generating Grounded Answers with Ollama Local LLM

Finally, the generator — your LLM — reads the composed prompt and generates a response grounded in the retrieved data.

In our implementation, this will be the stage where we integrate with Ollama, a local LLM runtime capable of running models (e.g., Llama 3, Mistral, or Gemma 2) on your machine.

But the design will stay framework-agnostic, so you can later swap Ollama for an API call to OpenAI, Claude, or an enterprise model running in-house.

What makes this step powerful is the synergy between retrieval and generation: the LLM isn’t hallucinating — it’s reasoning with evidence. If the context doesn’t contain the answer, it should politely say so, thanks to the strict vs. synthesis prompt patterns defined in config.py (Figure 6).

Adding Feedback Loops to Improve Retrieval Accuracy

In more advanced systems, RAG doesn’t end at generation. You can capture user feedback (e.g., thumbs-up/down or re-query actions) to fine-tune retrieval parameters, re-rank documents, or even re-embed sections of your corpus. This transforms a static RAG setup into a continually learning knowledge engine.

Putting It All Together

Figure 7 displays a conceptual flow that ties the 3 components together.

Each box in this pipeline maps directly to a section of your upcoming implementation.

In code, these steps will unfold through modular utilities and clean interfaces so you can swap retrievers, tweak prompt templates, or change models without rewriting the entire system.

Configuring Your Development Environment: Setting Up Ollama and FAISS for a Local RAG Pipeline

To follow this RAG pipeline guide, you’ll need several Python packages installed on your system. The tutorial builds upon semantic embeddings and vector search, requiring machine learning libraries, HTTP clients, and visualization tools.

$ pip install sentence-transformers==2.7.0
$ pip install faiss-cpu==1.8.0.post1
$ pip install numpy==1.26.4
$ pip install requests==2.32.3
$ pip install rich==13.8.1

Optional Dependencies

For visualization and enhanced functionality:

$ pip install scikit-learn==1.5.1
$ pip install matplotlib==3.9.2
$ pip install ollama>=0.1.0

This installs the Python client only. The Ollama runtime must be installed separately.

Local LLM Setup (Ollama)

The RAG pipeline uses Ollama for local language model inference. Install Ollama separately:

• Install Ollama: Visit ollama.ai and follow the installation instructions for your platform.
• Pull a model: Once Ollama is installed, download a model:

$ ollama pull llama3

• Verify installation:

$ ollama list

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!

Implementation Walkthrough

We’ll cover this in 3 parts:

• config.py: central configuration and prompt templates
• rag_utils.py: retrieval + LLM integration logic
• 03_rag_pipeline.py: driver script that ties everything together

Configuration (config.py)

The config.py module defines paths, constants, and templates that are used throughout the RAG pipeline. Think of it as the “control room” for your entire setup.

Directory and Path Setup

BASE_DIR = Path(__file__).resolve().parent.parent
DATA_DIR = BASE_DIR / "data"
INPUT_DIR = DATA_DIR / "input"
OUTPUT_DIR = DATA_DIR / "output"
INDEX_DIR = DATA_DIR / "indexes"
FIGURES_DIR = DATA_DIR / "figures"

Here, we define a consistent directory structure so that every script can find data, indexes, and output files, regardless of where it runs from.

This ensures reproducibility — a key trait for multi-script projects like this one.

Tip: Using Path(__file__).resolve().parent.parent automatically points to your project’s root directory, keeping all paths portable.

Corpus and Embedding Artifacts

CORPUS_PATH = INPUT_DIR / "corpus.txt"
CORPUS_META_PATH = INPUT_DIR / "corpus_metadata.json"
EMBEDDINGS_PATH = OUTPUT_DIR / "embeddings.npy"
METADATA_ALIGNED_PATH = OUTPUT_DIR / "metadata_aligned.json"
DIM_REDUCED_PATH = OUTPUT_DIR / "pca_2d.npy"

These paths represent:

• Corpus files: your input text and metadata
• Embedding artifacts: precomputed vectors and PCA-reduced coordinates for visualization

We also include environment variable overrides (i.e., CORPUS_PATH, CORPUS_META_PATH) to make it easy to point to new datasets without editing code.

Index Artifacts

FLAT_INDEX_PATH = INDEX_DIR / "faiss_flat.index"
HNSW_INDEX_PATH = INDEX_DIR / "faiss_hnsw.index"

These define storage for your Flat (exact) and HNSW (approximate) FAISS indexes.

They’re generated in Lesson 2 and reused here for retrieval.

Model and General Settings

EMBED_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
SEED = 42
DEFAULT_TOP_K = 5
SIM_THRESHOLD = 0.35

• Sentence Transformer model: the same compact model used for embedding queries and documents
• SEED: ensures deterministic sampling
• DEFAULT_TOP_K: number of chunks retrieved per question
• SIM_THRESHOLD: a similarity cut-off to filter weak matches

Prompt Templates for RAG

STRICT_SYSTEM_PROMPT = (
    "You are a concise assistant. Use ONLY the provided context."
    " If the answer is not contained verbatim or explicitly, say you do not know."
)
SYNTHESIZING_SYSTEM_PROMPT = (
    "You are a concise assistant. Rely ONLY on the provided context, but you MAY synthesize"
    " an answer by combining or paraphrasing the facts present. If the context truly lacks"
    " sufficient evidence, say you do not know instead of guessing."
)

The following 2 templates control LLM behavior:

• Strict mode: purely extractive, no paraphrasing
• Synthesizing mode: allows combining retrieved snippets to form explanatory answers

This distinction is critical when testing retrieval quality versus generation quality.

Intelligent Prompt Builder

def build_prompt(context_chunks, question: str, allow_synthesis: bool = False) -> str:
    system_prompt = SYNTHESIZING_SYSTEM_PROMPT if allow_synthesis else STRICT_SYSTEM_PROMPT
    context_str = "\n\n".join(context_chunks)
    return f"System: {system_prompt}\n{CONTEXT_HEADER}\n{context_str}\n\n" + USER_QUESTION_TEMPLATE.format(question=question)

This function assembles the final prompt fed into the LLM.

It concatenates retrieved context snippets, appends the system instructions, and ends with the user query.

Tip: The key here is flexibility — by toggling allow_synthesis, you can dynamically switch between closed-book and open-book answering styles.

Directory Bootstrap

for d in (OUTPUT_DIR, INDEX_DIR, FIGURES_DIR):
    d.mkdir(parents=True, exist_ok=True)

Ensures that all critical folders exist before any writing occurs — a small but essential safeguard for production stability (Figure 8).

At this point, the configuration module provides the foundation for the next step: actually retrieving and generating answers.

Integrating Ollama with FAISS Vector Search for RAG

Now that our FAISS index is ready to serve embeddings, the next step is to connect it with an LLM — the final reasoning layer that generates natural-language answers based on retrieved context.

The rag_utils.py file is where retrieval meets generation.

It ties together the embedding search results, builds prompts, calls the LLM (Ollama by default), and even adds explainability through citations and per-sentence support scoring.

Overview and Setup

Let’s start by looking at the top of the file:

import os, json, re, requests
import numpy as np
from typing import List, Dict, Tuple, Any

try:
    import ollama  # type: ignore
except ImportError:
    ollama = None

At the core, this script:

• Uses Ollama for local LLM inference, but gracefully falls back to HTTP if the Python client isn’t installed.
• Imports NumPy for fast vector math, requests for API calls, and typing hints for readability.

Then, it configures Ollama’s endpoints:

OLLAMA_BASE_URL = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434")
OLLAMA_API_URL = f"{OLLAMA_BASE_URL}/api/generate"
OLLAMA_TAGS_URL = f"{OLLAMA_BASE_URL}/api/tags"

Tip: You can override OLLAMA_BASE_URL with an environment variable — handy when deploying on remote servers or Docker containers (Figure 9).

Health Check and Model Discovery

Before we make any generation calls, it’s good practice to confirm that Ollama is actually reachable.

def ollama_available() -> bool:
    try:
        r = requests.get(OLLAMA_TAGS_URL, timeout=2)
        return r.status_code == 200
    except requests.RequestException:
        return False

If this returns False, your RAG pipeline will still work — it will simply skip generation or return a warning message.

Similarly, you can list all locally available models:

def list_ollama_models() -> List[str]:
    """Return a list of available local Ollama model names (empty if unreachable)."""
    resp = requests.get(OLLAMA_TAGS_URL, timeout=2)
    resp.raise_for_status()
    data = resp.json()
    models = []
    for m in data.get("models", []):
        name = m.get("name", "")
        if name.endswith(":latest"):
            name = name.rsplit(":", 1)[0]
        if name:
            models.append(name)
    return sorted(set(models))

This lets you dynamically query what’s installed (e.g., llama3, mistral, or gemma2).

If you’re running an interactive RAG app, this list can populate a dropdown for user selection.

Making the Ollama Call

Here’s the heart of your LLM connector:

def call_ollama(model: str, prompt: str, stream: bool = False) -> str:
    """Call Ollama using python client if available else raw HTTP."""
    if ollama is not None:
        try:
            if stream:
                out = []
                for chunk in ollama.generate(model=model, prompt=prompt, stream=True):
                    out.append(chunk.get("response", ""))
                return "".join(out)
            else:
                resp = ollama.generate(model=model, prompt=prompt)
                return resp.get("response", "")
        except Exception:
            pass

• If the ollama library is installed, the function uses its official Python client for better efficiency and streaming support.
• If not, it falls back to a manual HTTP request:

payload = {"model": model, "prompt": prompt, "stream": stream}
resp = requests.post(OLLAMA_API_URL, json=payload, timeout=120, stream=stream)

It even supports streaming tokens one by one — useful for building chat UIs or dashboards that display the answer as it’s generated.

Why this dual approach?

Not all environments (e.g., Docker containers or lightweight cloud runners) have the ollama Python package installed, but they can still access the REST (Representational State Transfer) API.

Optional: Cloud Fallback (OpenAI)

There’s a commented-out section providing an optional fallback to OpenAI’s API.

If uncommented, you can quickly switch between local and cloud models (e.g., gpt-4o-mini).

# OPENAI_MODEL = os.getenv("OPENAI_MODEL", "gpt-4o-mini")
# openai.api_key = os.getenv("OPENAI_API_KEY")
# def call_openai(prompt: str, model: str = OPENAI_MODEL) -> str:
#     ...

This flexibility lets you deploy the same RAG logic on-premises (Ollama) or in the cloud (OpenAI).

Selecting the Top-k Relevant Chunks

Once a user asks a question, we compute its embedding and retrieve similar text chunks:

def select_top_k(question_emb, embeddings, texts, metadata, k=5, sim_threshold=0.35):
    sims = embeddings @ question_emb  # cosine if normalized
    ranked = np.argsort(-sims)
    results = []
    for idx in ranked[:k * 2]:
        score = float(sims[idx])
        if score < sim_threshold and len(results) >= k:
            break
        results.append({
            "id": metadata[idx]["id"],
            "text": texts[idx],
            "score": score,
            "topic": metadata[idx].get("topic", "unknown")
        })
        if len(results) >= k:
            break
    return results

This function:

• Computes cosine similarities between the query and all embeddings.
• Ranks them, filters by a similarity threshold, and returns the top-k chunks with metadata.

This lightweight retrieval replaces the need to re-query FAISS every time — perfect for quick experiments or small datasets.

Splitting Answers into Sentences

Once the LLM produces an answer, we may want to analyze it sentence-by-sentence.

def _sentence_split(text: str) -> List[str]:
    raw = re.split(r'(?<=[.!?])\s+|\n+', text.strip())
    return [s.strip() for s in raw if s and not s.isspace()]

This regex-based approach avoids heavy NLP libraries and still performs well for most English prose.

Computing Sentence Support

A unique feature of this pipeline is its ability to score each sentence in the LLM’s answer by how well it aligns with the retrieved context chunks.

This helps determine which parts of the generated answer are actually supported by the retrieved evidence — forming the basis for citations such as [1], [2].

def _compute_support(sentences, retrieved, metadata, embeddings, model):
    id_to_idx = {m["id"]: i for i, m in enumerate(metadata)}
    chunk_vecs, ranks = [], []
    for rank, r in enumerate(retrieved, start=1):
        idx = id_to_idx.get(r["id"])
        if idx is None:
            continue
        chunk_vecs.append(embeddings[idx])
        ranks.append(rank)
    if not chunk_vecs:
        return [], sentences

    chunk_matrix = np.vstack(chunk_vecs)
    sent_embs = model.encode(sentences, normalize_embeddings=True, convert_to_numpy=True)

Each sentence is embedded and compared to the embeddings of the top-k retrieved chunks.

This yields 2 useful artifacts:

• support_rows: structured table of support scores
• cited_sentences: answer text annotated with citations such as [1], [2]

Example: Sentence-to-Context Alignment

For example, suppose the user asked:

“What is Streamlit used for?”

The retriever would return the top-k most relevant chunks for that query.

Each sentence in the model’s generated answer is then compared to the retrieved chunks to determine how well it is supported (Table 3).

Note: The context ranks come from the retrieval step based on the query “What is Streamlit used for?”. The similarity scores show how strongly each sentence aligns with those retrieved chunks — indicating how well each part of the generated answer is supported by evidence.

Formatting and Styling

To display results nicely, the _apply_style() helper supports different output styles:

def _apply_style(answer, style, cited_sentences):
    if style == "bullets" and cited_sentences:
        return "\n" + "\n".join(f"- {s}" for s in cited_sentences)
    return answer

This allows both paragraph and bullet-point summaries with inline citations — perfect for user-facing dashboards.

The Core: generate_rag_response()

Finally, the star of this file — the main RAG generation function:

def generate_rag_response(question, model, embeddings, texts, metadata,
                          llm_model_name="llama3", top_k=5,
                          allow_synthesis=False, force_strict=False,
                          add_citations=False, compute_support=False,
                          style="paragraph") -> Dict:

This function orchestrates the full retrieval-generation pipeline:

Step 1: Detect intent and embeddings

It embeds the question and automatically decides whether to allow synthesis:

if any(pat in q_lower for pat in config.AUTO_SYNTHESIS_PATTERNS):
    allow_synthesis = True
    heuristic_triggered = True

So if a query contains words like “why” or “benefits”, the model automatically switches to a paraphrasing mode instead of strict extraction.

Step 2: Retrieve top-k chunks

top = select_top_k(q_emb, embeddings, texts, metadata, k=top_k)
prompt = build_prompt([r["text"] for r in top], question, allow_synthesis=allow_synthesis)

Step 3: Generate via LLM

if not ollama_available():
    answer = "[Ollama not available at base URL.]"
else:
    answer = call_ollama(llm_model_name, prompt)

Step 4: Optional post-processing

If citations or support scoring are enabled:

sentences = _sentence_split(answer)
support_rows, cited_sentences = _compute_support(sentences, top, metadata, embeddings, model)
answer = _apply_style(answer, style, cited_sentences)

Finally, it returns a structured dictionary — containing everything from the retrieved context to the generated answer and support metrics.

Summary of the Utilities

The rag_utils.py file provides a robust and extensible RAG backbone:

• Local-first design: works seamlessly with Ollama or over HTTP
• Hybrid retrieval: embedding search + FAISS indexes
• Explainable outputs: sentence-level support and citations
• Prompt control: configurable synthesis vs. strict modes
• Output flexibility: paragraph or bullet styles, JSON export

Running a Local RAG Pipeline with Ollama and FAISS

Imports and Module Wiring

"""

Steps:
1. Load embeddings & indexes (or build fallbacks)
2. Accept user question(s)
3. Retrieve top-k relevant chunks
4. Construct prompt & call Ollama (fallback to placeholder if unavailable)
5. Display answer with retrieved context & scores
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
from rich import print
from rich.table import Table

from pyimagesearch import config
from pyimagesearch.embeddings_utils import load_embeddings, load_corpus, get_model, generate_embeddings
from pyimagesearch.vector_search_utils import build_flat_index, load_index, build_hnsw_index
from pyimagesearch.rag_utils import generate_rag_response, list_ollama_models, ollama_available

What this sets up:

• CLI (command line interface) flags (argparse), pretty terminal output (rich), NumPy for arrays.
• Pulls in config paths, embedding helpers, FAISS index builders and loaders, the RAG core (generate_rag_response), and Ollama helpers.

Ensure Embeddings (load or build once)

def ensure_embeddings(corpus_path=None, meta_path=None):
    if config.EMBEDDINGS_PATH.exists():
        emb, meta = load_embeddings()
        texts, _ = load_corpus(corpus_path or config.CORPUS_PATH, meta_path or config.CORPUS_META_PATH)
        return emb, meta, texts
    texts, meta = load_corpus(corpus_path or config.CORPUS_PATH, meta_path or config.CORPUS_META_PATH)
    model = get_model()
    emb = generate_embeddings(texts, model=model)
    from pyimagesearch.embeddings_utils import save_embeddings
    save_embeddings(emb, meta)
    return emb, meta, texts

What it does (and why):

• If data/output/embeddings.npy is present, it loads the embeddings and aligned metadata, then reads the current corpus to ensure your text list is up to date.
• If not present, it embeds the corpus with SentenceTransformer and caches both artifacts to disk for speed on re-runs.

Ensure Indexes (Flat must exist; HNSW is optional)

def ensure_indexes(embeddings):
    # Try load flat
    idx = None
    if config.FLAT_INDEX_PATH.exists():
        try:
            from pyimagesearch.vector_search_utils import load_index
            idx = load_index(config.FLAT_INDEX_PATH)
        except Exception:
            idx = None
    if idx is None:
        idx = build_flat_index(embeddings)
    # Optional: attempt HNSW
    hnsw = None
    if config.HNSW_INDEX_PATH.exists():
        try:
            hnsw = load_index(config.HNSW_INDEX_PATH)
        except Exception:
            hnsw = None
    else:
        try:
            hnsw = build_hnsw_index(embeddings)
        except Exception:
            hnsw = None
    return idx, hnsw

What it does (and why):

• Flat index (exact, inner product): Attempts to load from disk; if missing, builds from the embedding matrix. This guarantees you always have a correct baseline.
• HNSW (approximate, fast): Loads if available; otherwise builds the index. If FAISS isn’t installed with HNSW support, it fails gracefully and returns None.
• Returns: A tuple (flat, hnsw) for downstream use.

Interactive Q&A Loop — Optional Mode

def interactive_loop(model, embeddings, texts, metadata, llm_model: str, top_k: int, allow_synth: bool):
    print("[bold cyan]Enter questions (type 'exit' to quit).[/bold cyan]")
    while True:
        try:
            q = input("Question> ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\n[red]Exiting.[/red]")
            break
        if not q:
            continue
        if q.lower() in {"exit", "quit"}:
            break
        result = generate_rag_response(q, model, embeddings, texts, metadata, llm_model_name=llm_model, top_k=top_k, allow_synthesis=allow_synth)
        show_result(result)

What it does (and why):

• Lets you chat with your local RAG system.
• For each typed question, calls generate_rag_response(...) — retrieves context → builds the prompt → calls Ollama → formats the answer — and prints a rich table of the results.

Pretty Printing the Answer and Context (optional prompt/support)

def show_result(result, show_prompt: bool = False, show_support: bool = False):
    print("\n[bold green]Answer[/bold green]:")
    print(result["answer"].strip())
    synth_flag = "yes" if result.get("synthesis_used") else "no"
    if result.get("synthesis_used") and result.get("synthesis_heuristic"):
        print(f"[dim]Synthesis: {synth_flag} (auto-enabled by heuristic)\n[/dim]")
    else:
        print(f"[dim]Synthesis: {synth_flag}\n[/dim]")
    table = Table(title="Retrieved Context")
    table.add_column("Rank")
    table.add_column("ID")
    table.add_column("Score", justify="right")
    table.add_column("Snippet")
    for i, r in enumerate(result["retrieved"], start=1):
        snippet = r["text"][:80] + ("..." if len(r["text"]) > 80 else "")
        table.add_row(str(i), r["id"], f"{r['score']:.3f}", snippet)
    print(table)
    if show_prompt:
        print("[bold yellow]\n--- Prompt Sent to LLM ---[/bold yellow]")
        print(result.get("prompt", "[prompt missing]"))
    if show_support and result.get("support"):
        support_table = Table(title="Sentence Support Scores")
        support_table.add_column("Sentence")
        support_table.add_column("Rank")
        support_table.add_column("Score", justify="right")
        for row in result["support"]:
            support_table.add_row(row["sentence"], str(row["citation_rank"]), f"{row['support_score']:.3f}")
        print(support_table)

What it does (and why):

• Prints the final answer and indicates whether synthesis was used (including whether it was auto-enabled by the heuristic).
• Renders a Retrieved Context table showing rank, ID, similarity score, and a clean snippet.
• If --show-prompt is used, prints the full prompt for transparency.
• If --support-scores is enabled, shows per-sentence support strength against the retrieved chunks — useful for debugging groundedness.

CLI Entry Point (main) — flags, loading, answering

def main():
    parser = argparse.ArgumentParser(description="Minimal RAG pipeline demo")
    parser.add_argument("--llm-model", default="llama3", help="Ollama model name (must be pulled beforehand, e.g. 'ollama pull llama3')")
    parser.add_argument("--top-k", type=int, default=config.DEFAULT_TOP_K)
    parser.add_argument("--corpus-path", type=str, help="Override corpus file path")
    parser.add_argument("--corpus-meta-path", type=str, help="Override corpus metadata path")
    parser.add_argument("--question", type=str, help="Single question to answer (skip interactive mode)")
    parser.add_argument("--allow-synthesis", action="store_true", help="Permit model to synthesize answer by combining provided context facts")
    parser.add_argument("--list-models", action="store_true", help="List available local Ollama models and exit")
    parser.add_argument("--show-prompt", action="store_true", help="Display the full constructed prompt for debugging/teaching")
    parser.add_argument("--strict", action="store_true", help="Force strict extractive mode (disable synthesis even if heuristic matches)")
    parser.add_argument("--citations", action="store_true", help="Annotate sentences with citation indices")
    parser.add_argument("--style", choices=["paragraph", "bullets"], default="paragraph", help="Answer formatting style")
    parser.add_argument("--support-scores", action="store_true", help="Compute and display per-sentence support scores")
    parser.add_argument("--json", action="store_true", help="Output full result JSON to stdout (suppresses pretty tables except retrieved context)")
    args = parser.parse_args()
    if args.list_models:
        if not ollama_available():
            print("[red]Ollama not reachable at default base URL. Start Ollama to list models.[/red]")
            return
        models = list_ollama_models()
        if not models:
            print("[yellow]No models returned. Pull some with: ollama pull llama3[/yellow]")
        else:
            print("[bold cyan]Available Ollama models:[/bold cyan]")
            for m in models:
                print(f" - {m}")
        return
    print(f"[bold magenta]Using LLM model:[/bold magenta] {args.llm_model}")
    print("[bold magenta]Loading embeddings...[/bold magenta]")
    embeddings, metadata, texts = ensure_embeddings(corpus_path=args.corpus_path, meta_path=args.corpus_meta_path)
    model = get_model()
    print("[bold magenta]Preparing indexes (flat + optional hnsw)...[/bold magenta]")
    flat, hnsw = ensure_indexes(embeddings)
    # NOTE: We use embedding matrix directly for retrieval selection in rag_utils (cosine) for transparency.
    if args.question:
        result = generate_rag_response(
            args.question,
            model,
            embeddings,
            texts,
            metadata,
            llm_model_name=args.llm_model,
            top_k=args.top_k,
            allow_synthesis=args.allow_synthesis,
            force_strict=args.strict,
            add_citations=args.citations,
            compute_support=args.support_scores,
            style=args.style,
        )
        if args.json:
            import json as _json
            print(_json.dumps(result, indent=2))
        show_result(result, show_prompt=args.show_prompt, show_support=args.support_scores)
    else:
        # For interactive mode we keep previous behavior (could extend flags similarly if desired)
        interactive_loop(model, embeddings, texts, metadata, args.llm_model, args.top_k, args.allow_synthesis)

    print("[green]\nFinished RAG demo.\n[/green]")

What it does (and why):

• Defines a rich set of flags to control the model, retrieval depth, strictness vs. synthesis, prompt visibility, citations, style, and JSON output.
• --list-models lets you sanity-check your local Ollama setup without running the full pipeline.
• Loads or creates embeddings, prepares indexes, then either:
  • answers a single question (--question ...), or
  • launches the interactive loop.
• Optional JSON output is useful for scripting or automated tests.

Standard Python Entrypoint

if __name__ == "__main__":
    main()

What it does:

• Runs the CLI when you execute python 03_rag_pipeline.py.

Tiny Gotchas and Tips

• If FAISS was installed without HNSW support, ensure_indexes will still work — it just will not provide an HNSW index. The Flat index is always available.
• Make sure the Ollama model you request (e.g., llama3) is pulled first:

ollama pull llama3

• You can view exactly what the model saw with:

python 03_rag_pipeline.py --question "What is IVF indexing?" --show-prompt

• For teaching and debugging groundedness:

python 03_rag_pipeline.py --question "Why normalize embeddings?" --citations --support-scores

How to Run a Local RAG System with Ollama and FAISS

Now that everything’s wired up — embeddings, FAISS indexes, and the RAG utilities — it’s time to see the full pipeline in action.

You can start by verifying your local Ollama setup and ensuring the model (e.g., Llama 3) is pulled:

ollama pull llama3

Then, from your project root, launch the RAG pipeline:

python 03_rag_pipeline.py --question "What is FAISS?" --show-prompt --support-scores

If you’d rather chat interactively:

python 03_rag_pipeline.py

You’ll be greeted with a prompt like:

Question> Why do we normalize embeddings?

and can exit at any time with exit or Ctrl+C.

Example Output

Here’s what a typical run looks like inside your terminal (Figure 11).

What You Learned: Building a Production-Ready Local RAG System with Ollama and FAISS

By the end of this tutorial, you will have built and tested a complete, local Retrieval-Augmented Generation (RAG) system:

• Connected the FAISS vector store built in Lesson 2 to a local LLM served by Ollama.
• Used embeddings to retrieve semantically relevant chunks from your corpus.
• Constructed prompts dynamically and generated grounded answers, optionally including citations and synthesis.

This closes the loop of your vector → retrieval → generation workflow — forming the foundation for more advanced, production-ready RAG pipelines.

What's next? We recommend PyImageSearch University.

Course information:
86+ total classes • 115+ hours hours of on-demand code walkthrough videos • Last updated: April 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:

• ✓ 86+ courses on essential computer vision, deep learning, and OpenCV topics
• ✓ 86 Certificates of Completion
• ✓ 115+ hours 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 final lesson, you brought everything together (i.e., embeddings, vector search, and generation) to build a complete Retrieval-Augmented Generation (RAG) pipeline from scratch. You began by understanding how retrieval connects to language models, bridging the gap between semantic search and contextual reasoning.

Next, you explored how the system uses SentenceTransformer embeddings and FAISS indexes to fetch relevant context from a corpus before generating an answer. You then examined the RAG utilities in detail — from ollama_available() and call_ollama(), which handle model calls and fallbacks, to select_top_k(), which performs the crucial retrieval step by ranking and filtering results based on cosine similarity. You also saw how automatic synthesis heuristics determine when to allow the LLM to combine information creatively, adding flexibility to the pipeline.

Then came the driver script, where the theoretical pieces transformed into a working application. You walked through the full flow — loading embeddings, preparing indexes, retrieving the top-k most relevant chunks, and generating context-aware answers via Ollama. You also learned how to add citations, measure support scores, and switch between strict and synthesis modes for transparent reasoning.

Finally, you ran the pipeline locally, queried your own data, and observed meaningful, grounded responses generated by a local LLM. With this, you completed a true end-to-end workflow — from encoding and indexing knowledge to retrieving and generating answers — running fully offline and powered by FAISS and Ollama.

In short, you did not just learn RAG — you built it.

Citation Information

Singh, V. “Vector Search Using Ollama for Retrieval-Augmented Generation (RAG),” PyImageSearch, P. Chugh, S. Huot, A. Sharma, and P. Thakur, eds., 2026, https://pyimg.co/q68nv

@incollection{Singh_2026_vector-search-using-ollama-for-rag,
  author = {Vikram Singh},
  title = {{Vector Search Using Ollama for Retrieval-Augmented Generation (RAG)}},
  booktitle = {PyImageSearch},
  editor = {Puneet Chugh and Susan Huot and Aditya Sharma and Piyush Thakur},
  year = {2026},
  url = {https://pyimg.co/q68nv},
}

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 Vector Search Using Ollama for Retrieval-Augmented Generation (RAG) appeared first on PyImageSearch.

Vector Search with FAISS: Approximate Nearest Neighbor (ANN) Explained

In this tutorial, you’ll learn how vector databases achieve lightning-fast retrieval using Approximate Nearest Neighbor (ANN) algorithms. You’ll explore how FAISS structures your embeddings into efficient indexes (e.g., Flat, IVF, and HNSW), benchmark their performance, and see how recall and latency trade off as your dataset scales.

This lesson is the 2nd of a 3-part series on Retrieval-Augmented Generation (RAG):

1. TF-IDF vs. Embeddings: From Keywords to Semantic Search
2. Vector Search with FAISS: Approximate Nearest Neighbor (ANN) Explained (this tutorial)
3. Lesson 3

To learn how to scale semantic search from milliseconds to millions of vectors, just keep reading.

Looking for the source code to this post?

Jump Right To The Downloads Section

From Exact to Approximate: Why Indexing Matters

In the previous lesson, you learned how to turn text into embeddings — compact, high-dimensional vectors that capture semantic meaning.

By computing cosine similarity between these vectors, you could find which sentences or paragraphs were most alike.

That worked beautifully for a small handcrafted corpus of 30-40 paragraphs.

But what if your dataset grows to millions of documents or billions of image embeddings?

Suddenly, your brute-force search breaks down — and that’s where Approximate Nearest Neighbor (ANN) methods come to the rescue.

The Trouble with Brute-Force Search

In Lesson 1, we compared every query vector against all stored embeddings using cosine similarity:

If your dataset contains vectors, this means dot products for each query — an operation that scales linearly as .

That’s fine for 500 vectors, but catastrophic for 50 million.

Let’s do a quick back-of-the-envelope estimate:

Even on modern hardware capable of tens of gigaflops (GFLOPS) per second, brute-force search over billions of vectors still incurs multi-second latency per query — before accounting for memory bandwidth limits — making indexing essential for production-scale retrieval.

The Curse of Dimensionality

High-dimensional data introduces another subtle problem: distances begin to blur.

In 384-dimensional space, most points are almost equally distant from each other.

This phenomenon, called the curse of dimensionality, makes “nearest neighbor” less meaningful unless your metric is very carefully chosen and your embeddings are normalized.

🧠 Intuition: imagine scattering points on a 2D plane — you can easily tell which ones are close.

In 384D, everything is “far away,” so naive distance computations lose discriminative power.

That’s why normalization (L2) is critical — it keeps embeddings on a unit hypersphere, letting cosine similarity focus purely on angle, not length.

But even with normalization, you still can’t escape the sheer computational load of checking every vector.

Enter the Approximate Nearest Neighbor (ANN)

So, how do we search faster without comparing against every vector?

Approximate Nearest Neighbor (ANN) algorithms do exactly that.

They trade a small amount of accuracy (recall) for huge latency gains — often 100× faster.

Conceptually, they build an index that helps jump directly to promising regions of vector space.

Think of it like finding a book in a massive library:

• Brute-force search: opening every book until you find the right topic.
• Indexed search: using the catalog — you might not find the exact match every time, but you’ll get very close, much faster.

Accuracy vs. Latency: The Core Trade-Off

ANN methods exist in many flavors — graph-based (HNSW), cluster-based (IVF), tree-based (LSH, KD-Tree), and hybrid approaches — but they all share a single philosophy:

They do so by:

• Pre-structuring the vector space (clustering, graphs, hashing).
• Restricting comparisons to only nearby candidates.
• Accepting that some true neighbors might be missed — in exchange for massive speed-ups.

You’ll often see performance expressed using recall@k, the fraction of true neighbors recovered among the top-k results.

A good ANN index achieves 0.9 – 0.99 recall@k while being orders of magnitude faster than brute force.

The Flat index delivers exact results but with high latency, whereas HNSW and IVF-Flat balance speed and accuracy for scalable retrieval.

What You’ll Learn Next

In the next section, we’ll open the black box and look inside FAISS — the industry-standard library that powers vector search at scale.

You’ll understand how these indexes (i.e., Flat, IVF, and HNSW) are built, what parameters control their trade-offs, and how to construct them yourself in code.

By the end, you’ll see how these algorithms let you scale semantic search from a few dozen vectors to millions — without loss of meaning.

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.

Inside FAISS: How Vector Indexing Works

In Lesson 1, we learned how to generate embeddings and compute semantic similarity.

In Section 1 of this post, we discovered why brute-force search breaks down as our vector collection grows.

Now, let’s open the black box and explore how FAISS makes large-scale vector search practical.

What Is FAISS?

FAISS (Facebook AI Similarity Search) is an open-source library developed by Meta AI for efficient similarity search and clustering of dense vectors.

It’s implemented in highly optimized C++ with Python bindings and optional GPU support.

In simple terms, FAISS is the NumPy of vector search — it provides a consistent interface for building, training, querying, and persisting vector indexes, with dozens of back-ends optimized for both accuracy and speed.

Why We Need Index Structures

Every FAISS index is a data structure that helps you locate nearest neighbors efficiently in a large vector space.

Instead of storing raw vectors in a flat array and scanning them linearly, FAISS builds one of several possible structures that act as “shortcuts” through the space.

Let’s explore the three we will implement in this tutorial — Flat, IVF-Flat, and HNSW — and see how each balances speed and accuracy.

Flat Index: The Exact Baseline

The Flat index is the simplest and most precise approach.

It stores every embedding as is and computes cosine similarity (or inner product) between the query and each vector.

This is essentially our brute-force baseline from Lesson 1.

It achieves perfect recall (1.0) but scales poorly — every query requires scanning all vectors.

IVF-Flat: Coarse Clustering for Speed

IVF stands for Inverted File Index.

The idea is simple but powerful: partition the vector space into coarse clusters and search only a few of them at query time.

Step-by-step:

• FAISS trains a coarse quantizer using k-means on your embeddings. Each cluster center is a “list” in the index — there are nlist of them.
• Each vector is assigned to its closest centroid.
• During search, instead of scanning all lists, FAISS only probes nprobe of them (the most promising clusters).

That reduces search complexity from O(N) to roughly O(N / nlist × nprobe).

HNSW: Navigable Graph Search

The Hierarchical Navigable Small World (HNSW) index uses a graph instead of clusters.

Each embedding is a node connected to its nearest neighbors.

At query time, the algorithm navigates the graph — starting from a random entry point and greedily moving to closer nodes until it reaches a local minimum.

This search is then refined by probing nearby connections to find better candidates.

Intuition: It’s like finding a friend’s house by asking neighbors for directions instead of checking every street yourself.

Comparing the Three Index Types

What You’ll Do Next

Now that you understand what each index does, you’ll see how to implement them with FAISS.

In the next section, we’ll dive into the code — explaining config.py, vector_search_utils.py, then building and benchmarking indexes side by side.

Configuring Your Development Environment

To follow this guide, you’ll need to install several Python libraries for vector search and approximate nearest neighbor (ANN) indexing.

The core dependencies are:

$ pip install sentence-transformers==2.7.0
$ pip install faiss-cpu==1.8.0
$ pip install numpy==1.26.4
$ pip install rich==13.8.1

Note on FAISS Installation

• Use faiss-cpu for CPU-only environments (most common)
• If you have a CUDA-compatible GPU, you can install faiss-gpu instead for better performance with large datasets

Optional Dependencies

If you plan to extend the examples with your own experiments:

$ pip install matplotlib  # For creating performance visualizations

Verifying Your Installation

You can verify FAISS is properly installed by running:

import faiss
print(f"FAISS version: {faiss.__version__}")

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!

Implementation Walkthrough

Alright, time to get our hands dirty.

You’ve seen why indexing matters and how FAISS structures (e.g., Flat, HNSW, and IVF) work conceptually; now we’ll implement them and run them side-by-side. We’ll build each index, query it with the same vectors, and compare recall@k, average query latency, build time, and a rough memory footprint.

To keep things clean, we’ll proceed in 3 parts:

• config: paths and constants
• FAISS utilities: builders, query/save/load, benchmarking helpers
• driver script: ties it all together and prints a clear, comparable report

pyimagesearch/config.py (what we use in Lesson 2)

These constants keep paths, filenames, and defaults in one place, so the rest of the code stays clean and portable.

from pathlib import Path
import os

# Base paths
BASE_DIR = Path(__file__).resolve().parent.parent
DATA_DIR = BASE_DIR / "data"
INPUT_DIR = DATA_DIR / "input"
OUTPUT_DIR = DATA_DIR / "output"
INDEX_DIR = DATA_DIR / "indexes"
FIGURES_DIR = DATA_DIR / "figures"

# Corpus files (allow environment overrides for flexibility)
_CORPUS_OVERRIDE = os.getenv("CORPUS_PATH")
_CORPUS_META_OVERRIDE = os.getenv("CORPUS_META_PATH")
CORPUS_PATH = Path(_CORPUS_OVERRIDE) if _CORPUS_OVERRIDE else INPUT_DIR / "corpus.txt"
CORPUS_META_PATH = Path(_CORPUS_META_OVERRIDE) if _CORPUS_META_OVERRIDE else INPUT_DIR / "corpus_metadata.json"

# Embedding artifacts
EMBEDDINGS_PATH = OUTPUT_DIR / "embeddings.npy"
METADATA_ALIGNED_PATH = OUTPUT_DIR / "metadata_aligned.json"
DIM_REDUCED_PATH = OUTPUT_DIR / "pca_2d.npy"

# Index artifacts
FLAT_INDEX_PATH = INDEX_DIR / "faiss_flat.index"
HNSW_INDEX_PATH = INDEX_DIR / "faiss_hnsw.index"

# Models & defaults
EMBED_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
SEED = 42
DEFAULT_TOP_K = 5
SIM_THRESHOLD = 0.35  # not used in this lesson, but fine to keep

# Ensure dirs exist
for d in (OUTPUT_DIR, INDEX_DIR, FIGURES_DIR):
    d.mkdir(parents=True, exist_ok=True)

What’s happening here (and why):

• We compute BASE_DIR and derive data/ subfolders so the repo runs the same on any machine.
• CORPUS_PATH and CORPUS_META_PATH can be overridden via environment variables — a nice quality-of-life touch if you swap datasets.
• EMBEDDINGS_PATH is the cache produced in Lesson 1; we load it here to avoid re-embedding.
• FLAT_INDEX_PATH and HNSW_INDEX_PATH are where we’ll persist FAISS indexes after benchmarking. (IVF can be added similarly if you want on-disk reuse.)
• The directory loop at the bottom makes sure first runs never fail due to missing folders.

Note: The file also contains prompt-related constants for Lesson 3 (RAG). We don’t use them here — safe to ignore in this lesson.

pyimagesearch/vector_search_utils.py

This module is your FAISS “Swiss Army Knife.” We’ll go through every function you shared, including the ones you asked me not to miss: brute_force_search, measure_latency, and estimate_index_memory_bytes, plus query/save/load.

a) FAISS Import Guard

try:
    import faiss  # type: ignore
except ImportError:
    faiss = None  # type: ignore

def require_faiss():
    if faiss is None:
        raise ImportError("faiss not installed. Please pip install faiss-cpu (or faiss-gpu).")

This is a friendly safety check. Any time we build or query an index, we call require_faiss(). If FAISS isn’t installed, you get a clear error with the exact package to install.

b) Flat (exact) Index

def build_flat_index(embeddings: np.ndarray):
    require_faiss()
    d = embeddings.shape[1]
    index = faiss.IndexFlatIP(d)
    index.add(embeddings.astype(np.float32))
    return index

What it does:

We create an IndexFlatIP, where IP = inner product. Because our embeddings are L2-normalized (from Lesson 1), inner product equals cosine similarity. We then .add() all vectors (as float32) to the index. This gives perfect recall but scans everything (slow at scale). This equivalence holds only if both stored embeddings and query vectors are L2-normalized. If not, inner product will not equal cosine similarity.

c) HNSW (graph) Index

def build_hnsw_index(embeddings: np.ndarray, m: int = 32, ef_construction: int = 128, ef_search: int = 64):
    require_faiss()
    d = embeddings.shape[1]
    index = faiss.IndexHNSWFlat(d, m, faiss.METRIC_INNER_PRODUCT)
    index.hnsw.efConstruction = ef_construction
    index.hnsw.efSearch = ef_search
    index.add(embeddings.astype(np.float32))
    return index

What it does:

HNSW builds a multi-layer graph over your vectors:

• m: sets connectivity per node (more links → better recall and higher memory)
• efConstruction: controls how widely we explore during graph building
• efSearch: controls breadth at query time (higher → better recall, slower)

We add vectors and return a ready-to-query graph index.

d) IVF-Flat (clustered) Index

def build_ivf_flat_index(embeddings: np.ndarray, nlist: int = 8, nprobe: int = 4):
    require_faiss()
    d = embeddings.shape[1]
    quantizer = faiss.IndexFlatIP(d)
    index = faiss.IndexIVFFlat(quantizer, d, nlist, faiss.METRIC_INNER_PRODUCT)
    if not index.is_trained:
        index.train(embeddings.astype(np.float32))
    index.add(embeddings.astype(np.float32))
    index.nprobe = nprobe
    return index

What it does:

IVF partitions your space into nlist coarse clusters (centroids via k-means). You must train first, then add vectors. Queries only probe a handful of lists (nprobe) near the query, drastically reducing the number of distance comparisons.

• Tune nlist up as your dataset grows (more clusters).
• Tune nprobe for recall vs. latency.

e) Unified Query Wrapper

def query_index(index, query_vectors: np.ndarray, top_k: int = config.DEFAULT_TOP_K) -> Tuple[np.ndarray, np.ndarray]:
    distances, indices = index.search(query_vectors.astype(np.float32), top_k)
    return indices, distances

What it does:

FAISS returns (distances, indices), but most of our code wants (indices, distances). This wrapper normalizes that behavior and ensures the input dtype is float32. Use this anywhere you want a consistent top-k interface.

f) Save and Load Indexes

def save_index(index, path: Path):
    require_faiss()
    faiss.write_index(index, str(path))

def load_index(path: Path):
    require_faiss()
    return faiss.read_index(str(path))

What it does:

Build once, reuse many times. These functions serialize an index to disk and read it back instantly — essential when your dataset is large or when you run repeated experiments.

g) Brute-Force Baseline (ground truth)

def brute_force_search(embeddings: np.ndarray, queries: np.ndarray, top_k: int) -> Tuple[np.ndarray, np.ndarray]:
    sims = queries @ embeddings.T  # cosine if normalized
    top_indices = np.argsort(-sims, axis=1)[:, :top_k]
    top_scores = np.take_along_axis(sims, top_indices, axis=1)
    return top_indices, top_scores

What it does:

This computes exact cosine similarities (via matrix multiplication) and sorts them to get the true top-k. We use this to measure recall@k of ANN indexes.

h) Latency measurement helper

def measure_latency(func, *args, repeat: int = 3, **kwargs) -> Dict[str, Any]:
    import time
    times = []
    for _ in range(repeat):
        start = time.perf_counter()
        func(*args, **kwargs)
        times.append(time.perf_counter() - start)
    return {"avg_s": float(np.mean(times)), "stdev_s": float(np.std(times)), "runs": times}

What it does:

This runs the same query callable a few times and returns the average and standard deviation. This gives stable latency numbers for fair comparisons between Flat, HNSW, and IVF.

i) Rough Memory Estimator (educational)

def estimate_index_memory_bytes(index_type: str, n: int, d: int, m: int = 32, nlist: int = 0) -> int:
    base = n * d * 4
    if index_type == "Flat":
        return base
    if index_type == "HNSW":
        return base + n * m * 4
    if index_type == "IVF-Flat":
        return base + nlist * d * 4
    return base

What it does:

Provides a lower-bound estimate (in bytes) for apples-to-apples reporting.

• Flat: just the raw vectors (float32 = 4 bytes).
• HNSW: vectors + approximate link list.
• IVF-Flat: vectors + centroid matrix.

FAISS has internal overhead, so treat this as a teaching aid, not an exact profiler.

02_vector_search_ann.py (driver)

This script glues everything together:

• loads embeddings
• builds indexes
• benchmarks recall and latency
• displays results
• persists indexes

a) Ensuring Embeddings Exist

def ensure_embeddings():
    if config.EMBEDDINGS_PATH.exists():
        emb, meta = load_embeddings()
        texts, _ = load_corpus()
        return emb, meta, texts
    texts, meta = load_corpus()
    model = get_model()
    emb = generate_embeddings(texts, model=model)
    from pyimagesearch.embeddings_utils import save_embeddings
    save_embeddings(emb, meta)
    return emb, meta, texts

First run? We load corpus → embed → save. Subsequent runs? We simply load embeddings.npy. This keeps the ANN lesson focused on indexing rather than recomputing vectors.

b) Query Sampling and Recall

def sample_query_indices(n: int, total: int, seed: int = config.SEED) -> np.ndarray:
    rng = np.random.default_rng(seed)
    return rng.choice(total, size=min(n, total), replace=False)

def compute_recall(brute_force_idx: np.ndarray, ann_idx: np.ndarray) -> float:
    matches = 0
    total = brute_force_idx.shape[0] * brute_force_idx.shape[1]
    for row_true, row_test in zip(brute_force_idx, ann_idx):
        matches += len(set(row_true.tolist()) & set(row_test.tolist()))
    return matches / total if total else 0.0

We pick some random query rows from the embedding matrix and compare ANN results against the ground truth to compute recall@k.

c) Benchmark Core

def benchmark(
    embeddings: np.ndarray,
    queries: np.ndarray,
    k: int = config.DEFAULT_TOP_K,
    hnsw_m: int = 32,
    hnsw_ef_search: int = 64,
    ivf_nlist: int = 8,
    ivf_nprobe: int = 4,
    auto_adjust_ivf: bool = True,
):
    # 1) Ground truth
    bf_idx, bf_scores = brute_force_search(embeddings, queries, k)
    results = []

    # 2) Flat
    build_start = time.perf_counter()
    flat_index = build_flat_index(embeddings)
    build_time_flat = time.perf_counter() - build_start
    flat_index.search(queries[:1].astype(np.float32), k)  # warm-up
    latency_flat = measure_latency(flat_index.search, queries.astype(np.float32), k)
    flat_dist, flat_idx_res = flat_index.search(queries.astype(np.float32), k)
    recall_flat = compute_recall(bf_idx, flat_idx_res)
    if recall_flat < 0.99:
        print(f"[red]Warning: Flat recall {recall_flat:.3f} < 0.99. Check normalization or brute force implementation.[/red]")
    results.append(("Flat", recall_flat, latency_flat["avg_s"], build_time_flat,
                    estimate_index_memory_bytes("Flat", embeddings.shape[0], embeddings.shape[1])))

    # 3) HNSW
    try:
        build_start = time.perf_counter()
        hnsw_index = build_hnsw_index(embeddings, m=hnsw_m, ef_search=hnsw_ef_search)
        build_time_hnsw = time.perf_counter() - build_start
        hnsw_index.search(queries[:1].astype(np.float32), k)
        latency_hnsw = measure_latency(hnsw_index.search, queries.astype(np.float32), k)
        hnsw_dist, hnsw_idx_res = hnsw_index.search(queries.astype(np.float32), k)
        recall_hnsw = compute_recall(bf_idx, hnsw_idx_res)
        results.append(("HNSW", recall_hnsw, latency_hnsw["avg_s"], build_time_hnsw,
                        estimate_index_memory_bytes("HNSW", embeddings.shape[0], embeddings.shape[1], m=hnsw_m)))
    except Exception as e:
        results.append(("HNSW (unavailable)", 0.0, 0.0, 0.0, 0))
        print(f"[yellow]HNSW build failed: {e}[/yellow]")

    # 4) IVF-Flat (+ tiny-corpus auto-tuning)
    try:
        effective_nlist = ivf_nlist
        if auto_adjust_ivf:
            N = embeddings.shape[0]
            if N < ivf_nlist * 5:
                shrunk = max(2, min(ivf_nlist, max(2, N // 2)))
                if shrunk != ivf_nlist:
                    print(f"[yellow]Adjusting nlist from {ivf_nlist} to {shrunk} for tiny corpus (N={N}).[/yellow]")
                    effective_nlist = shrunk
        build_start = time.perf_counter()
        ivf_index = build_ivf_flat_index(embeddings, nlist=effective_nlist, nprobe=ivf_nprobe)
        build_time_ivf = time.perf_counter() - build_start
        ivf_index.search(queries[:1].astype(np.float32), k)
        latency_ivf = measure_latency(ivf_index.search, queries.astype(np.float32), k)
        ivf_dist, ivf_idx_res = ivf_index.search(queries.astype(np.float32), k)
        recall_ivf = compute_recall(bf_idx, ivf_idx_res)
        results.append(("IVF-Flat", recall_ivf, latency_ivf["avg_s"], build_time_ivf,
                        estimate_index_memory_bytes("IVF-Flat", embeddings.shape[0], embeddings.shape[1], nlist=effective_nlist)))
    except Exception as e:
        results.append(("IVF-Flat (failed)", 0.0, 0.0, 0.0, 0))
        print(f"[yellow]IVF build failed: {e}[/yellow]")

    return bf_idx, results

The pipeline is:

• create a ground-truth baseline
• build each index
• warm up
• measure latency
• compute recall vs. the ground-truth baseline
• record a row of metrics

The IVF block also includes a tiny-corpus auto-adjust so you don’t “over-partition” small datasets.

d) Results Table

def display_results(results):
    table = Table(title="ANN Benchmark - Recall / Query(ms) / Build(ms) / Mem(KB)")
    table.add_column("Index")
    table.add_column("Recall", justify="right")
    table.add_column("Query (ms)", justify="right")
    table.add_column("Build (ms)", justify="right")
    table.add_column("Mem (KB)", justify="right")
    for name, recall, q_latency, build_time, mem_bytes in results:
        table.add_row(
            name,
            f"{recall:.3f}",
            f"{q_latency*1000:.2f}",
            f"{build_time*1000:.1f}",
            f"{mem_bytes/1024:.1f}",
        )
    print(table)

This produces the clean terminal table you’ve seen earlier — great for screenshots.

e) Persisting Indexes

def persist_indexes(embeddings: np.ndarray):
    from pyimagesearch.vector_search_utils import save_index, build_flat_index
    flat = build_flat_index(embeddings)
    save_index(flat, config.FLAT_INDEX_PATH)
    try:
        from pyimagesearch.vector_search_utils import build_hnsw_index
        hnsw = build_hnsw_index(embeddings)
        save_index(hnsw, config.HNSW_INDEX_PATH)
    except Exception:
        pass

We persist Flat and HNSW by default, so Lesson 3 can load them directly. (Add IVF persistence if you want — same pattern.)

f) CLI Entrypoint

def main():
    parser = argparse.ArgumentParser(description="ANN benchmark demo")
    parser.add_argument("--k", type=int, default=5)
    parser.add_argument("--queries", type=int, default=5)
    parser.add_argument("--hnsw-m", type=int, default=32)
    parser.add_argument("--hnsw-ef-search", type=int, default=64)
    parser.add_argument("--ivf-nlist", type=int, default=8)
    parser.add_argument("--ivf-nprobe", type=int, default=4)
    parser.add_argument("--no-auto-adjust-ivf", action="store_true")
    args = parser.parse_args()

    print("[bold magenta]Loading embeddings...[/bold magenta]")
    embeddings, metadata, texts = ensure_embeddings()

    q_idx = sample_query_indices(args.queries, embeddings.shape[0])
    queries = embeddings[q_idx]

    print("[bold magenta]Benchmarking indexes...[/bold magenta]")
    bf_idx, results = benchmark(
        embeddings,
        queries,
        k=args.k,
        hnsw_m=args.hnsw_m,
        hnsw_ef_search=args.hnsw_ef_search,
        ivf_nlist=args.ivf_nlist,
        ivf_nprobe=args.ivf_nprobe,
        auto_adjust_ivf=not args.no_auto_adjust_ivf,
    )
    display_results(results)

    print("[bold magenta]Persisting indexes...[/bold magenta]")
    persist_indexes(embeddings)

    print("[green]\nDone. Proceed to Post 3 for RAG integration.\n[/green]")

Run it like this:

python 02_vector_search_ann.py --queries 10 --k 5 --ivf-nlist 32 --ivf-nprobe 8 --hnsw-m 32 --hnsw-ef-search 128

Terminal Output (example)

Benchmarking and Analyzing Results

After running the script with:

python 02_vector_search_ann.py --queries 10 --k 5 --ivf-nlist 32 --ivf-nprobe 8 --hnsw-m 32 --hnsw-ef-search 128

FAISS prints a benchmark table comparing Flat, HNSW, and IVF-Flat across recall, query latency, build time, and memory footprint.

How to Interpret These Results

Let’s break down what this means.

• Flat Index:
  This performs an exact nearest-neighbor search. Every query compares against all stored embeddings. It gives perfect recall (1.0) but becomes slow at scale because it’s linear in the number of vectors.
• HNSW (Hierarchical Navigable Small World):
  A graph-based index that builds layered shortcuts between points. It maintains near-perfect recall while cutting down query latency by 10-50× on large datasets.
  Here, because our dataset is tiny (41 vectors), the difference isn’t dramatic — but with 1M+ vectors, HNSW scales gracefully.
• IVF-Flat (Inverted File Index):
  This structure first clusters vectors into “lists” (centroids), then searches only a few clusters (nprobe).
  You can see a small recall drop (0.98) because some true neighbors fall outside probed clusters — but query speed improves noticeably.
  The build time is slightly higher since training the quantizer (k-means step) adds overhead.

Takeaways

• For small datasets, Flat is perfectly fine — it’s fast and exact.
• For large-scale retrieval (e.g., thousands of documents), HNSW offers the best accuracy — latency balance.
• For ultra-large datasets with millions of vectors, IVF-Flat or IVF-PQ can drastically cut memory and latency with minimal recall loss.
• Always benchmark with your actual data — FAISS lets you trade off accuracy and performance based on your use case.

What's next? We recommend PyImageSearch University.

Course information:
86+ total classes • 115+ hours hours of on-demand code walkthrough videos • Last updated: April 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:

• ✓ 86+ courses on essential computer vision, deep learning, and OpenCV topics
• ✓ 86 Certificates of Completion
• ✓ 115+ hours 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 tutorial, you moved beyond understanding embeddings and stepped into the practical world of vector search. You began by revisiting the limitations of brute-force similarity search and saw why indexing is crucial when datasets grow larger. Then, you explored how FAISS, Meta’s open-source library for efficient similarity search, builds the backbone of modern vector databases.

From there, you implemented and compared 3 major index types (i.e., Flat, HNSW, and IVF-Flat) each offering a different trade-off between speed, recall, and memory. You wrote code to benchmark their performance, measuring recall@k, query latency, build time, and memory usage, and visualized the results directly in your terminal. Along the way, you saw how HNSW maintains near-perfect accuracy while drastically improving query speed, and how IVF-Flat compresses the search space to achieve sub-millisecond responses even with slight precision loss.

By the end, you not only understood how approximate nearest neighbor (ANN) search works but also saw why it matters — enabling scalable, real-time retrieval without compromising too much on accuracy. You also saved your indexes, preparing them for use in the next lesson, where we’ll connect them to a Retrieval-Augmented Generation (RAG) pipeline and bring everything full circle: embeddings, vector search, and large language models working together in action.

Citation Information

Singh, V. “Vector Search with FAISS: Approximate Nearest Neighbor (ANN) Explained,” PyImageSearch, P. Chugh, S. Huot, A. Sharma, and P. Thakur, eds., 2026, https://pyimg.co/htl5f

@incollection{Singh_2026_vector-search-with-faiss-ann-explained,
  author = {Vikram Singh},
  title = {{Vector Search with FAISS: Approximate Nearest Neighbor (ANN) Explained}},
  booktitle = {PyImageSearch},
  editor = {Puneet Chugh and Susan Huot and Aditya Sharma and Piyush Thakur},
  year = {2026},
  url = {https://pyimg.co/htl5f},
}

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 Vector Search with FAISS: Approximate Nearest Neighbor (ANN) Explained appeared first on PyImageSearch.