Overview and Architecture

This guide covers how to build a full machine learning pipeline that combines PyTorch for model training with Milvus as a vector database — a common architecture for embedding-based retrieval, semantic search, recommendation systems, and similarity learning.

Why Combine Milvus and PyTorch?

PyTorch is the de facto framework for training deep learning models. Milvus is a purpose-built vector database optimized for storing, indexing, and querying high-dimensional embeddings at scale. Together they enable architectures that were previously impractical:

  • Retrieval-Augmented Training (RAT): During training, fetch the most similar embeddings from Milvus to use as hard negatives or context.
  • Approximate Nearest Neighbor (ANN) Search: After training, your model’s outputs can be indexed in Milvus for millisecond-latency semantic search across billions of vectors.
  • Continual Learning: As your model improves, periodically re-index updated embeddings without retraining from scratch.
  • Few-Shot and Metric Learning: Use Milvus to retrieve labeled examples at inference time, enabling k-NN classification without retraining.

High-Level Data Flow

flowchart TD
    A[Raw Data] --> B["PyTorch Encoder
    CNN / Transformer / MLP"]
    B -->|"forward pass produces embedding vectors"| C[("Milvus Collection")]
    C -->|"ANN search returns similar vectors"| D["Hard Negative Mining
    Retrieval Context"]
    D --> E["Loss Function
    Contrastive / Triplet / InfoNCE"]
    E --> F["Optimizer Step
    Adam / SGD"]
    F -->|"periodic re-indexing"| C

The loop above describes a self-supervised or metric learning scenario where Milvus actively participates in the training loop. In simpler use cases (e.g., plain classification), Milvus is used purely at inference time after training completes.

Prerequisites

Before starting, make sure you are comfortable with the following:

  • Python 3.8+ — all code examples use Python 3.10 idioms.
  • Basic PyTorch — tensors, nn.Module, DataLoader, training loops.
  • Docker — Milvus is most easily run via Docker Compose.
  • Linear algebra fundamentals — understanding of cosine similarity and dot products helps when choosing distance metrics.

Knowledge of Key Concepts

Concept Why You Need It
Embedding vectors The output of your model that gets stored in Milvus
ANN indexing (HNSW, IVF) Determines query speed vs. recall trade-off in Milvus
Contrastive / triplet loss Common loss functions that benefit from Milvus-backed negative mining
Batch normalization Important for stable embedding norms

Setting Up Your Environment

Python Dependencies

Create a virtual environment and install all required packages:

python -m venv venv
source venv/bin/activate        # On Windows: venv\Scripts\activate

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
pip install pymilvus>=2.4.0
pip install numpy scipy scikit-learn
pip install tqdm matplotlib
Note

Replace the PyTorch CUDA index URL with the appropriate version for your GPU. For CPU-only training, use pip install torch without the extra index.

Verify Installations

import torch
import pymilvus

print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
print(f"PyMilvus version: {pymilvus.__version__}")

Expected output (versions may vary):

PyTorch version: 2.3.0+cu121
CUDA available: True
PyMilvus version: 2.4.1

Understanding Milvus

Core Concepts

Before writing any code it is important to understand how Milvus organizes data:

Collection The top-level container in Milvus, analogous to a table in a relational database. A collection has a fixed schema that defines the fields stored for each entity.

Schema Defines the structure of a collection. Every collection must have a primary key field and at least one FLOAT_VECTOR or BINARY_VECTOR field.

Partition An optional subdivision of a collection. Partitions let you scope queries to a subset of data (e.g., partitioning by category or split — train, val, test).

Index An ANN data structure built over a vector field. Without an index, Milvus falls back to brute-force search. Common index types:

Index Description Best For
FLAT Exact brute-force search Small datasets (<1M), maximum recall needed
IVF_FLAT Inverted file, flat quantization General-purpose, medium datasets
IVF_SQ8 IVF with scalar quantization Memory-constrained environments
HNSW Hierarchical Navigable Small World High-recall, high-throughput production use
DISKANN Disk-based ANN Billion-scale datasets that don’t fit in RAM

Metric Type The distance function used when comparing vectors. Your choice here must match how your model was trained:

  • L2 — Euclidean distance (lower = more similar)
  • IP — Inner product / cosine similarity (higher = more similar, when vectors are L2-normalized)
  • COSINE — Cosine similarity (Milvus 2.4+, vectors normalized automatically)

Milvus Architecture Overview

flowchart LR
    subgraph Client["Client Layer"]
        PY[PyMilvus SDK]
    end
    subgraph Milvus["Milvus Standalone"]
        direction TB
        PROXY[Proxy] --> ROOT[Root Coord]
        PROXY --> DATA[Data Coord]
        PROXY --> INDEX[Index Coord]
        PROXY --> QUERY[Query Coord]
    end
    subgraph Storage["Persistent Storage"]
        ETCD[(etcd Metadata)]
        MINIO[(MinIO Object Store)]
    end
    PY --> PROXY
    DATA --> MINIO
    INDEX --> MINIO
    ROOT --> ETCD

PyMilvus Client Modes

PyMilvus 2.4 introduced a new MilvusClient that simplifies connection management. There is also the older connections + Collection API. This guide primarily uses MilvusClient for simplicity but shows the lower-level API where needed.

from pymilvus import MilvusClient

# Connect to a local Milvus instance
client = MilvusClient(uri="http://localhost:19530")

# For in-process Milvus Lite (no server required, good for development)
client = MilvusClient(uri="./milvus_local.db")

Milvus Lite is a zero-dependency embedded version ideal for experimentation. It stores data in a local SQLite-backed file. Switch to the full Milvus server when your dataset exceeds ~1M vectors or when you need concurrent read/write.

Setting Up Milvus

Option A: Milvus Lite (Development)

No Docker required. Milvus Lite is bundled with pymilvus>=2.4:

from pymilvus import MilvusClient

client = MilvusClient("./dev_embeddings.db")
print("Milvus Lite ready.")

Option B: Docker Compose (Production-like)

Save the following as docker-compose.yml:

version: '3.5'

services:
  etcd:
    container_name: milvus-etcd
    image: quay.io/coreos/etcd:v3.5.5
    environment:
      - ETCD_AUTO_COMPACTION_MODE=revision
      - ETCD_AUTO_COMPACTION_RETENTION=1000
      - ETCD_QUOTA_BACKEND_BYTES=4294967296
      - ETCD_SNAPSHOT_COUNT=50000
    command: etcd -advertise-client-urls=http://127.0.0.1:2379 -listen-client-urls http://0.0.0.0:2379 --data-dir /etcd
    volumes:
      - ${DOCKER_VOLUME_DIRECTORY:-.}/volumes/etcd:/etcd

  minio:
    container_name: milvus-minio
    image: minio/minio:RELEASE.2023-03-20T20-16-18Z
    environment:
      MINIO_ACCESS_KEY: minioadmin
      MINIO_SECRET_KEY: minioadmin
    volumes:
      - ${DOCKER_VOLUME_DIRECTORY:-.}/volumes/minio:/minio_data
    command: minio server /minio_data --console-address ":9001"
    healthcheck:
      test: ["CMD", "curl", "-f", "http://localhost:9000/minio/health/live"]
      interval: 30s
      timeout: 20s
      retries: 3

  standalone:
    container_name: milvus-standalone
    image: milvusdb/milvus:v2.4.5
    command: ["milvus", "run", "standalone"]
    environment:
      ETCD_ENDPOINTS: etcd:2379
      MINIO_ADDRESS: minio:9000
    volumes:
      - ${DOCKER_VOLUME_DIRECTORY:-.}/volumes/milvus:/var/lib/milvus
    ports:
      - "19530:19530"
      - "9091:9091"
    depends_on:
      - etcd
      - minio

Start the stack:

docker compose up -d
docker compose ps     # verify all three containers are running

Creating a Collection

from pymilvus import MilvusClient, DataType

client = MilvusClient(uri="http://localhost:19530")

COLLECTION_NAME = "image_embeddings"
EMBEDDING_DIM   = 512     # must match your model's output dimension

# Drop if it already exists (useful during development)
if client.has_collection(COLLECTION_NAME):
    client.drop_collection(COLLECTION_NAME)

# Define schema
schema = client.create_schema(
    auto_id=True,            # Milvus assigns primary keys automatically
    enable_dynamic_field=True  # allows storing arbitrary extra fields
)

schema.add_field(
    field_name="id",
    datatype=DataType.INT64,
    is_primary=True,
    auto_id=True
)

schema.add_field(
    field_name="embedding",
    datatype=DataType.FLOAT_VECTOR,
    dim=EMBEDDING_DIM
)

# Optional scalar fields for metadata / filtering
schema.add_field(field_name="label",     datatype=DataType.INT32)
schema.add_field(field_name="split",     datatype=DataType.VARCHAR, max_length=16)
schema.add_field(field_name="file_path", datatype=DataType.VARCHAR, max_length=512)

# Create collection
client.create_collection(
    collection_name=COLLECTION_NAME,
    schema=schema
)
print(f"Collection '{COLLECTION_NAME}' created.")

# Create an HNSW index on the vector field
index_params = client.prepare_index_params()
index_params.add_index(
    field_name="embedding",
    index_type="HNSW",
    metric_type="IP",       # inner product (assumes L2-normalized vectors)
    params={"M": 16, "efConstruction": 200}
)

client.create_index(COLLECTION_NAME, index_params)
client.load_collection(COLLECTION_NAME)
print("Index built and collection loaded into memory.")

Building the PyTorch Model

Defining the Encoder

We will build a generic image encoder that maps images to 512-dimensional L2-normalized embeddings. The same principles apply to text encoders (BERT, etc.) or tabular encoders (MLP).

Model Architecture

flowchart LR
    IMG["Input Image
    3 x 224 x 224"] --> BB["ResNet-50 Backbone
    (frozen or fine-tuned)"]
    BB --> FLAT["Flatten
    2048-dim"]
    FLAT --> LIN1["Linear
    2048 -> 2048"]
    LIN1 --> BN1["BatchNorm1d"]
    BN1 --> RELU["ReLU"]
    RELU --> LIN2["Linear
    2048 -> 512"]
    LIN2 --> BN2["BatchNorm1d"]
    BN2 --> NORM["L2 Normalize"]
    NORM --> EMB["Embedding
    512-dim unit vector"]


import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models as models


class ImageEncoder(nn.Module):
    """
    A ResNet-50 backbone with a trainable projection head.
    Outputs L2-normalized embeddings of dimension `embed_dim`.
    """

    def __init__(self, embed_dim: int = 512, pretrained: bool = True):
        super().__init__()

        # Load backbone
        weights = models.ResNet50_Weights.IMAGENET1K_V2 if pretrained else None
        backbone = models.resnet50(weights=weights)

        # Remove the original classification head
        self.backbone = nn.Sequential(*list(backbone.children())[:-1])
        backbone_out_dim = 2048  # ResNet-50 penultimate layer

        # Projection head: 2048 -> embed_dim
        self.projector = nn.Sequential(
            nn.Linear(backbone_out_dim, backbone_out_dim),
            nn.BatchNorm1d(backbone_out_dim),
            nn.ReLU(inplace=True),
            nn.Linear(backbone_out_dim, embed_dim),
            nn.BatchNorm1d(embed_dim)
        )

        self.embed_dim = embed_dim

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: (batch, 3, H, W)
        features = self.backbone(x)               # (batch, 2048, 1, 1)
        features = features.flatten(start_dim=1)  # (batch, 2048)
        embeddings = self.projector(features)     # (batch, embed_dim)
        return F.normalize(embeddings, p=2, dim=1)  # L2 normalize
Tip

Why L2 normalize? When vectors are unit-norm, inner product equals cosine similarity. This makes the IP metric in Milvus equivalent to cosine similarity, which is what most contrastive losses optimize.

Loss Functions

Contrastive Loss (SimCLR-style)

class NTXentLoss(nn.Module):
    """
    Normalized Temperature-scaled Cross Entropy Loss.
    Used in SimCLR: https://arxiv.org/abs/2002.05709
    """

    def __init__(self, temperature: float = 0.07):
        super().__init__()
        self.temperature = temperature
        self.criterion = nn.CrossEntropyLoss()

    def forward(self, z_i: torch.Tensor, z_j: torch.Tensor) -> torch.Tensor:
        """
        z_i, z_j: two augmented views of the same batch (both L2-normalized).
        Shape: (batch_size, embed_dim)
        """
        batch_size = z_i.shape[0]
        z = torch.cat([z_i, z_j], dim=0)          # (2B, D)
        sim = torch.mm(z, z.T) / self.temperature  # (2B, 2B)

        # Mask out self-similarity on the diagonal
        mask = torch.eye(2 * batch_size, dtype=torch.bool, device=z.device)
        sim.masked_fill_(mask, float('-inf'))

        # Positive pairs: (i, i+B) and (i+B, i)
        labels = torch.cat([
            torch.arange(batch_size, 2 * batch_size),
            torch.arange(0, batch_size)
        ]).to(z.device)

        return self.criterion(sim, labels)

Triplet Loss with Hard Negative Mining

class TripletLoss(nn.Module):
    """
    Triplet margin loss with online hard negative mining.
    """

    def __init__(self, margin: float = 0.3):
        super().__init__()
        self.margin = margin

    def forward(
        self,
        anchors: torch.Tensor,
        positives: torch.Tensor,
        negatives: torch.Tensor
    ) -> torch.Tensor:
        d_pos = 1.0 - (anchors * positives).sum(dim=1)  # 1 - cosine sim
        d_neg = 1.0 - (anchors * negatives).sum(dim=1)
        losses = F.relu(d_pos - d_neg + self.margin)
        return losses.mean()

DataLoader Setup

import torchvision.transforms as T
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader

def get_transforms(split: str):
    if split == "train":
        return T.Compose([
            T.RandomResizedCrop(224, scale=(0.6, 1.0)),
            T.RandomHorizontalFlip(),
            T.ColorJitter(0.4, 0.4, 0.4, 0.1),
            T.RandomGrayscale(p=0.2),
            T.ToTensor(),
            T.Normalize(mean=[0.485, 0.456, 0.406],
                        std=[0.229, 0.224, 0.225]),
        ])
    else:
        return T.Compose([
            T.Resize(256),
            T.CenterCrop(224),
            T.ToTensor(),
            T.Normalize(mean=[0.485, 0.456, 0.406],
                        std=[0.229, 0.224, 0.225]),
        ])

train_dataset = ImageFolder(root="./data/train", transform=get_transforms("train"))
val_dataset   = ImageFolder(root="./data/val",   transform=get_transforms("val"))

train_loader = DataLoader(
    train_dataset,
    batch_size=128,
    shuffle=True,
    num_workers=8,
    pin_memory=True,
    drop_last=True        # important for contrastive loss
)

val_loader = DataLoader(
    val_dataset,
    batch_size=256,
    shuffle=False,
    num_workers=4,
    pin_memory=True
)

Generating and Storing Embeddings in Milvus

Before using Milvus during the training loop, you typically do an initial embedding pass over your dataset to populate the index.

Embedding Insertion Flow

flowchart TD
    DL[DataLoader batch] --> FWD[Model forward pass]
    FWD --> EMB[Embeddings tensor]
    EMB --> BUF[Append to buffer]
    BUF --> CHECK{Buffer size >= 1000?}
    CHECK -- Yes --> INS[client.insert]
    INS --> CLR[Clear buffer]
    CLR --> DL
    CHECK -- No --> DL
    DL -- Exhausted --> LAST{Buffer empty?}
    LAST -- No --> INS2[client.insert remainder]
    INS2 --> FLUSH[client.flush]
    LAST -- Yes --> FLUSH
    FLUSH --> DONE([Done])


import numpy as np
from tqdm import tqdm


def embed_dataset(
    model: nn.Module,
    loader: DataLoader,
    device: torch.device,
    split: str,
    client: MilvusClient,
    collection_name: str,
    batch_size_insert: int = 1000
):
    """
    Run the model over `loader`, then upsert embeddings into Milvus.
    Returns a list of all assigned Milvus IDs.
    """
    model.eval()
    all_ids = []
    buffer = []

    with torch.no_grad():
        for images, labels in tqdm(loader, desc=f"Embedding [{split}]"):
            images = images.to(device, non_blocking=True)
            embeddings = model(images)  # (B, D), L2-normalized

            # Build list of dicts for pymilvus insert
            for emb, lbl in zip(embeddings.cpu().numpy(), labels.numpy()):
                buffer.append({
                    "embedding": emb.tolist(),
                    "label": int(lbl),
                    "split": split,
                })

            # Batch insert to avoid holding too much in RAM
            if len(buffer) >= batch_size_insert:
                result = client.insert(
                    collection_name=collection_name,
                    data=buffer
                )
                all_ids.extend(result["ids"])
                buffer.clear()

    # Flush remaining
    if buffer:
        result = client.insert(
            collection_name=collection_name,
            data=buffer
        )
        all_ids.extend(result["ids"])

    # Flush to make data searchable immediately
    client.flush(collection_name)
    print(f"Inserted {len(all_ids)} embeddings for split='{split}'.")
    return all_ids


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model  = ImageEncoder(embed_dim=512, pretrained=True).to(device)

train_ids = embed_dataset(model, train_loader, device, "train", client, COLLECTION_NAME)
val_ids   = embed_dataset(model, val_loader,   device, "val",   client, COLLECTION_NAME)

Training with Milvus-Backed Retrieval

This section covers the core loop where Milvus actively participates in training by supplying hard negatives — examples that are semantically similar to an anchor but belong to a different class. Hard negatives are far more informative than random negatives for metric learning.

The Hard Negative Mining Pattern

flowchart LR
    A[Anchor embedding] --> SEARCH["Milvus ANN search
    top-k results"]
    SEARCH --> FILTER{Same label as anchor?}
    FILTER -- Yes --> SKIP[Skip]
    FILTER -- No --> SELECT[Select as hard negative]
    SKIP --> FILTER
    SELECT --> LOSS["Triplet Loss
    anchor, positive, hard_neg"]

For each anchor in the batch:

  1. Compute the anchor’s current embedding.
  2. Query Milvus for the top-k most similar vectors that have a different label.
  3. Use those returned vectors as negatives in the loss function.
def mine_hard_negatives(
    client: MilvusClient,
    collection_name: str,
    anchor_embeddings: np.ndarray,   # (B, D)
    anchor_labels: np.ndarray,        # (B,)
    top_k: int = 64
) -> np.ndarray:
    """
    For each anchor, find the hardest negative (closest vector with a different label).
    Returns hard_negatives of shape (B, D).
    """
    hard_negatives = np.zeros_like(anchor_embeddings)

    results = client.search(
        collection_name=collection_name,
        data=anchor_embeddings.tolist(),
        anns_field="embedding",
        search_params={"metric_type": "IP", "params": {"ef": 128}},
        limit=top_k,
        output_fields=["embedding", "label"],
        partition_names=["train"]     # only search the training split
    )

    for i, (query_results, anchor_label) in enumerate(zip(results, anchor_labels)):
        chosen = None
        for hit in query_results:
            if hit["entity"]["label"] != int(anchor_label):
                chosen = np.array(hit["entity"]["embedding"], dtype=np.float32)
                break

        if chosen is None:
            # Fallback: use the farthest result in the returned set
            chosen = np.array(query_results[-1]["entity"]["embedding"], dtype=np.float32)

        hard_negatives[i] = chosen

    return hard_negatives

Full Training Loop

import os
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR


def train(
    model: nn.Module,
    train_loader: DataLoader,
    client: MilvusClient,
    collection_name: str,
    num_epochs: int = 50,
    reindex_every: int = 5,
    device: torch.device = torch.device("cuda"),
    checkpoint_dir: str = "./checkpoints"
):
    os.makedirs(checkpoint_dir, exist_ok=True)

    optimizer = AdamW(model.parameters(), lr=3e-4, weight_decay=1e-4)
    scheduler = CosineAnnealingLR(optimizer, T_max=num_epochs)
    triplet_loss = TripletLoss(margin=0.3)

    for epoch in range(1, num_epochs + 1):
        model.train()
        epoch_loss = 0.0

        for batch_idx, (images, labels) in enumerate(tqdm(train_loader, desc=f"Epoch {epoch}")):
            images = images.to(device, non_blocking=True)
            labels_np = labels.numpy()

            # --- Forward pass ---
            embeddings = model(images)           # (B, 512), L2-normalized
            emb_np = embeddings.detach().cpu().numpy()

            # --- Mine hard negatives from Milvus ---
            hard_neg_np = mine_hard_negatives(
                client, collection_name, emb_np, labels_np
            )
            hard_neg = torch.tensor(hard_neg_np, dtype=torch.float32, device=device)
            hard_neg = F.normalize(hard_neg, p=2, dim=1)

            # --- Construct positives (augmented duplicate in same batch) ---
            # Shuffle within same-label groups to get positives
            positive_idx = torch.zeros(len(labels), dtype=torch.long)
            for lbl in labels.unique():
                mask = (labels == lbl).nonzero(as_tuple=True)[0]
                if len(mask) >= 2:
                    shuffled = mask[torch.randperm(len(mask))]
                    positive_idx[mask] = shuffled
                else:
                    positive_idx[mask] = mask  # fallback: self-positive (degenerate)
            positives = embeddings[positive_idx]

            # --- Compute loss ---
            loss = triplet_loss(embeddings, positives, hard_neg)

            # --- Backward pass ---
            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()

            epoch_loss += loss.item()

        scheduler.step()
        avg_loss = epoch_loss / len(train_loader)
        print(f"Epoch {epoch:3d} | Loss: {avg_loss:.4f} | LR: {scheduler.get_last_lr()[0]:.2e}")

        # --- Periodically re-index to keep Milvus up-to-date ---
        if epoch % reindex_every == 0:
            print(f"Re-indexing after epoch {epoch}...")
            client.delete(collection_name, filter=f'split == "train"')
            embed_dataset(model, train_loader, device, "train", client, collection_name)

        # --- Save checkpoint ---
        ckpt_path = os.path.join(checkpoint_dir, f"encoder_epoch{epoch:03d}.pt")
        torch.save({
            "epoch": epoch,
            "model_state_dict": model.state_dict(),
            "optimizer_state_dict": optimizer.state_dict(),
            "loss": avg_loss
        }, ckpt_path)

Updating the Index During Training

Keeping the Milvus index synchronized with a changing model is one of the more nuanced aspects of this architecture. Here are three strategies:

Strategy Comparison

flowchart TD
    START([Choose re-index strategy]) --> Q1{Dataset size}
    Q1 -- "< 500K" --> FULL[Strategy 1: Full Re-Index - Delete all, re-embed all]
    Q1 -- "500K - 10M" --> INCR[Strategy 2: Incremental Update - Re-embed high-loss samples only]
    Q1 -- "> 10M" --> MOM[Strategy 3: Momentum Encoder - EMA model writes to Milvus continuously]
    FULL --> SIMPLE[Simple and consistent - Expensive per epoch]
    INCR --> MEDIUM[Efficient and targeted - Requires loss tracking]
    MOM --> STABLE[Stable negatives - Most complex to implement]

Strategy 1: Periodic Full Re-Index (Simple)

As shown in the training loop above, delete all training embeddings every N epochs and recompute them. This is simple but expensive if your dataset is large.

def full_reindex(model, loader, device, client, collection_name, split="train"):
    client.delete(collection_name, filter=f'split == "{split}"')
    client.flush(collection_name)
    embed_dataset(model, loader, device, split, client, collection_name)

When to use: Datasets under ~500K samples, re-index every 5–10 epochs.

Strategy 2: Incremental Update (Efficient)

Maintain a mapping from dataset index to Milvus ID. After each epoch, re-embed only items whose loss exceeded a threshold, then call upsert:

def incremental_update(
    high_loss_indices: list,
    dataset,
    model: nn.Module,
    id_map: dict,       # dataset_index -> milvus_id
    client: MilvusClient,
    collection_name: str,
    device: torch.device
):
    model.eval()
    with torch.no_grad():
        for idx in high_loss_indices:
            image, label = dataset[idx]
            tensor = image.unsqueeze(0).to(device)
            emb = model(tensor).squeeze(0).cpu().numpy()

            client.upsert(
                collection_name=collection_name,
                data=[{
                    "id": id_map[idx],
                    "embedding": emb.tolist(),
                    "label": int(label),
                    "split": "train"
                }]
            )
    client.flush(collection_name)

When to use: Large datasets where full re-index is too costly; track per-sample loss and only update stale embeddings.

Strategy 3: Momentum Encoder (Advanced)

Inspired by MoCo (He et al., 2020), maintain two model weights:

  • model (online encoder): updated by gradients each step.
  • momentum_model (key encoder): updated as an exponential moving average of model.

Milvus stores embeddings produced by momentum_model, which provides more stable negatives:

import copy

momentum_model = copy.deepcopy(model)
momentum_model.eval()

MOMENTUM = 0.999

def update_momentum_encoder(model, momentum_model, m=MOMENTUM):
    with torch.no_grad():
        for param_q, param_k in zip(model.parameters(), momentum_model.parameters()):
            param_k.data = m * param_k.data + (1.0 - m) * param_q.data

Call update_momentum_encoder at the end of each optimizer step. Use momentum_model for generating the embeddings inserted into Milvus.

Productionizing the Pipeline

Checkpointing and Resuming

def load_checkpoint(checkpoint_path: str, model: nn.Module, optimizer=None):
    ckpt = torch.load(checkpoint_path, map_location="cpu")
    model.load_state_dict(ckpt["model_state_dict"])
    if optimizer and "optimizer_state_dict" in ckpt:
        optimizer.load_state_dict(ckpt["optimizer_state_dict"])
    start_epoch = ckpt.get("epoch", 0) + 1
    print(f"Resumed from epoch {ckpt['epoch']}, loss={ckpt['loss']:.4f}")
    return start_epoch

Export Paths

flowchart LR
    MODEL[Trained PyTorch Model] --> JIT[TorchScript Export .torchscript.pt]
    MODEL --> ONNX[ONNX Export .onnx]
    JIT --> SERVE_JIT[Serve via TorchServe / LibTorch]
    ONNX --> SERVE_ONNX[Serve via ONNX Runtime / Triton]
    SERVE_JIT --> MILVUS[(Milvus Index)]
    SERVE_ONNX --> MILVUS

Exporting the Encoder

# PyTorch JIT (TorchScript)
scripted = torch.jit.script(model.eval().cpu())
scripted.save("encoder.torchscript.pt")

# ONNX export
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
    model.eval().cpu(),
    dummy_input,
    "encoder.onnx",
    input_names=["image"],
    output_names=["embedding"],
    dynamic_axes={"image": {0: "batch_size"}, "embedding": {0: "batch_size"}},
    opset_version=17
)

Milvus in Production

Once your collection is stable and you have finished training, tune the HNSW search parameters for your recall / latency target:

# Higher ef = better recall, slower query
search_params = {"metric_type": "IP", "params": {"ef": 256}}

# Compact storage (removes deleted vectors permanently)
client.compact(COLLECTION_NAME)

# Check collection stats
stats = client.get_collection_stats(COLLECTION_NAME)
print(f"Row count: {stats['row_count']}")

Scaling Considerations

Concern Recommendation
Dataset > 10M vectors Use DISKANN index or shard across multiple collections
Query latency < 10ms Deploy Milvus with GPU indexing (GPU_IVF_FLAT or GPU_CAGRA)
High write throughput Increase dataCoord.segment.maxSize in Milvus config
Multiple model versions Use separate collections or partitions per model checkpoint

Common Pitfalls and Debugging

Embedding Dimensions Don’t Match

Symptom: pymilvus.exceptions.MilvusException: ... vector size 768 not match collection dim 512

Fix: Ensure your model’s embed_dim matches the dim set in the collection schema. If you change the architecture, drop and recreate the collection.

# Always print the shape of a sample embedding before inserting
sample = model(torch.randn(1, 3, 224, 224).to(device))
print(f"Embedding shape: {sample.shape}")   # should be (1, 512)

Vectors Are Not Normalized

Symptom: Cosine similarity scores > 1 or nonsensical; loss doesn’t converge.

Fix: Always call F.normalize(embeddings, p=2, dim=1) before inserting into Milvus when using IP metric. Verify:

norms = embeddings.norm(dim=1)
assert torch.allclose(norms, torch.ones_like(norms), atol=1e-5), \
    f"Embeddings not unit-norm. Min={norms.min():.4f}, Max={norms.max():.4f}"

Milvus Returns Stale Embeddings

Symptom: Hard negatives seem too easy despite many training epochs.

Fix: Either increase re-indexing frequency or use the momentum encoder strategy. Also call client.flush() after inserting to ensure newly inserted vectors are searchable.

Out of Memory During Re-Index

Symptom: CUDA OOM during embedding pass.

Fix: Use torch.no_grad(), reduce batch_size during the embedding pass, and use pin_memory=True + non_blocking=True for efficient GPU transfers.

Index Not Loaded

Symptom: MilvusException: collection not loaded

Fix: After creating a collection or reconnecting, always call:

client.load_collection(COLLECTION_NAME)

Collections are loaded lazily into memory. In production, use Milvus’s autoLoad configuration to load on startup.

Full End-to-End Example

Below is a self-contained script that ties everything together. Assumes you have your data in ./data/train and ./data/val as ImageFolder-compatible directories.

"""
end_to_end_milvus_pytorch.py

Full training pipeline: PyTorch image encoder + Milvus hard negative mining.
"""

import os
import copy
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models as models
import torchvision.transforms as T
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR
from tqdm import tqdm
from pymilvus import MilvusClient, DataType

# -- Config --
EMBED_DIM        = 512
NUM_EPOCHS       = 30
BATCH_SIZE       = 64
REINDEX_EVERY    = 5
LEARNING_RATE    = 3e-4
WEIGHT_DECAY     = 1e-4
TRIPLET_MARGIN   = 0.3
MILVUS_URI       = "./training_embeddings.db"   # Milvus Lite
COLLECTION_NAME  = "encoder_embeddings"
DATA_ROOT        = "./data"
CHECKPOINT_DIR   = "./checkpoints"
DEVICE           = torch.device("cuda" if torch.cuda.is_available() else "cpu")

os.makedirs(CHECKPOINT_DIR, exist_ok=True)


# -- Model --
class ImageEncoder(nn.Module):
    def __init__(self, embed_dim=512):
        super().__init__()
        backbone = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)
        self.backbone  = nn.Sequential(*list(backbone.children())[:-1])
        self.projector = nn.Sequential(
            nn.Linear(2048, 2048), nn.BatchNorm1d(2048), nn.ReLU(inplace=True),
            nn.Linear(2048, embed_dim), nn.BatchNorm1d(embed_dim)
        )

    def forward(self, x):
        x = self.backbone(x).flatten(1)
        return F.normalize(self.projector(x), p=2, dim=1)


# -- Loss --
class TripletLoss(nn.Module):
    def __init__(self, margin=0.3):
        super().__init__()
        self.margin = margin

    def forward(self, a, p, n):
        return F.relu(1 - (a * p).sum(1) - (1 - (a * n).sum(1)) + self.margin).mean()


# -- Milvus helpers --
def setup_milvus(uri, collection_name, embed_dim):
    client = MilvusClient(uri=uri)
    if client.has_collection(collection_name):
        client.drop_collection(collection_name)
    schema = client.create_schema(auto_id=True, enable_dynamic_field=True)
    schema.add_field("id",        DataType.INT64,        is_primary=True, auto_id=True)
    schema.add_field("embedding", DataType.FLOAT_VECTOR, dim=embed_dim)
    schema.add_field("label",     DataType.INT32)
    schema.add_field("split",     DataType.VARCHAR,      max_length=16)
    client.create_collection(collection_name, schema=schema)
    idx = client.prepare_index_params()
    idx.add_index("embedding", "HNSW", "IP", params={"M": 16, "efConstruction": 200})
    client.create_index(collection_name, idx)
    client.load_collection(collection_name)
    return client


def embed_and_store(model, loader, split, client, collection_name, device):
    model.eval()
    buf = []
    with torch.no_grad():
        for imgs, lbls in tqdm(loader, desc=f"Indexing [{split}]"):
            embs = model(imgs.to(device)).cpu().numpy()
            for e, l in zip(embs, lbls.numpy()):
                buf.append({"embedding": e.tolist(), "label": int(l), "split": split})
            if len(buf) >= 1000:
                client.insert(collection_name, buf); buf.clear()
    if buf:
        client.insert(collection_name, buf)
    client.flush(collection_name)


def mine_negatives(client, collection_name, embs, labels, k=32):
    results = client.search(
        collection_name, embs.tolist(), "embedding",
        {"metric_type": "IP", "params": {"ef": 64}},
        limit=k, output_fields=["embedding", "label"],
        filter='split == "train"'
    )
    negs = np.zeros_like(embs)
    for i, (hits, lbl) in enumerate(zip(results, labels)):
        for h in hits:
            if h["entity"]["label"] != int(lbl):
                negs[i] = np.array(h["entity"]["embedding"], dtype=np.float32)
                break
        else:
            negs[i] = np.array(hits[-1]["entity"]["embedding"], dtype=np.float32)
    return negs


# -- Main --
def main():
    # Data
    train_tf = T.Compose([
        T.RandomResizedCrop(224), T.RandomHorizontalFlip(),
        T.ColorJitter(0.4, 0.4, 0.4, 0.1),
        T.ToTensor(), T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    val_tf = T.Compose([
        T.Resize(256), T.CenterCrop(224),
        T.ToTensor(), T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    train_ds = ImageFolder(f"{DATA_ROOT}/train", transform=train_tf)
    val_ds   = ImageFolder(f"{DATA_ROOT}/val",   transform=val_tf)
    train_dl = DataLoader(train_ds, BATCH_SIZE, shuffle=True,  num_workers=4, pin_memory=True, drop_last=True)
    val_dl   = DataLoader(val_ds,  256,         shuffle=False, num_workers=4, pin_memory=True)

    # Model + Milvus
    model  = ImageEncoder(EMBED_DIM).to(DEVICE)
    client = setup_milvus(MILVUS_URI, COLLECTION_NAME, EMBED_DIM)
    embed_and_store(model, train_dl, "train", client, COLLECTION_NAME, DEVICE)

    # Optimizer
    opt     = AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY)
    sched   = CosineAnnealingLR(opt, T_max=NUM_EPOCHS)
    loss_fn = TripletLoss(TRIPLET_MARGIN)

    for epoch in range(1, NUM_EPOCHS + 1):
        model.train()
        total_loss = 0.0

        for imgs, lbls in tqdm(train_dl, desc=f"Epoch {epoch:3d}"):
            imgs = imgs.to(DEVICE)
            embs = model(imgs)
            emb_np, lbl_np = embs.detach().cpu().numpy(), lbls.numpy()

            # Hard negative mining via Milvus
            neg_np = mine_negatives(client, COLLECTION_NAME, emb_np, lbl_np)
            neg = F.normalize(torch.tensor(neg_np, device=DEVICE), p=2, dim=1)

            # Build positives from same-class shuffling
            pos_idx = torch.arange(len(lbls))
            for lbl in lbls.unique():
                m = (lbls == lbl).nonzero(as_tuple=True)[0]
                if len(m) > 1:
                    pos_idx[m] = m[torch.randperm(len(m))]
            pos = embs[pos_idx]

            loss = loss_fn(embs, pos, neg)
            opt.zero_grad(); loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
            total_loss += loss.item()

        sched.step()
        print(f"Epoch {epoch:3d} | Loss={total_loss/len(train_dl):.4f}")

        if epoch % REINDEX_EVERY == 0:
            print("Re-indexing training set...")
            client.delete(COLLECTION_NAME, filter='split == "train"')
            client.flush(COLLECTION_NAME)
            embed_and_store(model, train_dl, "train", client, COLLECTION_NAME, DEVICE)

        torch.save({"epoch": epoch, "model": model.state_dict()},
                   f"{CHECKPOINT_DIR}/encoder_{epoch:03d}.pt")

    print("Training complete.")


if __name__ == "__main__":
    main()

Summary

This guide walked through the complete lifecycle of combining PyTorch and Milvus:

  • Milvus acts as an external memory that stores your model’s evolving embedding space and enables fast approximate nearest neighbor retrieval during training.
  • PyTorch handles the gradient-based optimization and forward/backward passes.
  • Hard negative mining via Milvus dramatically accelerates metric learning convergence compared to random negatives.
  • Periodic re-indexing keeps the stored embeddings consistent with the changing model weights.
  • The same Milvus collection used during training transitions seamlessly into a production search index after training completes.

For further reading, see: