Pytest Tutorial: MLOps Testing, Fixtures, and Locust Load Testing

In this lesson, you will learn how to make ML systems reliable, correct, and production-ready through structured testing and validation. You will walk through unit tests, integration tests, load and performance checks, fixtures, code quality tools, and automated test runs, giving you everything you need to ensure your ML API behaves predictably under real-world conditions.

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

1. FastAPI for MLOps: Python Project Structure and API Best Practices
2. Pytest Tutorial: MLOps Testing, Fixtures, and Locust Load Testing (this tutorial)

To learn how to test, validate, and stress-test your ML services like a professional MLOps engineer, just keep reading.

Looking for the source code to this post?

Jump Right To The Downloads Section

Introduction to MLOps Testing: Building Reliable ML Systems with Pytest

Testing is the backbone of reliable MLOps. A model might look great in a notebook, but once wrapped in services, APIs, configs, and infrastructure, dozens of things can break silently: incorrect inputs, unexpected model outputs, missing environment variables, slow endpoints, and downstream failures. This lesson ensures you never ship those problems into production.

In this lesson, you will learn the complete testing workflow for machine learning (ML) systems: from small, isolated unit tests to full API integration checks and load testing your endpoints under real traffic conditions. You will also understand how to structure your tests, how each type of test fits into the MLOps lifecycle, and how to design a test suite that grows cleanly as your project evolves.

To learn how to validate, benchmark, and harden your ML applications for production, just keep reading.

Why Testing Is Non-Negotiable in MLOps

Machine learning adds layers of unpredictability on top of regular software engineering. Models drift, inputs vary, inference latency can increase, and small code changes can ripple into major behavioral shifts. Without testing, you have no safety net. Proper tests make your system observable, predictable, and safe to deploy.

What You Will Learn: Pytest, Fixtures, and Load Testing for MLOps

You will walk through a practical testing workflow tailored for ML applications: writing unit tests for inference logic, validating API endpoints end-to-end, using fixtures to isolate environments, verifying configuration behavior, and running load tests to understand real-world performance. Each example connects directly to the codebase you built earlier.

From FastAPI to Testing: Extending Your MLOps Pipeline with Validation

Previously, you learned how to structure a clean ML codebase, configure environments, separate services, and expose reliable API endpoints. Now, you will stress-test that foundation. This lesson transforms your structured application into a validated, production-ready system with tests that catch issues before users ever see them.

Test-Driven MLOps: Applying Software Testing Best Practices to ML Pipelines

Test-driven development (TDD) matters even more in ML because models introduce uncertainty on top of normal software complexity. A single mistake in preprocessing, an incorrect model version, or a slow endpoint can break your application in ways that are hard to detect without a structured testing strategy. Test-driven MLOps gives you a predictable workflow: write tests, run them often, and let failures guide improvements.

What to Test in MLOps Pipelines: Models, APIs, and Configurations

ML systems require testing across multiple layers because issues can appear anywhere: in preprocessing logic, service code, configuration loading, API endpoints, or the model itself. You should verify that your inference service behaves correctly with both valid and invalid inputs, that your API returns consistent responses, that your configuration behaves as expected, and that the entire pipeline works end-to-end. Even when using a dummy model, testing ensures that the structure of your system remains correct as the real model is swapped in later.

Unit vs Integration vs Performance Testing

Unit tests focus on the smallest pieces of your system: functions, helper modules, and the inference service. They run fast and break quickly when a small change introduces an error. Integration tests validate how components work together: routes, services, configs, and the FastAPI layer. They ensure your API behaves consistently no matter what changes inside the codebase. Performance tests simulate real user traffic, evaluating latency, throughput, and failure rates under load. Together, these 3 types of tests create full confidence in your ML application.

The Software Testing Pyramid for MLOps: Unit, Integration, and Load Testing

The testing pyramid helps prioritize effort: many unit tests at the bottom, fewer integration tests in the middle, and a small number of heavy performance tests at the top. ML systems especially benefit from this structure because most failures occur in smaller utilities and service functions, not in the final API layer. By weighting your test suite correctly, you get fast feedback during development while still validating the entire system before deployment.

Project Structure and Test Layout

A clean testing layout makes your ML system predictable, scalable, and easy to maintain. By separating tests into clear categories (e.g., unit, integration, and performance), you ensure that each kind of test has a focused purpose and a natural home inside the repository. This structure also mirrors how real production MLOps teams organize their work, making your project easier to extend as your system grows.

Test Directory Structure for MLOps: unit, integration, and performance

Your Lesson 2 repository includes a dedicated tests/ directory with 3 subfolders:

tests/
│── unit/
│── integration/
└── performance/

• unit/: holds small, fast tests that validate individual pieces such as the DummyModel, the inference service, or helper functions.
• integration/: contains tests that spin up the FastAPI app and verify endpoints like /health, /predict, and the OpenAPI docs.
• performance/: includes Locust load testing scripts that simulate real traffic hitting your API to measure latency, throughput, and error rates.

This layout ensures that each type of test is separated by intent and runtime cost, giving you a clean way to scale your test suite over time.

Understanding Pytest Fixtures: Using conftest.py for Reusable Test Setup

The conftest.py file is the backbone of your testing environment. Pytest automatically loads fixtures defined here and makes them available across all test files without explicit imports.

Your project uses conftest.py to provide:

• FastAPI TestClient fixture: allows integration tests to call your API exactly the way a real HTTP client would.
• Sample input data: keeps repeated values out of your test files.
• Expected outputs: help tests stay focused on behavior rather than setup.

This shared setup reduces duplication, keeps tests clean, and ensures consistent test behavior across the entire suite.

Where to Place Tests in MLOps Projects: Unit vs Integration vs Performance

A simple rule-of-thumb keeps your test organization disciplined:

• Put tests in unit/ when the code under test does not require a running API or external system.
  Example: testing that the DummyModel.predict() returns “positive” for the word great.
• Put tests in integration/ when the test needs the full FastAPI app running.
  Example: calling /predict and checking that the API returns a JSON response.
• Put tests in performance/ when measuring speed, concurrency limits, or error behavior under load.
  Example: Locust scripts simulating dozens of users sending /predict requests at once.

Following this pattern ensures your tests remain stable, fast, and easy to reason about as the project grows.

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.

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!

Unit Testing in MLOps with Pytest

Unit tests are your first safety net in MLOps. Before you hit the API, spin up Locust, or ship to production, you want to know: Does my core prediction code behave exactly the way I think it does?

In this lesson, you do that by testing 2 things in isolation:

• inference service: services/inference_service.py
• dummy model: models/dummy_model.py

All of that is captured in tests/unit/test_inference_service.py.

The Code Under Test: Inference Service and Dummy Model

First, recall what you are testing.

services/inference_service.py

"""
Simple inference service for making model predictions.
"""
from models.dummy_model import DummyModel
from core.logger import logger

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


def predict(input_text: str) -> str:
    """
    Make a prediction using the loaded model.
   
    Args:
        input_text: Input text for prediction
       
    Returns:
        Prediction result as string
    """
    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 file does 3 things:

• Initializes a DummyModel once at import time and logs that it loaded.
• Exposes a predict(input_text: str) -> str function that:
  • Logs the incoming input (truncated to 50 chars).
  • Calls model.predict(...).
  • Logs and returns the prediction.
• Catches any exception, logs the error, and re-raises it so failures are visible.

You are not testing FastAPI here, just pure Python logic: given some text, does this function consistently return the correct label?

models/dummy_model.py

"""
Placeholder dummy model class.
"""
from typing import Any


class DummyModel:
    """
    A placeholder ML model class that returns fixed predictions.
    """
   
    def __init__(self) -> None:
        """Initialize the dummy model."""
        self.model_name = "dummy_classifier"
        self.version = "1.0.0"
   
    def predict(self, input_data: Any) -> str:
        """
        Make a prediction (returns a fixed string for demonstration).
       
        Args:
            input_data: Input data for prediction
           
        Returns:
            Fixed prediction string
        """
        text = str(input_data).lower()
        if "good" in text or "great" in text:
            return "positive"
        return "negative"

This model is deliberately simple:

• The constructor sets model_name and version for logging and version tracking.
• The predict() method:
  • Converts any input to lowercase text.
  • Returns "positive" if it sees "good" or "great" in the text.
  • Returns "negative" otherwise.

Your unit tests will assert that both the service and model behave exactly like this.

Writing Pytest Unit Tests for MLOps: test_inference_service.py

Here is the full unit test module:

"""
Unit tests for the inference service.
"""
import pytest
from services.inference_service import predict
from models.dummy_model import DummyModel


class TestInferenceService:
    """Test class for inference service."""
   
    def test_predict_returns_string(self):
        """Test that predict() returns a string."""
        result = predict("some input text")
        assert isinstance(result, str)
   
    def test_predict_positive_input(self):
        """Test prediction with positive input."""
        result = predict("This is good")
        assert result == "positive"
   
    def test_predict_negative_input(self):
        """Test prediction with negative input."""
        result = predict("This is bad")
        assert result == "negative"


class TestDummyModel:
    """Test class for DummyModel."""
   
    def test_model_initialization(self):
        """Test that the model initializes correctly."""
        model = DummyModel()
        assert model.model_name == "dummy_classifier"
        assert model.version == "1.0.0"
   
    def test_predict_with_good_word(self):
        """Test that the model returns positive for 'good'."""
        model = DummyModel()
        result = model.predict("This is good")
        assert result == "positive"
   
    def test_predict_with_great_word(self):
        """Test that the model returns positive for 'great'."""
        model = DummyModel()
        result = model.predict("This is great")
        assert result == "positive"
   
    def test_predict_without_keywords(self):
        """Test that the model returns negative without keywords."""
        model = DummyModel()
        test_inputs = ["test", "random text", "negative sentiment"]
        for input_text in test_inputs:
            result = model.predict(input_text)
            assert result == "negative"

Let us break it down.

Testing the Inference Service with Pytest (MLOps Unit Tests)

The first test class focuses on the service function, not the API:

class TestInferenceService:
    """Test class for inference service."""
   
    def test_predict_returns_string(self):
        """Test that predict() returns a string."""
        result = predict("some input text")
        assert isinstance(result, str)

• This test ensures predict() always returns a string, no matter what you pass in.
• If someone later changes predict() to return a dict, tuple, or Pydantic model, this test will fail immediately.

    def test_predict_positive_input(self):
        """Test prediction with positive input."""
        result = predict("This is good")
        assert result == "positive"
   
    def test_predict_negative_input(self):
        """Test prediction with negative input."""
        result = predict("This is bad")
        assert result == "negative"

These 2 tests verify the happy-path behavior:

• Text containing "good" should be classified as "positive".
• Text without "good" or "great" should default to "negative".

Notice what’s not happening here:

• No FastAPI client.
• No HTTP calls.
• No environment or config loading.

This is pure, fast, deterministic testing of the core service logic.

Testing ML Models in Isolation with Pytest

The second test class targets the model directly:

class TestDummyModel:
    """Test class for DummyModel."""
   
    def test_model_initialization(self):
        """Test that the model initializes correctly."""
        model = DummyModel()
        assert model.model_name == "dummy_classifier"
        assert model.version == "1.0.0"

• This verifies that your model is initialized correctly.
• In real projects, this might include loading weights, setting up devices, or configuration. Here, it is just model_name and version, but the pattern is the same.

    def test_predict_with_good_word(self):
        """Test that the model returns positive for 'good'."""
        model = DummyModel()
        result = model.predict("This is good")
        assert result == "positive"
   
    def test_predict_with_great_word(self):
        """Test that the model returns positive for 'great'."""
        model = DummyModel()
        result = model.predict("This is great")
        assert result == "positive"

• These tests assert that the keyword-based classification logic works: both "good" and "great" map to "positive".

    def test_predict_without_keywords(self):
        """Test that the model returns negative without keywords."""
        model = DummyModel()
        test_inputs = ["test", "random text", "negative sentiment"]
        for input_text in test_inputs:
            result = model.predict(input_text)
            assert result == "negative"

• This test loops over several neutral and negative phrases to make sure the model consistently returns “negative” when no positive keywords are present.
• This is your guardrail against accidental changes to the keyword logic.

How to Run Pytest Unit Tests for MLOps Projects

To run just these tests:

pytest tests/unit/ -v

Or with Poetry:

poetry run pytest tests/unit/ -v

You will see output similar to:

tests/unit/test_inference_service.py::TestInferenceService::test_predict_returns_string PASSED
tests/unit/test_inference_service.py::TestInferenceService::test_predict_positive_input PASSED
tests/unit/test_inference_service.py::TestInferenceService::test_predict_negative_input PASSED
tests/unit/test_inference_service.py::TestDummyModel::test_model_initialization PASSED
...

When everything is green, you know:

• Your core prediction logic is stable.
• The dummy model behaves exactly as designed.
• You can now safely move on to integration tests and performance tests in later sections.

Integration Testing in MLOps

Unit tests validate your core Python logic, but integration tests answer a different question:

“Does the entire application behave correctly when all components work together?”

This means testing:

• FastAPI app
• routing layer
• service functions
• model
• configuration loaded at runtime

All of this happens using FastAPI’s TestClient and your actual running application object (app from main.py).

Let’s break it down.

Using FastAPI TestClient for Integration Testing with Pytest

Your conftest.py defines a reusable client fixture:

from fastapi.testclient import TestClient
from main import app

@pytest.fixture
def client():
    """Create a test client for the FastAPI app."""
    return TestClient(app)

How FastAPI TestClient Works for API Testing

• TestClient(app) spins up an in-memory FastAPI instance.
• No server is launched, no networking occurs.
• Every test receives a fresh client that behaves exactly like a real HTTP client or API consumer.

This lets you write code such as:

response = client.get("/health")

as if you were calling a real deployed API, but entirely offline and deterministic.

Testing API Endpoints (/health, /predict)

Here is the integration test code from your repo:

class TestHealthEndpoint:
    def test_health_check_returns_ok(self, client):
        response = client.get("/health")

        assert response.status_code == 200
        assert response.json() == {"status": "ok"}
   
    def test_health_check_has_correct_content_type(self, client):
        response = client.get("/health")

        assert response.status_code == 200
        assert "application/json" in response.headers["content-type"]

What Integration Tests Verify in an MLOps API

• Your /health route is reachable.
• It always returns a 200 response.
• It returns valid JSON.
• The content type is correct.

Here is the real FastAPI code being tested (main.py):

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

This alignment is exactly correct.

Testing the /predict Endpoint in an MLOps API

Your integration tests call the prediction endpoint:

class TestPredictEndpoint:

    def test_predict_endpoint(self, client):
        response = client.post("/predict", params={"input": "good movie"})
        assert response.status_code == 200
        assert "prediction" in response.json()
   
    def test_predict_positive(self, client):
        response = client.post("/predict", params={"input": "This is a great movie!"})
        assert response.status_code == 200
        assert response.json()["prediction"] == "positive"
   
    def test_predict_negative(self, client):
        response = client.post("/predict", params={"input": "This is bad"})
        assert response.status_code == 200
        assert response.json()["prediction"] == "negative"

This tests:

• The endpoint exists and accepts POST requests.
• The parameter is correctly passed using params={"input": ...}.
• The internal inference logic (service → model) behaves correctly end-to-end.

Here is the actual API endpoint in your main.py:

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

Perfect 1:1 match.

Testing Documentation Endpoints (/docs, /openapi.json)

These are built into FastAPI and must exist for production ML systems.

Your tests:

class TestAPIDocumentation:
    def test_openapi_schema_accessible(self, client):
        response = client.get("/openapi.json")

        assert response.status_code == 200
        schema = response.json()
        assert "openapi" in schema
        assert "info" in schema
   
    def test_swagger_ui_accessible(self, client):
        response = client.get("/docs")

        assert response.status_code == 200
        assert "text/html" in response.headers["content-type"]

What This Ensures

• The OpenAPI schema is generated.
• Swagger UI loads successfully.
• No misconfiguration broke the docs.
• Consumers (frontend teams, other ML services, monitoring) can introspect your API.

This is standard for production ML systems.

Testing Error Handling in FastAPI APIs with Pytest

Your code includes error tests that verify robustness:

class TestErrorHandling:
    def test_nonexistent_endpoint_returns_404(self, client):
        response = client.get("/nonexistent")
        assert response.status_code == 404
   
    def test_invalid_method_on_health_endpoint(self, client):
        response = client.post("/health")
        assert response.status_code == 405  # Method Not Allowed
   
    def test_malformed_requests_handled_gracefully(self, client):
        response = client.get("/health")
        assert response.status_code == 200

Integration Test Breakdown: What Each Test Validates

These tests ensure your service behaves consistently even when clients behave incorrectly.

How to Run Integration Tests with Pytest in MLOps

To run only the integration tests:

Using pytest directly

pytest tests/integration/ -v

With Poetry

poetry run pytest tests/integration/ -v

With Makefile

make test-integration

You will see output like:

tests/integration/test_api_routes.py::TestHealthEndpoint::test_health_check_returns_ok PASSED
tests/integration/test_api_routes.py::TestPredictEndpoint::test_predict_positive PASSED
tests/integration/test_api_routes.py::TestAPIDocumentation::test_swagger_ui_accessible PASSED
...

Green = your API works correctly end-to-end.

Performance and Load Testing with Locust

Performance testing is critical for ML systems because even a lightweight model can become slow, unstable, or unresponsive when many users hit the API at once. With Locust, you can simulate hundreds or thousands of concurrent users calling your ML inference endpoints and measure how your API behaves under pressure.

This section explains why load testing matters, how Locust works, how your actual test file is structured, and how to interpret its results.

Why Load Testing Is Essential for MLOps and ML APIs

ML inference services have unique scaling behaviors:

• Model loading requires significant memory.
• Inference latency grows non-linearly under load.
• CPU/GPU bottlenecks show up only when multiple users hit the system.
• Thread starvation can cause cascading failures.
• Autoscaling decisions depend on real-world load patterns.

A service that performs well for one user may fail miserably at 50 users.

Load testing ensures:

• The API stays responsive under traffic.
• Latency stays under acceptable thresholds.
• No unexpected failures or timeouts occur.
• You understand the system’s scaling limits before going to production.

Locust is perfect for this because it is lightweight, Python-based, and designed for web APIs.

Locust Load Testing Concepts: Users, Spawn Rate, and Tasks Explained

Locust simulates user behavior using simple Python classes.

Users

A “user” is an independent client that continuously makes requests to your API.

Example:

• 10 users = 10 active clients repeatedly calling /predict.

Spawn rate

How quickly Locust ramps up users.

Example:

• spawn rate 2 = add 2 users per second until target is reached.

This helps simulate realistic traffic spikes instead of instantly launching all users.

Tasks

Each simulated user executes a set of tasks (e.g., repeatedly calling the /predict endpoint).

Every task can have a weight:

• Higher weight = more frequent calls.

This lets you mimic real user patterns like:

• 90% predict calls
• 10% health checks

Your project does exactly this.

Writing the locustfile.py

from locust import HttpUser, task, between

class MLAPIUser(HttpUser):
    """
    Locust user class for testing the ML API.
   
    Simulates a user making requests to the API endpoints.
    """
   
    # Wait between 1 and 3 seconds between requests
    wait_time = between(1, 3)
   
    @task(10)
    def test_predict(self):
        """
        Test the predict endpoint.
       
        This task has weight 10, making it the most frequently called.
        """
        payload = {"input": "The movie was good"}
        with self.client.post("/predict", params=payload, catch_response=True) as response:
            if response.status_code == 200:
                response_data = response.json()
                if "prediction" in response_data:
                    response.success()
                else:
                    response.failure(f"Missing prediction in response: {response_data}")
            else:
                response.failure(f"HTTP {response.status_code}")
   
    def on_start(self):
        """
        Called when a user starts testing.
       
        Used for setup tasks like authentication.
        """
        # Verify the API is reachable
        response = self.client.get("/health")
        if response.status_code != 200:
            print(f"Warning: API health check failed with status {response.status_code}")

What This Locust Load Test Validates in an MLOps API

• Creates a simulated user (MLAPIUser) that calls /predict.
• Gives the /predict task a weight of 10, making it the dominant request.
• Sends realistic input (“The movie was good”).
• Validates:
  • Response code is 200.
  • JSON contains “prediction”.
• Marks failures explicitly for clean reporting.
• On startup, each user verifies that /health works.

This matches your API perfectly:

• /predict is POST with query parameter input=...
• /health is GET and returns status OK

Nothing needs to be changed; this is production-quality.

Running Locust: Headless Mode vs Web UI Dashboard

Locust supports two modes.

A. Web UI Mode (Interactive Dashboard)

Launch Locust:

locust -f tests/performance/locustfile.py --host=http://localhost:8000

Then open:

http://localhost:8089

You will see a dashboard where you can:

• Set number of users
• Set spawn rate
• Start/stop tests
• View real-time stats

B. Headless Mode (Automated CI/CD or scripting)

You already have a script:

software-engineering-mlops-lesson2/scripts/run_locust.sh

Run:

./scripts/run_locust.sh http://localhost:8000 10 2 5m

This executes:

• 10 users
• spawn rate 2 users per second
• run time 5 minutes
• save HTML report

No UI; perfect for pipelines.

Generating Locust Load Testing Reports for ML APIs

Your script uses:

--html="reports/locust_reports/locust_report_<timestamp>.html"

Which produces files like:

reports/locust_reports/locust_report_20251030_031331.html

Each report includes:

• Requests per second (RPS)
• Failure stats
• Full latency distribution
• Percentiles (50th, 95th, 99th)
• Charts of active users and response times

These HTML reports are great for:

• Comparing deployments
• Regression testing API performance
• Flagging slow model versions
• Archiving performance history

Everything is already correctly set up in your repo.

Understanding Test Metrics (RPS, failures, latency, P95/P99)

Locust gives several performance metrics you must understand for ML systems.

Requests per Second (RPS)

How many inference calls your API can handle per second.

• CPU-bound models lead to low RPS
• Simple models lead to high RPS

Increasing users will show where your model and server saturates.

Failures

Locust marks a request as failed when:

• Status code ≠ 200
• Response JSON does not contain "prediction"
• Timeout occurs
• Server returns an internal error

Your catch_response=True logic handles this explicitly.

This prevents “hidden” failures.

Latency (ms)

Response time per request, typically measured in milliseconds.

For ML, latency is the most important metric.

You will see:

• Average latency
• Median (P50)
• Slowest (max latency)

P95 / P99 (Tail Latency)

The 95th and 99th percentile response times.

These capture worst-case behavior.

Example:

• P50 = 40 ms
• P95 = 210 ms
• P99 = 540 ms

This means:

Most users see fast responses, but a small % experience major slowdowns.

This is common in ML workloads due to:

• Model warmup
• Thread contention
• Python GIL blockage
• Model cache misses

Production Service Level Objectives (SLOs) usually track P95 and P99, not averages.

MLOps Test Configuration: YAML and Environment Variables

ML systems behave differently across production, development, and testing environments.

Your Lesson 2 codebase separates these environments cleanly using:

• A test-specific YAML config
• A modified BaseSettings loader
• .env overrides for test mode

This ensures that tests run quickly, deterministically, and without polluting real environment settings.

Let’s break down how this works.

Understanding test_config.yaml for MLOps Testing

# Test Configuration
environment: "test"
log_level: "DEBUG"

# API Configuration
api_host: "127.0.0.1"
api_port: 8000
debug: true

# Performance Testing
performance:
  baseline_users: 10
  spawn_rate: 2
  test_duration: "5m"

# Model Configuration
model:
  name: "dummy_classifier"
  version: "1.0.0"

What test_config.yaml Controls in MLOps Pipelines

This config prevents tests from accidentally picking up production configs.

Overriding Application Configuration in Test Mode

Your test environment uses a special configuration loader inside:

core/config.py

Here is the real code:

def load_config() -> Settings:
    # Load base settings from environment
    settings = Settings()
   
    # Load additional configuration from YAML if it exists
    config_path = "configs/test_config.yaml"
    if os.path.exists(config_path):
        yaml_config = load_yaml_config(config_path)
       
        # Override settings with YAML values if they exist
        for key, value in yaml_config.items():
            if hasattr(settings, key):
                setattr(settings, key, value)
   
    return settings

How Configuration Overrides Work: YAML and Environment Variables

• Step 1: BaseSettings loads environment variables
  (.env, operating system (OS) variables, defaults)
• Step 2: YAML configuration overrides them
  test_config.yaml replaces any matching fields in Settings.
• Final output:
  The application is now in test mode, completely isolated from development and production environments.

Why Configuration Management Matters in MLOps Testing

• Integration tests always use the same port, host, and log settings.
• Tests are repeatable and deterministic.
• You never accidentally load production API keys or endpoints.
• CI/CD pipelines get consistent behavior.

This pattern is very common in real-world MLOps systems.

Using Environment Variables for Test Isolation

Your test environment uses a .env.example file:

# API Configuration
API_PORT=8000
API_HOST=0.0.0.0
DEBUG=true

# Environment
ENVIRONMENT=test

# Logging
LOG_LEVEL=DEBUG

During setup, users run:

cp .env.example .env

This creates the .env used during tests.

Why Test-Specific .env Variables Matter

Combined with YAML Overrides

.env → applies defaults

test_config.yaml → overrides final values

This gives you a flexible and safe configuration stack.

Code Quality in MLOps: Linting, Formatting, and Static Analysis Tools

Testing ensures correctness, but code quality tools ensure that your ML system remains maintainable as it grows.

In Lesson 2, you introduce a full suite of professional-quality tooling:

• flake8 for linting
• Black for auto-formatting
• isort for import ordering
• MyPy for static typing
• Makefile for automation consistency

Together, they enforce the same engineering discipline used on real production ML teams at scale.

Linting Python Code with flake8

Linting catches code smells, stylistic issues, and subtle bugs before they hit production.

Your repository includes a real .flake8 file:

[flake8]
max-line-length = 88
extend-ignore = E203, W503
exclude =
    .git,
    __pycache__,
    .venv,
    venv,
    env,
    build,
    dist,
    *.egg-info,
    .pytest_cache,
    .mypy_cache
per-file-ignores =
    __init__.py:F401
max-complexity = 10

What your flake8 setup enforces

• 88-character line limit (matches Black)
• Ignores stylistic warnings that Black also overrides (E203, W503)
• Avoids checking generated or virtual-env directories
• Allows unused imports only in __init__.py files
• Enforces a maximum complexity score of 10

Run flake8 manually

poetry run flake8 .

Or via Makefile

make lint

Linting becomes part of your day-to-day workflow and prevents style drift across your ML services.

Formatting Python Code with Black Pipelines

Black is an automatic code formatter; it rewrites Python code into a consistent style.

Your Lesson 2 pyproject.toml includes:

[tool.black]
line-length = 88
target-version = ['py39']
include = '\.pyi?$'

This means:

• All Python files (.py) are formatted.
• Max line length is 88 chars.
• py39 syntax is allowed.

Format all code:

poetry run black .

Or using the Makefile shortcut:

make format

Black removes tedious decisions about spacing, commas, and line breaks, ensuring all contributors share the same style.

Using isort to Manage Python Imports

isort automatically manages import sorting and grouping.

Your pyproject.toml contains:

[tool.isort]
profile = "black"
multi_line_output = 3

This aligns isort’s output with Black’s formatting rules, avoiding conflicts.

How to Run isort for Clean Python Imports

poetry run isort .

Or via Makefile:

make format

Why This Matters

As ML services grow, import lists become messy. isort keeps them clean and consistent, improving readability exponentially.

Static Type Checking with MyPy for MLOps Codebases

Static typing is increasingly important in MLOps systems, especially when passing models, configs, and data structures between services.

Your repo contains a full mypy.ini:

[mypy]
python_version = 3.9
warn_return_any = True
warn_unused_configs = True
disallow_untyped_defs = False
ignore_missing_imports = True

[mypy-tests.*]
disallow_untyped_defs = False

[mypy-locust.*]
ignore_missing_imports = True

What This Config Enforces

• Flags functions that return Any
• Warns about unused config options
• Does not require type hints everywhere (reasonable for ML codebases)
• Skips type-checking external packages (common in ML pipelines)
• Allows untyped defs in tests

Run MyPy

poetry run mypy .

Or via Makefile:

make type-check

Why MyPy Is Critical in ML Systems

• Prevents silent type errors (e.g., passing a list where a tensor is expected)
• Catches config mistakes before runtime
• Improves refactor safety for large ML codebases

Using a Makefile to Automate MLOps Testing and Code Quality

Your Makefile automates all key development tasks:

make test          # Run all tests
make test-unit     # Unit tests only
make test-integration
make format        # Black + isort
make lint          # flake8
make type-check    # mypy
make load-test     # Locust performance tests
make clean         # Reset environment

This ensures:

• Every developer uses the same commands
• CI/CD pipelines can call the same interface
• Tooling stays consistent across machines

Example workflow for contributors:

make format
make lint
make type-check
make test

If all commands pass, you know your code is clean, consistent, and ready for production.

Automating Testing with a Pytest Test Runner Script

As your ML system grows, running dozens of unit, integration, and performance tests manually becomes tedious and error-prone.

Lesson 2 includes a fully automated test runner (scripts/run_tests.sh) that enforces a predictable, repeatable workflow for your entire test suite.

This script acts like a miniature CI pipeline that you can run locally. It prints structured logs, enforces failure conditions, and ensures that no test is accidentally skipped.

Running Automated Tests with run_tests.sh

Your repository includes a fully functional test runner:

#!/bin/bash

# Test Runner Script for MLOps Lesson 2

set -e

echo "🧪 Running MLOps Lesson 2 Tests..."

# Colors for output
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
RED='\033[0;31m'
NC='\033[0m'

print_status() {
    echo -e "${GREEN}✅ $1${NC}"
}

print_warning() {
    echo -e "${YELLOW}⚠️  $1${NC}"
}

print_error() {
    echo -e "${RED}❌ $1${NC}"
}

# Run unit tests
echo ""
echo "📝 Running unit tests..."
poetry run pytest tests/unit/ -v
if [ $? -eq 0 ]; then
    print_status "Unit tests passed"
else
    print_error "Unit tests failed"
    exit 1
fi

# Run integration tests
echo ""
echo "🔗 Running integration tests..."
poetry run pytest tests/integration/ -v
if [ $? -eq 0 ]; then
    print_status "Integration tests passed"
else
    print_error "Integration tests failed"
    exit 1
fi

echo ""
print_status "All tests completed successfully!"

How to Run It

./scripts/run_tests.sh

or, via Makefile:

make test

What It Does

• Runs unit tests
• Runs integration tests
• Stops immediately (set -e) if anything fails
• Prints colored output for clarity
• Provides a clear pass/fail summary

This mirrors real CI pipelines where a failing test stops deployment.

Understanding Pytest Output and Test Results

When you run the script, you will typically see output like this:

🧪 Running MLOps Lesson 2 Tests...

📝 Running unit tests...
============================= test session starts ==============================
collected 7 items

tests/unit/test_inference_service.py::TestInferenceService::test_predict_returns_string PASSED
tests/unit/test_inference_service.py::TestInferenceService::test_predict_positive_input PASSED
tests/unit/test_inference_service.py::TestInferenceService::test_predict_negative_input PASSED
tests/unit/test_inference_service.py::TestDummyModel::test_model_initialization PASSED
tests/unit/test_inference_service.py::TestDummyModel::test_predict_with_good_word PASSED
tests/unit/test_inference_service.py::TestDummyModel::test_predict_with_great_word PASSED
tests/unit/test_inference_service.py::TestDummyModel::test_predict_without_keywords PASSED

============================== 7 passed in 0.45s ===============================
✅ Unit tests passed

Then integration tests:

🔗 Running integration tests...

tests/integration/test_api_routes.py::TestHealthEndpoint::test_health_check_returns_ok PASSED
tests/integration/test_api_routes.py::TestPredictEndpoint::test_predict_positive PASSED
tests/integration/test_api_routes.py::TestAPIDocumentation::test_swagger_ui_accessible PASSED
tests/integration/test_api_routes.py::TestErrorHandling::test_nonexistent_endpoint_returns_404 PASSED

============================== 8 passed in 0.78s ===============================
✅ Integration tests passed

Finally:

✅ All tests completed successfully!

Why Automated Testing Workflows Matter in MLOps

• You see exactly which tests failed.
• You immediately know whether the API is healthy.
• You build the habit of treating tests as a gatekeeper before shipping ML code.

This is foundational MLOps workflow discipline.

Integrating Pytest into CI/CD Pipelines

Your test runner is already written as if it were part of CI.

Very soon, you will plug this into:

• GitHub Actions
• GitLab CI
• CircleCI
• AWS CodeBuild
• Azure DevOps

A typical GitHub Actions step would look like:

- name: Run Tests
  run: ./scripts/run_tests.sh

Since your script exits with non-zero status on failures, the CI job fails automatically.

What this enables in production ML workflows:

• No pull request gets merged unless tests pass
• Deployments are blocked if integration tests fail
• Load testing can be added as a gated step
• Test failures provide early feedback on regressions
• Teams enforce consistent standards across developers

You already have everything CI needs:

• A deterministic test runner
• A strict exit-on-fail system
• Separate unit and integration test layers
• Makefile wrappers for automation
• Poetry ensuring repeatable environments

Once you introduce CI/CD in later lessons, these scripts plug in seamlessly.

Automating Load Testing in MLOps with Locust Scripts

Performance testing becomes essential once an ML API starts supporting real traffic. You want confidence that your inference service will not collapse under load, that p95/p99 latencies remain acceptable, and that the system behaves predictably when scaling horizontally.

Manually running Locust is fine for experimentation, but production MLOps requires automated, repeatable load tests. Lesson 2 provides a dedicated script (run_locust.sh) which allows you to run performance tests in a single line and automatically generate HTML reports for analysis.

Running Automated Locust Load Tests with run_locust.sh

#!/bin/bash

# Simple Locust Load Testing Script for MLOps Lesson 2

set -e

echo "🚀 Starting Locust Load Testing..."

# Configuration
HOST=${1:-"http://localhost:8000"}
USERS=${2:-10}
SPAWN_RATE=${3:-2}
RUN_TIME=${4:-"5m"}

echo "🔧 Configuration: $USERS users, spawn rate $SPAWN_RATE, run time $RUN_TIME"

# Create reports directory
mkdir -p reports/locust_reports

# Check if the API is running
echo "🏥 Checking if API is running..."
if ! curl -s "$HOST/health" > /dev/null; then
    echo "❌ API is not reachable at $HOST"
    echo "Please start the API server first with: python main.py"
    exit 1
fi

echo "✅ API is reachable"

# Run Locust load test
echo "🧪 Starting load test..."

TIMESTAMP=$(date +"%Y%m%d_%H%M%S")
HTML_REPORT="reports/locust_reports/locust_report_$TIMESTAMP.html"

poetry run locust \
    -f tests/performance/locustfile.py \
    --host="$HOST" \
    --users="$USERS" \
    --spawn-rate="$SPAWN_RATE" \
    --run-time="$RUN_TIME" \
    --html="$HTML_REPORT" \
    --headless

echo "✅ Load test completed!"
echo "📊 Report: $HTML_REPORT"

How to Run It

Basic load test:

./scripts/run_locust.sh

10 users, spawn rate 2 users/sec, run for 5 minutes.

Custom parameters:

./scripts/run_locust.sh http://localhost:8000 30 5 2m

This means:

• 30 users total
• 5 users per second spawn rate
• 2-minute runtime
• Tests /predict endpoint repeatedly (because of locustfile.py)

What This Script Automates

• API health check before running
• Creates timestamped report directories
• Runs Locust in headless mode
• Stores HTML reports for analysis
• Fails gracefully when API is unreachable

This gives you a push-button reproducible performance test, a key requirement in professional MLOps.

Automatically Generating Load Testing Reports for ML APIs

Every run creates a unique HTML report:

reports/locust_reports/
    locust_report_20251203_031331.html
    locust_report_20251203_041215.html
    ...

This file includes:

• Requests per second (RPS)
• Response time percentiles (p50, p90, p95, p99)
• Failure rates
• Total requests
• Charts for concurrency vs performance
• Per-endpoint performance metrics

You can open the report in your browser:

open reports/locust_reports/locust_report_20251203_031331.html

(Windows)

start reports\locust_reports\locust_report_XXXX.html

Why This Is Important

Performance regressions are one of the most common ML service failures:

• model upgrades slow down inference unintentionally
• logging overhead increases latency
• new preprocessing increases CPU usage
• hardware changes alter throughput

By keeping each test run stored, you can compare historical performance.

This is the foundation of automatic performance regression detection.

Preparing Load Testing for CI/CD and Cloud MLOps Pipelines

Your load testing script is already CI-ready.

Here is how it fits into a production MLOps pipeline.

Option 1 — GitHub Actions

- name: Run Load Tests
  run: ./scripts/run_locust.sh http://localhost:8000 20 5 1m

Since the script exits non-zero on error, it becomes a gated step:

• Deployment is blocked if the API cannot sustain the expected load.
• Only performant builds reach production.

Option 2 — Nightly Performance Jobs

Teams often run Locust nightly to catch degradations early:

• baseline: 20 users
• alert if p95 > 300 ms
• alert if failures > 1%

Reports are archived automatically via your script.

Option 3 — Cloud Load Testing (AWS/GCP/Azure)

Your script can run inside:

• AWS CodeBuild
• Azure Pipelines
• Google CloudBuild

Simply modify the host:

./scripts/run_locust.sh https://staging.mycompany.com/api 50 10 10m

Why CI Load Tests Matter

• Prevents slow releases from being deployed
• Ensures model swaps do not tank performance
• Protects SLAs (Service Level Agreements)
• Helps capacity planning and autoscaling decisions
• Detects bottlenecks before customers do

Your repository already contains everything needed to industrialize performance testing.

Test Coverage in MLOps: Measuring and Improving Code Coverage

Even with strong unit, integration, and performance testing, you still need a way to quantify how much of your codebase is actually exercised. This is where test coverage comes in. Coverage tools show you which lines are tested, which are skipped, and where hidden bugs may still be lurking. This is especially important in ML systems, where subtle code paths (error handling, preprocessing, retry logic) can easily be missed.

Your Lesson 2 environment includes pytest-cov, allowing you to generate detailed coverage reports in a single command.

Using pytest-cov to Measure Test Coverage

Coverage is enabled simply by adding --cov flags to pytest.

Basic usage:

pytest --cov=.

Your repo’s pyproject.toml installs pytest-cov automatically under [tool.poetry.group.dev.dependencies], so coverage works out of the box.

A more detailed command:

pytest --cov=. --cov-report=term-missing

This reports:

• total coverage percentage
• which lines were executed
• which lines were missed
• hints for improving coverage

Example output you might see:

---------- coverage: platform linux, python 3.9 ----------
Name                                Stmts   Miss  Cover
--------------------------------------------------------
services/inference_service.py          22      0   100%
models/dummy_model.py                  16      0   100%
core/config.py                         40      8    80%
core/logger.py                         15      0   100%
tests/unit/test_inference_service.py   28      0   100%
--------------------------------------------------------
TOTAL                                 121      8    93%

This gives immediate visibility into which modules need more test attention.

How to Measure Code Coverage in MLOps Projects

To formally measure coverage for Lesson 2, run:

pytest -v --cov=. --cov-report=html

This generates a full HTML report inside:

htmlcov/index.html

Open it in your browser:

open htmlcov/index.html

(Windows)

start htmlcov\index.html

The HTML report visualizes:

• executed vs missed lines
• branch coverage
• per-module summaries
• clickable source code with line highlighting

This is the gold standard report format used in industry pipelines.

Integrating Coverage into Your Workflow

Your Makefile could easily support it:

make coverage

But even without that, pytest-cov gives you everything you need to evaluate test completeness.

How to Increase Test Coverage in MLOps Pipelines

ML systems often have unusual testing challenges:

• multiple code paths depending on data
• dynamic model loading
• error cases that only appear in production
• preprocessing/postprocessing steps
• branching logic based on config values
• retry and timeout logic
• logging behavior that might hide bugs

To increase coverage meaningfully:

1. Test failure modes

Example: model not loaded, invalid input, exceptions in service layer.

2. Test alternative branches

For example., your dummy model has:

if "good" in text or "great" in text:
    return "positive"
return "negative"

Coverage increases when you test:

• positive branch
• fallback branch
• edge cases like empty strings

3. Test configuration-dependent behavior

Since your system loads from:

• .env
• YAML
• runtime values

Try testing scenarios where each layer overrides the next.

4. Test logging paths

Logging is crucial in MLOps, and ensuring logs appear where expected also contributes to coverage.

5. Test the API under different payloads

Missing parameters, malformed types, unexpected values.

6. Test integration between modules

Even simple ML systems can break across module boundaries, so testing interactions raises coverage dramatically.

Recommended Test Coverage Targets for MLOps Systems

High coverage is good, but perfection is unrealistic and unnecessary.

Here are industry-grade ML-specific targets:

Why You Do Not Aim for 100%

• ML models are often treated as black boxes
• Some branches (especially failure conditions) are difficult to simulate
• Performance code paths are not always practical to test

A strong MLOps system targets:

Overall coverage: 80-90%

This ensures the most important logic is covered while avoiding diminishing returns.

Critical paths: 100%

Inference, preprocessing, conversion, routing, safety checks.

Performance-sensitive code: covered via load tests

This is why Locust complements pytest rather than replacing it.

What's next? We recommend PyImageSearch University.

Course information:
86+ total classes • 115+ hours hours of on-demand code walkthrough videos • Last updated: May 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 make ML systems safe, correct, and production-ready through a full testing and validation workflow. You started by understanding why ML services need far more than “just unit tests,” and how a layered approach (unit, integration, and performance tests) creates confidence in both the code and the behavior of the system. You then explored a real test layout with dedicated folders, fixtures, and isolation, and saw how each type of test validates a different piece of the pipeline.

From there, you implemented unit tests for the inference service and dummy model, followed by integration tests that exercise real FastAPI endpoints, documentation routes, and error handling. You also learned how to perform load testing with Locust, simulate concurrent users, generate performance reports, and interpret latency and failure metrics. This is an essential skill for production ML APIs.

Finally, you covered the tools that keep an ML codebase clean and maintainable: linting, formatting, static typing, and the Makefile commands that tie everything together. You closed with automated test runners, load-test scripts, and coverage reporting, giving you an end-to-end workflow that mirrors real MLOps engineering practice.

By now, you have seen how professional ML systems are tested, validated, measured, and maintained. This sets you up for the next module, where we will begin building data pipelines and reproducible ML workflows.

Citation Information

Singh, V. “Pytest Tutorial: MLOps Testing, Fixtures, and Locust Load Testing,” PyImageSearch, S. Huot, A. Sharma, and P. Thakur, eds., 2026, https://pyimg.co/4ztdu

@incollection{Singh_2026_pytest-tutorial-mlops-testing-fixtures-locust-load-testing,
  author = {Vikram Singh},
  title = {{Pytest Tutorial: MLOps Testing, Fixtures, and Locust Load Testing}},
  booktitle = {PyImageSearch},
  editor = {Susan Huot and Aditya Sharma and Piyush Thakur},
  year = {2026},
  url = {https://pyimg.co/4ztdu},
}

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 Pytest Tutorial: MLOps Testing, Fixtures, and Locust Load Testing appeared first on PyImageSearch.

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: May 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: May 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: May 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: May 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: May 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: May 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.