Targeted Fine-Tuning from Model Mispredictions with Milvus

Overview and Architecture

This guide walks through the full pipeline for identifying images your model predicts incorrectly, storing and querying them using Milvus as the vector database, re-annotating them with correct labels, extending your training dataset, and fine-tuning the model so it learns from its own mistakes.

Milvus plays a central role throughout: it stores dense float embeddings from your model’s penultimate layer alongside scalar metadata (ground truth label, predicted label, confidence, annotation status), and lets you query them with both scalar filters and vector search for finding the most diverse and informative failure cases.

Why Milvus for This Pipeline?

Traditional approaches store predictions in flat CSV files and run pandas queries to find errors. This breaks down when your dataset grows beyond a few hundred thousand images, when you want to find visually similar hard examples rather than just statistically similar ones, when you need to deduplicate hard examples to avoid adding hundreds of near-identical failure images, and when you are running multiple fine-tuning iterations and need to track annotation state per sample across iterations.

Milvus solves all of these. Scalar filters handle rule-based selection and ANN search handles embedding-space diversity sampling, both against the same store.

Full Data Flow

flowchart TD
    A([Images]) --> B[Model Inference]
    B --> C1[(Milvus Collection)]
    B --> C2[Embeddings + Metadata]
    C2 --> C1
    C1 --> D{Scalar Filter}
    D -->|correct == False| E[Ranked Error Set]
    E --> F[ANN Diversity Search]
    F --> G[Hard Example Subset]
    G --> H[Export to Label Studio]
    H --> I[Human Annotation]
    I --> J[Write Corrected Labels to Milvus]
    J --> K[Query Annotated Samples]
    K --> L[Build Fine-Tuning Dataset]
    L --> M[Fine-Tune Model]
    M --> N[Re-index Embeddings]
    N --> C1

Loop Iteration Flow

flowchart LR
    T0[Loop 0: Train Model] --> I0[Infer on Eval Set]
    I0 --> E0[Find Errors in Milvus]
    E0 --> A0[Annotate Hard Examples]
    A0 --> T1[Loop 1: Fine-Tune]
    T1 --> I1[Re-index with New Model]
    I1 --> E1[Find Remaining Errors]
    E1 --> A1[Annotate New Hard Examples]
    A1 --> T2[Loop 2: Fine-Tune Again]

Prerequisites and Setup

Required Packages

pip install pymilvus torch torchvision transformers \
            Pillow tqdm pandas numpy scikit-learn \
            label-studio

Start Milvus with Docker Compose

The easiest way to run Milvus locally is with the official standalone Docker Compose:

curl -sfL https://raw.githubusercontent.com/milvus-io/milvus/master/scripts/standalone_embed.sh \
  | bash -s start

docker ps | grep milvus

By default Milvus listens on localhost:19530. For production deployments refer to the Milvus operator docs.

Verify the Connection

from pymilvus import connections

connections.connect(host="localhost", port="19530")
print("Connected to Milvus")

Step 1 — Milvus Schema Design

Design the collection schema before inserting any data. The schema captures everything needed to query, curate, and track samples across loop iterations.

Field Definitions

Field	Type	Notes
`id`	INT64 (primary key, auto)	Unique row ID
`filepath`	VARCHAR(512)	Absolute path to the image file
`embedding`	FLOAT_VECTOR(dim)	Penultimate-layer embedding from the model
`gt_label`	VARCHAR(128)	Ground-truth class name
`pred_label`	VARCHAR(128)	Model predicted class name
`confidence`	FLOAT	Top-1 softmax probability
`correct`	BOOL	Whether prediction matches ground truth
`split`	VARCHAR(32)	One of `train`, `val`, `eval`, or `pool`
`annotated`	BOOL	Whether a human has reviewed this sample
`corrected_label`	VARCHAR(128)	Human-corrected label, empty until annotated
`loop_iter`	INT32	Which fine-tuning loop iteration indexed this sample

Schema Lifecycle

stateDiagram-v2
    [*] --> Indexed: run_inference_and_index()
    Indexed --> QueryResult: scalar filter query
    QueryResult --> StagedForAnnotation: copy to staging dir
    StagedForAnnotation --> Annotated: human labels in Label Studio
    Annotated --> UpdatedInMilvus: delete + re-insert with corrected_label
    UpdatedInMilvus --> TrainingData: pulled into fine-tune dataset
    TrainingData --> [*]

Create the Collection

from pymilvus import (
    connections, FieldSchema, CollectionSchema,
    DataType, Collection, utility
)

COLLECTION_NAME = "image_predictions"
EMBEDDING_DIM   = 2048   # match your model output size (ResNet-50 = 2048)

connections.connect(host="localhost", port="19530")

def create_collection(name: str, dim: int) -> Collection:
    if utility.has_collection(name):
        print(f"Collection '{name}' already exists -- loading.")
        return Collection(name)

    fields = [
        FieldSchema(name="id",              dtype=DataType.INT64,        is_primary=True, auto_id=True),
        FieldSchema(name="filepath",        dtype=DataType.VARCHAR,      max_length=512),
        FieldSchema(name="embedding",       dtype=DataType.FLOAT_VECTOR, dim=dim),
        FieldSchema(name="gt_label",        dtype=DataType.VARCHAR,      max_length=128),
        FieldSchema(name="pred_label",      dtype=DataType.VARCHAR,      max_length=128),
        FieldSchema(name="confidence",      dtype=DataType.FLOAT),
        FieldSchema(name="correct",         dtype=DataType.BOOL),
        FieldSchema(name="split",           dtype=DataType.VARCHAR,      max_length=32),
        FieldSchema(name="annotated",       dtype=DataType.BOOL),
        FieldSchema(name="corrected_label", dtype=DataType.VARCHAR,      max_length=128),
        FieldSchema(name="loop_iter",       dtype=DataType.INT32),
    ]

    schema     = CollectionSchema(fields, description="Image prediction records")
    collection = Collection(name, schema)

    index_params = {
        "metric_type": "L2",
        "index_type":  "IVF_FLAT",
        "params":      {"nlist": 256},
    }
    collection.create_index("embedding", index_params)
    collection.load()

    print(f"Created collection '{name}' with dim={dim}")
    return collection

collection = create_collection(COLLECTION_NAME, EMBEDDING_DIM)

Note

Choosing nlist: A value of 256 is reasonable for datasets up to around 500k images. For larger datasets use nlist = 4 * sqrt(N) as a rough heuristic.

Step 2 — Run Inference and Index Embeddings

During inference, extract both the prediction (class and confidence) and the embedding vector from the penultimate layer. Insert both into Milvus in batches.

Hook-Based Embedding Extraction

flowchart LR
    A[Input Image Batch] --> B[Backbone Layers]
    B --> C[avgpool Layer]
    C -->|forward hook| D[Embedding Vector 2048-d]
    C --> E[Classifier Head]
    E --> F[Logits / Softmax]
    D --> G[(Milvus Insert)]
    F --> G

import torch
import torch.nn.functional as F
from torchvision import transforms, datasets
from torch.utils.data import DataLoader
from pymilvus import Collection
from tqdm import tqdm
import numpy as np

DEVICE         = "cuda" if torch.cuda.is_available() else "cpu"
BATCH_SIZE     = 64
LOOP_ITER      = 0
SPLIT          = "eval"

transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])

def load_model(checkpoint_path: str, device: str):
    model = torch.load(checkpoint_path, map_location=device)
    model.eval()
    return model

def register_embedding_hook(model):
    embeddings = {}

    def hook(module, input, output):
        embeddings["vector"] = output.flatten(1).detach().cpu().numpy()

    handle = model.avgpool.register_forward_hook(hook)
    return embeddings, handle

def run_inference_and_index(
    model,
    data_dir: str,
    collection: Collection,
    split: str,
    loop_iter: int,
    device: str,
    batch_size: int = 64,
):
    dataset   = datasets.ImageFolder(data_dir, transform=transform)
    loader    = DataLoader(dataset, batch_size=batch_size,
                           shuffle=False, num_workers=4)
    idx2cls   = {v: k for k, v in dataset.class_to_idx.items()}
    emb_store, handle = register_embedding_hook(model)

    insert_buffer = {
        "filepath":        [],
        "embedding":       [],
        "gt_label":        [],
        "pred_label":      [],
        "confidence":      [],
        "correct":         [],
        "split":           [],
        "annotated":       [],
        "corrected_label": [],
        "loop_iter":       [],
    }
    FLUSH_EVERY = 512

    def flush(buf):
        collection.insert([
            buf["filepath"],
            buf["embedding"],
            buf["gt_label"],
            buf["pred_label"],
            buf["confidence"],
            buf["correct"],
            buf["split"],
            buf["annotated"],
            buf["corrected_label"],
            buf["loop_iter"],
        ])
        for v in buf.values():
            v.clear()

    with torch.no_grad():
        for batch_idx, (images, labels) in enumerate(tqdm(loader, desc="Indexing")):
            _      = model(images.to(device))
            embs   = emb_store["vector"]
            probs  = F.softmax(model(images.to(device)), dim=1).cpu().numpy()
            top_pred = probs.argmax(axis=1)
            top_conf = probs.max(axis=1)

            start = batch_idx * batch_size
            for j in range(len(images)):
                img_path = dataset.samples[start + j][0]
                gt_idx   = labels[j].item()
                pred_idx = int(top_pred[j])
                correct  = gt_idx == pred_idx

                insert_buffer["filepath"].append(img_path)
                insert_buffer["embedding"].append(embs[j].tolist())
                insert_buffer["gt_label"].append(idx2cls[gt_idx])
                insert_buffer["pred_label"].append(idx2cls[pred_idx])
                insert_buffer["confidence"].append(float(top_conf[j]))
                insert_buffer["correct"].append(correct)
                insert_buffer["split"].append(split)
                insert_buffer["annotated"].append(False)
                insert_buffer["corrected_label"].append("")
                insert_buffer["loop_iter"].append(loop_iter)

            if len(insert_buffer["filepath"]) >= FLUSH_EVERY:
                flush(insert_buffer)

    if insert_buffer["filepath"]:
        flush(insert_buffer)

    handle.remove()
    collection.flush()
    print(f"Indexed {collection.num_entities} total entities in '{collection.name}'")

model      = load_model("checkpoints/model_best.pt", DEVICE)
collection = Collection(COLLECTION_NAME)
collection.load()

run_inference_and_index(
    model      = model,
    data_dir   = "data/eval",
    collection = collection,
    split      = SPLIT,
    loop_iter  = LOOP_ITER,
    device     = DEVICE,
)

Important

Always call collection.flush() after inserting a batch. Milvus buffers inserts in memory before persisting to disk and unflushed data will not be queryable.

Step 3 — Query Mispredicted Samples

Use Milvus scalar filters to retrieve exactly the samples that were predicted incorrectly. No pandas, no CSV – just a query expression.

Query Decision Tree

flowchart TD
    A[All Indexed Samples] --> B{correct == False?}
    B -->|No| C[Skip - Correctly Predicted]
    B -->|Yes| D{confidence > threshold?}
    D -->|No - Low Confidence| E[Uncertain Errors]
    D -->|Yes - High Confidence| F[Confident Wrong Predictions]
    F --> G{annotated == False?}
    G -->|Already Done| H[Skip - Already Reviewed]
    G -->|Not Yet| I[Add to Annotation Queue]
    E --> G

All Wrong Predictions

from pymilvus import Collection

collection = Collection(COLLECTION_NAME)
collection.load()

results = collection.query(
    expr          = 'correct == False and split == "eval"',
    output_fields = ["id", "filepath", "gt_label", "pred_label", "confidence", "loop_iter"],
    limit         = 16384,
)

print(f"Total mispredictions: {len(results)}")

High-Confidence Wrong Predictions

confident_errors = collection.query(
    expr          = 'correct == False and confidence > 0.85 and split == "eval"',
    output_fields = ["id", "filepath", "gt_label", "pred_label", "confidence"],
    limit         = 4096,
)
confident_errors.sort(key=lambda r: r["confidence"], reverse=True)

print(f"High-confidence errors (>0.85): {len(confident_errors)}")
for r in confident_errors[:10]:
    print(f"  {r['gt_label']:20s} predicted as {r['pred_label']:20s}  conf={r['confidence']:.4f}")

Per-Class Confusion Breakdown

from collections import Counter

all_errors = collection.query(
    expr          = 'correct == False and split == "eval"',
    output_fields = ["gt_label", "pred_label"],
    limit         = 16384,
)

confusion_pairs = Counter(
    (r["gt_label"], r["pred_label"]) for r in all_errors
)

print("Top confused class pairs:")
for (gt, pred), count in confusion_pairs.most_common(15):
    print(f"  {gt:20s} -> {pred:20s}  ({count}x)")

Targeted Query for a Specific Error Pair

target_gt   = "cat"
target_pred = "dog"

focused = collection.query(
    expr = (
        f'correct == False '
        f'and gt_label == "{target_gt}" '
        f'and pred_label == "{target_pred}" '
        f'and annotated == False'
    ),
    output_fields = ["id", "filepath", "confidence"],
    limit         = 1000,
)
print(f"Unannotated {target_gt}->{target_pred} errors: {len(focused)}")

Step 4 — Rank and Diversify via ANN Search

Raw scalar filtering gives you all errors, but for annotation efficiency you want a diverse, representative subset – not 200 near-identical photos of the same dog in bad lighting.

Diversity Selection Strategy

flowchart TD
    A[All Error Embeddings] --> B[Compute Centroid]
    B --> C[ANN Search from Centroid]
    C --> D[Over-fetch Budget x 3 Candidates]
    D --> E[Greedy Furthest-Point Sampling]
    E --> F{Budget Reached?}
    F -->|No| G[Pick Candidate Furthest from Selected Set]
    G --> E
    F -->|Yes| H[Final Diverse Subset]

Compute the Error-Cluster Centroid

import numpy as np
from pymilvus import Collection

errors_with_emb = collection.query(
    expr          = 'correct == False and split == "eval" and annotated == False',
    output_fields = ["id", "filepath", "embedding", "gt_label", "pred_label", "confidence"],
    limit         = 16384,
)

embeddings = np.array([r["embedding"] for r in errors_with_emb])
centroid   = embeddings.mean(axis=0)

ANN Search from the Centroid

BUDGET = 500

search_params = {"metric_type": "L2", "params": {"nprobe": 32}}

diverse_hits = collection.search(
    data          = [centroid.tolist()],
    anns_field    = "embedding",
    param         = search_params,
    limit         = BUDGET * 3,
    expr          = 'correct == False and annotated == False',
    output_fields = ["id", "filepath", "gt_label", "pred_label", "confidence"],
)

candidates = [hit for hit in diverse_hits[0]]
print(f"Candidates retrieved: {len(candidates)}")

Greedy Furthest-Point Sampling

def greedy_diverse_select(candidates, embeddings_map: dict, budget: int):
    selected      = [candidates[0]]
    selected_embs = [embeddings_map[candidates[0].id]]
    remaining     = candidates[1:]

    while len(selected) < budget and remaining:
        best_idx  = -1
        best_dist = -1.0

        for i, cand in enumerate(remaining):
            emb      = embeddings_map[cand.id]
            min_dist = min(
                np.linalg.norm(emb - s) for s in selected_embs
            )
            if min_dist > best_dist:
                best_dist = min_dist
                best_idx  = i

        chosen = remaining.pop(best_idx)
        selected.append(chosen)
        selected_embs.append(embeddings_map[chosen.id])

    return selected

emb_map        = {r["id"]: np.array(r["embedding"]) for r in errors_with_emb}
diverse_subset = greedy_diverse_select(candidates, emb_map, budget=BUDGET)
print(f"Diverse hard examples selected: {len(diverse_subset)}")

Per-Class Quota with Hybrid Search

failing_classes  = list({r["gt_label"] for r in all_errors})
per_class_quota  = max(20, BUDGET // len(failing_classes))
selected_ids     = []

for cls in failing_classes:
    hits = collection.search(
        data          = [centroid.tolist()],
        anns_field    = "embedding",
        param         = search_params,
        limit         = per_class_quota,
        expr          = f'correct == False and annotated == False and gt_label == "{cls}"',
        output_fields = ["id", "filepath", "gt_label", "pred_label", "confidence"],
    )
    selected_ids.extend([h.id for h in hits[0]])

print(f"Selected {len(selected_ids)} hard examples across {len(failing_classes)} classes")

Step 5 — Export and Prepare for Annotation

Annotation Workflow

sequenceDiagram
    participant M as Milvus
    participant S as Staging Dir
    participant L as Label Studio
    participant A as Annotator

    M->>S: Copy selected images (id__class__filename.jpg)
    S->>L: Import label_studio_tasks.json
    L->>A: Present image + pre-filled GT suggestion
    A->>L: Confirm or override label
    L->>S: Export annotations_export.json
    S->>M: Write corrected_label back (delete + re-insert)

Fetch Full Records and Copy Images

import shutil
from pathlib import Path

id_expr = f"id in [{', '.join(str(i) for i in selected_ids)}]"

to_annotate = collection.query(
    expr          = id_expr,
    output_fields = ["id", "filepath", "gt_label", "pred_label", "confidence"],
)

STAGING_DIR = Path("hard_examples/images")
STAGING_DIR.mkdir(parents=True, exist_ok=True)

for rec in to_annotate:
    src  = Path(rec["filepath"])
    dest = STAGING_DIR / f"{rec['id']}__{rec['gt_label']}__{src.name}"
    shutil.copy2(src, dest)

print(f"Copied {len(to_annotate)} images to {STAGING_DIR}")

Generate Label Studio Import Manifest

import json

tasks = []
for rec in to_annotate:
    src       = Path(rec["filepath"])
    dest_name = f"{rec['id']}__{rec['gt_label']}__{src.name}"
    tasks.append({
        "data": {
            "image":      f"/data/local-files/?d=hard_examples/images/{dest_name}",
            "milvus_id":  rec["id"],
            "pred_label": rec["pred_label"],
            "confidence": rec["confidence"],
        },
        "predictions": [{
            "result": [{
                "type":  "choices",
                "value": {"choices": [rec["gt_label"]]},
            }]
        }]
    })

with open("hard_examples/label_studio_tasks.json", "w") as f:
    json.dump(tasks, f, indent=2)

print("Exported Label Studio tasks -> hard_examples/label_studio_tasks.json")

Launch Label Studio

label-studio start \
  --port 8080 \
  --data-dir ./hard_examples

After launching, create a project with Image Classification type, set the label list to match your training classes exactly, import label_studio_tasks.json, complete annotations, then export as JSON-MIN.

Warning

Annotation quality note: Do not tell annotators which images are model errors. Show them a mix of correct and incorrect predictions to prevent anchoring bias. For the highest-confidence wrong predictions, require two independent annotators and resolve disagreements before adding to the training set.

Step 6 — Re-annotate and Write Back to Milvus

After annotation, write the corrected labels back into Milvus. Because Milvus does not support in-place field updates, the standard pattern is to delete the old entity and re-insert it with the updated fields. Milvus 2.3+ supports native upsert() as a one-step alternative.

Milvus Update Pattern

flowchart LR
    A[annotations_export.json] --> B[Parse milvus_id + corrected_label]
    B --> C[Fetch Full Records by ID]
    C --> D[collection.delete by ID]
    D --> E[Re-insert with annotated=True and corrected_label set]
    E --> F[collection.flush]

Parse Label Studio Export

import json

with open("hard_examples/annotations_export.json") as f:
    ls_export = json.load(f)

corrections = []
for task in ls_export:
    milvus_id = task["data"]["milvus_id"]
    corrected = task["annotations"][0]["result"][0]["value"]["choices"][0]
    corrections.append({"milvus_id": int(milvus_id), "corrected_label": corrected})

print(f"Parsed {len(corrections)} annotations")

Delete and Re-Insert with Corrected Labels

from pymilvus import Collection

collection = Collection(COLLECTION_NAME)
collection.load()

corrected_ids = [c["milvus_id"] for c in corrections]
id_map        = {c["milvus_id"]: c["corrected_label"] for c in corrections}

old_records = collection.query(
    expr          = f"id in [{', '.join(str(i) for i in corrected_ids)}]",
    output_fields = ["id", "filepath", "embedding", "gt_label", "pred_label",
                     "confidence", "correct", "split", "loop_iter"],
)

collection.delete(f"id in [{', '.join(str(i) for i in corrected_ids)}]")

new_rows = {
    "filepath":        [],
    "embedding":       [],
    "gt_label":        [],
    "pred_label":      [],
    "confidence":      [],
    "correct":         [],
    "split":           [],
    "annotated":       [],
    "corrected_label": [],
    "loop_iter":       [],
}

for rec in old_records:
    corrected_label = id_map[rec["id"]]
    new_rows["filepath"].append(rec["filepath"])
    new_rows["embedding"].append(rec["embedding"])
    new_rows["gt_label"].append(rec["gt_label"])
    new_rows["pred_label"].append(rec["pred_label"])
    new_rows["confidence"].append(rec["confidence"])
    new_rows["correct"].append(rec["correct"])
    new_rows["split"].append(rec["split"])
    new_rows["annotated"].append(True)
    new_rows["corrected_label"].append(corrected_label)
    new_rows["loop_iter"].append(rec["loop_iter"])

collection.insert(list(new_rows.values()))
collection.flush()

print(f"Updated {len(old_records)} records in Milvus with corrected labels")

Step 7 — Build the Fine-Tuning Dataset from Milvus

Query Milvus to assemble exactly the data split you need for the next training run, mixing the original training samples with the newly annotated hard examples.

Dataset Assembly Flow

flowchart TD
    A[(Milvus Collection)] --> B{Query: split == train}
    A --> C{Query: annotated == True}
    B --> D[Original Training Records]
    C --> E[Annotated Hard Examples]
    D --> F[Merge and Deduplicate]
    E --> F
    F --> G[train_manifest_v2.csv]
    G --> H[ManifestDataset]
    H --> I[WeightedRandomSampler]
    I --> J[DataLoader for Fine-Tuning]

Query Original and Annotated Records

train_records = collection.query(
    expr          = 'split == "train"',
    output_fields = ["filepath", "gt_label"],
    limit         = 200000,
)
print(f"Original training samples: {len(train_records)}")

hard_records = collection.query(
    expr          = 'annotated == True and split == "eval"',
    output_fields = ["filepath", "corrected_label"],
    limit         = 16384,
)

for r in hard_records:
    r["gt_label"] = r.pop("corrected_label")

print(f"Annotated hard examples: {len(hard_records)}")

Merge into a Unified Manifest

import pandas as pd

df_train = pd.DataFrame([
    {"filepath": r["filepath"], "label": r["gt_label"]} for r in train_records
])

df_hard = pd.DataFrame([
    {"filepath": r["filepath"], "label": r["gt_label"]} for r in hard_records
])

df_merged = pd.concat([df_train, df_hard], ignore_index=True).drop_duplicates("filepath")

print(df_merged["label"].value_counts().to_string())
df_merged.to_csv("data/train_manifest_v2.csv", index=False)
print(f"Merged dataset: {len(df_merged)} samples")

PyTorch Dataset Backed by the Manifest

from torch.utils.data import Dataset
from torchvision import transforms
from PIL import Image
import pandas as pd

class ManifestDataset(Dataset):
    def __init__(self, manifest_path: str, transform=None):
        self.df        = pd.read_csv(manifest_path)
        self.transform = transform
        self.classes   = sorted(self.df["label"].unique().tolist())
        self.cls2idx   = {c: i for i, c in enumerate(self.classes)}

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        row   = self.df.iloc[idx]
        image = Image.open(row["filepath"]).convert("RGB")
        label = self.cls2idx[row["label"]]
        if self.transform:
            image = self.transform(image)
        return image, label

train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224, scale=(0.7, 1.0)),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(0.3, 0.3, 0.2),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])

train_dataset = ManifestDataset("data/train_manifest_v2.csv", transform=train_transform)

Class-Balanced Sampler from Milvus Counts

import numpy as np
from torch.utils.data import WeightedRandomSampler, DataLoader
from collections import Counter

class_counts_query = collection.query(
    expr          = 'split == "train"',
    output_fields = ["gt_label"],
    limit         = 200000,
)
class_counts = Counter(r["gt_label"] for r in class_counts_query)

for r in hard_records:
    class_counts[r["gt_label"]] += 1

sample_weights = [
    1.0 / class_counts[train_dataset.df.iloc[i]["label"]]
    for i in range(len(train_dataset))
]

sampler = WeightedRandomSampler(
    weights     = sample_weights,
    num_samples = len(train_dataset),
    replacement = True,
)

train_loader = DataLoader(
    train_dataset,
    batch_size  = 32,
    sampler     = sampler,
    num_workers = 4,
    pin_memory  = True,
)

Step 8 — Fine-Tune the Model

With the dataset assembled from Milvus, fine-tuning follows standard transfer learning practice with a lower learning rate and a frozen backbone warm-up phase to prevent catastrophic forgetting.

Training Phase Diagram

flowchart TD
    A[Load Checkpoint] --> B[Freeze Backbone]
    B --> C[Epochs 1 to FREEZE_EPOCHS]
    C --> D{Epoch == FREEZE_EPOCHS + 1?}
    D -->|Yes| E[Unfreeze Backbone, Reduce LR by 5x]
    D -->|No| F[Train Head Only]
    E --> G[Epochs FREEZE_EPOCHS+1 to END]
    F --> H{Val Acc Improved?}
    G --> H
    H -->|Yes| I[Save Best Checkpoint]
    H -->|No| J[Continue]
    I --> J
    J --> K{All Epochs Done?}
    K -->|No| C
    K -->|Yes| L[Done]

Fine-Tuning Script

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms, datasets
from torch.utils.data import DataLoader
from tqdm import tqdm

DEVICE         = "cuda" if torch.cuda.is_available() else "cpu"
CHECKPOINT_IN  = "checkpoints/model_best.pt"
CHECKPOINT_OUT = "checkpoints/model_finetuned_v2.pt"
EPOCHS         = 10
LR             = 1e-4
FREEZE_EPOCHS  = 3

val_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])

val_dataset = datasets.ImageFolder("data/val", transform=val_transform)
val_loader  = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=4)

model     = torch.load(CHECKPOINT_IN, map_location=DEVICE).to(DEVICE)
criterion = nn.CrossEntropyLoss()
optimizer = optim.AdamW(model.parameters(), lr=LR, weight_decay=1e-4)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS)

def set_backbone_grad(model, requires_grad: bool):
    for name, param in model.named_parameters():
        if not any(k in name for k in ["fc", "classifier", "head"]):
            param.requires_grad = requires_grad

set_backbone_grad(model, requires_grad=False)

best_val_acc = 0.0

for epoch in range(1, EPOCHS + 1):

    if epoch == FREEZE_EPOCHS + 1:
        print("Unfreezing backbone")
        set_backbone_grad(model, requires_grad=True)
        optimizer = optim.AdamW(model.parameters(), lr=LR / 5, weight_decay=1e-4)
        scheduler = optim.lr_scheduler.CosineAnnealingLR(
            optimizer, T_max=EPOCHS - FREEZE_EPOCHS
        )

    model.train()
    for images, labels in tqdm(train_loader, desc=f"E{epoch:02d} train"):
        images, labels = images.to(DEVICE), labels.to(DEVICE)
        optimizer.zero_grad()
        criterion(model(images), labels).backward()
        optimizer.step()
    scheduler.step()

    model.eval()
    correct = total = 0
    with torch.no_grad():
        for images, labels in val_loader:
            preds    = model(images.to(DEVICE)).argmax(1).cpu()
            correct += (preds == labels).sum().item()
            total   += labels.size(0)

    val_acc = correct / total
    print(f"Epoch {epoch:02d} | val_acc={val_acc:.4f}")

    if val_acc > best_val_acc:
        best_val_acc = val_acc
        torch.save(model, CHECKPOINT_OUT)
        print(f"  [SAVED] best checkpoint (val_acc={best_val_acc:.4f})")

print(f"Fine-tuning complete. Best val acc: {best_val_acc:.4f}")

Focal Loss for Hard Examples

import torch.nn.functional as F

class FocalLoss(nn.Module):
    def __init__(self, gamma=2.0):
        super().__init__()
        self.gamma = gamma

    def forward(self, inputs, targets):
        ce   = F.cross_entropy(inputs, targets, reduction="none")
        pt   = torch.exp(-ce)
        loss = ((1 - pt) ** self.gamma) * ce
        return loss.mean()

criterion = FocalLoss(gamma=2.0)

Step 9 — Evaluate and Update Milvus with New Predictions

After training, run inference with the new model and insert its predictions under the next loop_iter. This lets you compare results across iterations directly in Milvus without any external tracking system.

Cross-Iteration Comparison

flowchart LR
    A[(loop_iter == 0)] --> C[Query errors]
    B[(loop_iter == 1)] --> D[Query errors]
    C --> E[Per-class error count v0]
    D --> F[Per-class error count v1]
    E --> G[Delta Table]
    F --> G
    G --> H{delta < 0?}
    H -->|Yes| I[Regression Detected]
    H -->|No| J[Improvement Confirmed]

Error Count Comparison Across Iterations

errors_v0 = collection.query(
    expr          = 'correct == False and loop_iter == 0',
    output_fields = ["id", "gt_label", "pred_label", "confidence"],
    limit         = 16384,
)

errors_v1 = collection.query(
    expr          = 'correct == False and loop_iter == 1',
    output_fields = ["id", "gt_label", "pred_label", "confidence"],
    limit         = 16384,
)

print(f"Errors before fine-tuning (loop 0): {len(errors_v0)}")
print(f"Errors after fine-tuning  (loop 1): {len(errors_v1)}")
print(f"Reduction: {(1 - len(errors_v1)/len(errors_v0))*100:.1f}%")

Per-Class Delta

from collections import defaultdict

def error_rate_by_class(error_records):
    counts = defaultdict(int)
    for r in error_records:
        counts[r["gt_label"]] += 1
    return counts

er_v0 = error_rate_by_class(errors_v0)
er_v1 = error_rate_by_class(errors_v1)

all_classes = set(er_v0) | set(er_v1)
rows = sorted(
    [{"class": c, "v0": er_v0.get(c, 0), "v1": er_v1.get(c, 0)} for c in all_classes],
    key=lambda r: r["v0"] - r["v1"],
    reverse=True,
)

print(f"{'Class':<25} {'Errors v0':>10} {'Errors v1':>10} {'Delta':>10}")
print("-" * 55)
for r in rows:
    delta = r["v0"] - r["v1"]
    flag  = "[REGRESSION]" if delta < 0 else ""
    print(f"{r['class']:<25} {r['v0']:>10} {r['v1']:>10} {delta:>+10}  {flag}")

Verify Previously Hard Examples Were Fixed

prev_hard = collection.query(
    expr          = 'correct == False and loop_iter == 0 and annotated == True',
    output_fields = ["id", "filepath", "embedding"],
    limit         = 1000,
)

fixed = 0
still_wrong = 0

for rec in prev_hard:
    match = collection.query(
        expr          = f'filepath == "{rec["filepath"]}" and loop_iter == 1',
        output_fields = ["correct"],
        limit         = 1,
    )
    if match and match[0]["correct"]:
        fixed += 1
    else:
        still_wrong += 1

print(f"Previously hard examples now correct: {fixed}")
print(f"Previously hard examples still wrong: {still_wrong}")

Step 10 — Closing the Active Learning Loop

With the new model checkpointed and its predictions indexed under a new loop_iter, the loop is ready to repeat. Because all history lives in Milvus, you can reconstruct what was annotated in which iteration, diagnose regressions, and audit annotation quality at any time.

Loop Controller

flowchart TD
    A[Start Loop Iteration N] --> B[Query errors where loop_iter == N-1]
    B --> C{Any unannotated errors?}
    C -->|No| D[Loop Complete - Model Converged]
    C -->|Yes| E[ANN Diversity Selection]
    E --> F[Export to Label Studio]
    F --> G[Human Annotation]
    G --> H[Write Back to Milvus]
    H --> I[Build Dataset from Milvus]
    I --> J[Fine-Tune Model]
    J --> K[Re-index with New Model at loop_iter == N]
    K --> A

def active_learning_loop(
    collection,
    annotation_budget: int,
    loop_iter: int,
    model_checkpoint_in: str,
    model_checkpoint_out: str,
):
    print(f"=== Loop iteration {loop_iter} ===")

    errors = collection.query(
        expr          = f'correct == False and loop_iter == {loop_iter - 1} and annotated == False',
        output_fields = ["id", "filepath", "embedding", "gt_label", "pred_label", "confidence"],
        limit         = 16384,
    )

    if not errors:
        print("No unannotated errors found -- loop complete.")
        return

    emb_map  = {r["id"]: np.array(r["embedding"]) for r in errors}
    centroid = np.stack(list(emb_map.values())).mean(axis=0)

    # run ANN search + greedy selection (see Step 4)

    # export for annotation (see Step 5)

    # HUMAN ANNOTATION HAPPENS HERE

    # parse annotations and write back to Milvus (see Step 6)

    # build dataset and fine-tune (see Steps 7-8)

    run_inference_and_index(
        model      = torch.load(model_checkpoint_out),
        data_dir   = "data/eval",
        collection = collection,
        split      = "eval",
        loop_iter  = loop_iter,
        device     = DEVICE,
    )

    print(f"Loop {loop_iter} complete.")

Advanced Techniques

Milvus Partitions Per Loop Iteration

For large datasets, use Milvus partitions to isolate loop iterations and speed up queries:

from pymilvus import Collection

collection = Collection(COLLECTION_NAME)

for i in range(5):
    partition_name = f"loop_{i}"
    if not collection.has_partition(partition_name):
        collection.create_partition(partition_name)

collection.insert(data_rows, partition_name="loop_0")

results = collection.query(
    expr            = 'correct == False',
    output_fields   = ["filepath", "confidence"],
    partition_names = ["loop_0"],
    limit           = 4096,
)

Pseudo-Labelling High-Confidence Unlabelled Images

pseudo_labels = collection.query(
    expr          = 'split == "pool" and confidence > 0.98',
    output_fields = ["id", "filepath", "pred_label", "embedding"],
    limit         = 50000,
)

# Re-insert as auto-annotated training samples
for rec in pseudo_labels:
    rec["gt_label"]        = rec.pop("pred_label")
    rec["annotated"]       = True
    rec["corrected_label"] = rec["gt_label"]
    rec["split"]           = "train"
    # insert back into collection

Warning

Apply strict thresholds (0.98+) for pseudo-labelling and audit a random 5% sample to catch systematic errors before they corrupt the training set.

Semantic Near-Duplicate Detection via ANN

Before adding hard examples to training, use Milvus to check whether a near-identical image is already in the training set:

def is_near_duplicate(embedding, collection, threshold=0.05):
    hits = collection.search(
        data       = [embedding],
        anns_field = "embedding",
        param      = {"metric_type": "L2", "params": {"nprobe": 16}},
        limit      = 1,
        expr       = 'split == "train"',
    )
    return hits[0][0].distance < threshold if hits[0] else False

deduped = [
    r for r in to_annotate
    if not is_near_duplicate(r["embedding"], collection)
]
print(f"After dedup: {len(deduped)} / {len(to_annotate)} remain")

Tracking Label Noise Across Iterations

Use Milvus to flag training samples that the model consistently gets wrong across multiple loop iterations – these are likely mislabelled originals:

filepath_error_count = {}

for loop in range(LOOP_ITER):
    errors = collection.query(
        expr          = f'correct == False and loop_iter == {loop} and split == "train"',
        output_fields = ["filepath"],
        limit         = 50000,
    )
    for r in errors:
        filepath_error_count[r["filepath"]] = filepath_error_count.get(r["filepath"], 0) + 1

noise_candidates = [
    fp for fp, count in filepath_error_count.items()
    if count == LOOP_ITER
]
print(f"Potential label noise candidates: {len(noise_candidates)}")

Paginating Large Result Sets

Milvus query() has a hard limit parameter. Paginate with offset when you expect more results than a single call allows:

all_results = []
page_size   = 16384
offset      = 0

while True:
    page = collection.query(
        expr          = 'correct == False',
        output_fields = ["id", "filepath"],
        limit         = page_size,
        offset        = offset,
    )
    if not page:
        break
    all_results.extend(page)
    offset += page_size

print(f"Total results: {len(all_results)}")

Common Pitfalls

Schema Immutability

flowchart LR
    A[Need New Field?] --> B{Schema Already Created?}
    B -->|No| C[Add field in FieldSchema list before create_collection]
    B -->|Yes| D[Create new collection with updated schema]
    D --> E[Re-index all data from old collection to new]
    E --> F[Drop old collection]

Milvus collection schemas are fixed at creation time. You cannot add or rename fields after the fact. Plan your schema carefully upfront, or use the migration pattern shown above.

Collection Must Be Loaded Before Querying

# Always load before any query or search
collection = Collection(COLLECTION_NAME)
collection.load()   # required

Always Flush After Insert

collection.insert(rows)
collection.flush()   # do not skip -- unflushed data is not queryable

Avoid Catastrophic Forgetting

Always mix hard examples with original training data. A safe starting ratio is one hard example for every three original training samples. Never fine-tune exclusively on the hard-example set.

End-to-End Script

"""
milvus_active_learning.py

Full pipeline:
  1. Connect to Milvus and create collection
  2. Run inference and index embeddings + predictions
  3. Query hard examples with scalar filter
  4. Diverse selection via ANN search
  5. Export for annotation (Label Studio)
  6. Write corrected labels back to Milvus
  7. Build fine-tuning dataset from Milvus
  8. Fine-tune model
  9. Re-index with new model and compare results
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np
import pandas as pd
import shutil
import json
from pathlib import Path
from collections import Counter
from tqdm import tqdm
from torchvision import transforms, datasets
from torch.utils.data import DataLoader, WeightedRandomSampler, Dataset
from PIL import Image
from sklearn.metrics import classification_report

from pymilvus import (
    connections, FieldSchema, CollectionSchema,
    DataType, Collection, utility
)

# ----------------------------------------------------------------
# CONFIG
# ----------------------------------------------------------------
CFG = {
    "milvus_host":       "localhost",
    "milvus_port":       "19530",
    "collection":        "image_predictions",
    "embedding_dim":     2048,
    "checkpoint_in":     "checkpoints/model_best.pt",
    "checkpoint_out":    "checkpoints/model_finetuned.pt",
    "eval_dir":          "data/eval",
    "train_dir":         "data/train",
    "val_dir":           "data/val",
    "staging_dir":       "hard_examples",
    "ls_export":         "hard_examples/annotations_export.json",
    "annotation_budget": 500,
    "conf_threshold":    0.75,
    "loop_iter":         0,
    "batch_size":        32,
    "epochs":            10,
    "lr":                1e-4,
    "freeze_epochs":     3,
    "device":            "cuda" if torch.cuda.is_available() else "cpu",
}

# ----------------------------------------------------------------
# 1. MILVUS SETUP
# ----------------------------------------------------------------
connections.connect(host=CFG["milvus_host"], port=CFG["milvus_port"])

if not utility.has_collection(CFG["collection"]):
    fields = [
        FieldSchema("id",              DataType.INT64,        is_primary=True, auto_id=True),
        FieldSchema("filepath",        DataType.VARCHAR,      max_length=512),
        FieldSchema("embedding",       DataType.FLOAT_VECTOR, dim=CFG["embedding_dim"]),
        FieldSchema("gt_label",        DataType.VARCHAR,      max_length=128),
        FieldSchema("pred_label",      DataType.VARCHAR,      max_length=128),
        FieldSchema("confidence",      DataType.FLOAT),
        FieldSchema("correct",         DataType.BOOL),
        FieldSchema("split",           DataType.VARCHAR,      max_length=32),
        FieldSchema("annotated",       DataType.BOOL),
        FieldSchema("corrected_label", DataType.VARCHAR,      max_length=128),
        FieldSchema("loop_iter",       DataType.INT32),
    ]
    schema = CollectionSchema(fields)
    col    = Collection(CFG["collection"], schema)
    col.create_index("embedding", {"metric_type": "L2", "index_type": "IVF_FLAT",
                                    "params": {"nlist": 256}})
else:
    col = Collection(CFG["collection"])

col.load()
print(f"Milvus ready. Entities: {col.num_entities}")

# ----------------------------------------------------------------
# 2. INFERENCE + INDEX
# ----------------------------------------------------------------
transform_eval = transforms.Compose([
    transforms.Resize(256), transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])

model   = torch.load(CFG["checkpoint_in"], map_location=CFG["device"]).eval()
dataset = datasets.ImageFolder(CFG["eval_dir"], transform=transform_eval)
loader  = DataLoader(dataset, batch_size=64, shuffle=False, num_workers=4)
idx2cls = {v: k for k, v in dataset.class_to_idx.items()}

emb_store = {}
handle    = model.avgpool.register_forward_hook(
    lambda m, i, o: emb_store.update({"v": o.flatten(1).detach().cpu().numpy()})
)

buf = {k: [] for k in ["filepath", "embedding", "gt_label", "pred_label",
                         "confidence", "correct", "split", "annotated",
                         "corrected_label", "loop_iter"]}

with torch.no_grad():
    for bi, (imgs, labels) in enumerate(tqdm(loader, desc="Index")):
        _      = model(imgs.to(CFG["device"]))
        logits = model(imgs.to(CFG["device"]))
        probs  = F.softmax(logits, dim=1).cpu().numpy()
        top_p  = probs.argmax(1)
        top_c  = probs.max(1)
        start  = bi * loader.batch_size

        for j in range(len(imgs)):
            buf["filepath"].append(dataset.samples[start + j][0])
            buf["embedding"].append(emb_store["v"][j].tolist())
            buf["gt_label"].append(idx2cls[labels[j].item()])
            buf["pred_label"].append(idx2cls[int(top_p[j])])
            buf["confidence"].append(float(top_c[j]))
            buf["correct"].append(labels[j].item() == int(top_p[j]))
            buf["split"].append("eval")
            buf["annotated"].append(False)
            buf["corrected_label"].append("")
            buf["loop_iter"].append(CFG["loop_iter"])

        if len(buf["filepath"]) >= 512:
            col.insert(list(buf.values()))
            col.flush()
            for v in buf.values():
                v.clear()

if buf["filepath"]:
    col.insert(list(buf.values()))
    col.flush()

handle.remove()
print(f"Indexed. Total entities: {col.num_entities}")

# ----------------------------------------------------------------
# 3. QUERY HARD EXAMPLES
# ----------------------------------------------------------------
errors = col.query(
    expr=(
        f'correct == False and confidence > {CFG["conf_threshold"]} '
        f'and split == "eval" and loop_iter == {CFG["loop_iter"]}'
    ),
    output_fields=["id", "filepath", "embedding", "gt_label", "pred_label", "confidence"],
    limit=16384,
)
print(f"Hard examples: {len(errors)}")

# ----------------------------------------------------------------
# 4. DIVERSE SELECTION VIA ANN
# ----------------------------------------------------------------
embs     = np.array([r["embedding"] for r in errors])
centroid = embs.mean(axis=0)

hits = col.search(
    data=[centroid.tolist()],
    anns_field="embedding",
    param={"metric_type": "L2", "params": {"nprobe": 32}},
    limit=CFG["annotation_budget"] * 3,
    expr=f'correct == False and annotated == False and loop_iter == {CFG["loop_iter"]}',
    output_fields=["id", "filepath", "gt_label", "pred_label", "confidence"],
)[0]

emb_map   = {r["id"]: np.array(r["embedding"]) for r in errors}
selected  = [hits[0]]
sel_embs  = [emb_map.get(hits[0].id, centroid)]
remaining = list(hits[1:])

while len(selected) < CFG["annotation_budget"] and remaining:
    best_i, best_d = 0, -1.0
    for i, h in enumerate(remaining):
        e = emb_map.get(h.id, centroid)
        d = min(np.linalg.norm(e - s) for s in sel_embs)
        if d > best_d:
            best_d, best_i = d, i
    sel_embs.append(emb_map.get(remaining[best_i].id, centroid))
    selected.append(remaining.pop(best_i))

selected_ids = [h.id for h in selected]
print(f"Diverse subset: {len(selected_ids)} samples")

# ----------------------------------------------------------------
# 5. EXPORT FOR ANNOTATION
# ----------------------------------------------------------------
to_annotate = col.query(
    expr=f"id in [{','.join(str(i) for i in selected_ids)}]",
    output_fields=["id", "filepath", "gt_label", "pred_label", "confidence"],
)

staging = Path(CFG["staging_dir"]) / "images"
staging.mkdir(parents=True, exist_ok=True)
tasks   = []

for rec in to_annotate:
    src  = Path(rec["filepath"])
    name = f"{rec['id']}__{rec['gt_label']}__{src.name}"
    shutil.copy2(src, staging / name)
    tasks.append({
        "data": {"image": f"/data/local-files/?d=hard_examples/images/{name}",
                 "milvus_id": rec["id"]},
        "predictions": [{"result": [{"type": "choices",
                          "value": {"choices": [rec["gt_label"]]}}]}]
    })

with open(Path(CFG["staging_dir"]) / "label_studio_tasks.json", "w") as f:
    json.dump(tasks, f, indent=2)

print(f"Exported {len(tasks)} tasks -> {CFG['staging_dir']}/label_studio_tasks.json")
print("Annotate in Label Studio, then export JSON-MIN to:", CFG["ls_export"])

# ----------------------------------------------------------------
# 6. PARSE ANNOTATIONS + WRITE BACK TO MILVUS
# ----------------------------------------------------------------
with open(CFG["ls_export"]) as f:
    ls_data = json.load(f)

id_map = {
    int(t["data"]["milvus_id"]):
    t["annotations"][0]["result"][0]["value"]["choices"][0]
    for t in ls_data
}

old = col.query(
    expr=f"id in [{','.join(str(i) for i in id_map)}]",
    output_fields=["id", "filepath", "embedding", "gt_label", "pred_label",
                   "confidence", "correct", "split", "loop_iter"],
)
col.delete(f"id in [{','.join(str(i) for i in id_map)}]")

new_rows = {k: [] for k in ["filepath", "embedding", "gt_label", "pred_label",
                              "confidence", "correct", "split", "annotated",
                              "corrected_label", "loop_iter"]}
for r in old:
    new_rows["filepath"].append(r["filepath"])
    new_rows["embedding"].append(r["embedding"])
    new_rows["gt_label"].append(r["gt_label"])
    new_rows["pred_label"].append(r["pred_label"])
    new_rows["confidence"].append(r["confidence"])
    new_rows["correct"].append(r["correct"])
    new_rows["split"].append(r["split"])
    new_rows["annotated"].append(True)
    new_rows["corrected_label"].append(id_map[r["id"]])
    new_rows["loop_iter"].append(r["loop_iter"])

col.insert(list(new_rows.values()))
col.flush()
print(f"Updated {len(old)} records with corrected labels")

# ----------------------------------------------------------------
# 7. BUILD DATASET FROM MILVUS
# ----------------------------------------------------------------
train_recs = col.query(
    expr='split == "train"',
    output_fields=["filepath", "gt_label"],
    limit=200000,
)
hard_recs = col.query(
    expr='annotated == True and split == "eval"',
    output_fields=["filepath", "corrected_label"],
    limit=16384,
)

rows = [{"filepath": r["filepath"], "label": r["gt_label"]} for r in train_recs]
rows += [{"filepath": r["filepath"], "label": r["corrected_label"]} for r in hard_recs]
df   = pd.DataFrame(rows).drop_duplicates("filepath")
df.to_csv("data/train_manifest_v2.csv", index=False)
print(f"Dataset: {len(df)} samples  Classes: {df['label'].nunique()}")

# ----------------------------------------------------------------
# 8. FINE-TUNE
# ----------------------------------------------------------------
class ManifestDataset(Dataset):
    def __init__(self, path, transform):
        self.df      = pd.read_csv(path)
        self.classes = sorted(self.df["label"].unique())
        self.c2i     = {c: i for i, c in enumerate(self.classes)}
        self.tf      = transform

    def __len__(self):
        return len(self.df)

    def __getitem__(self, i):
        r = self.df.iloc[i]
        x = Image.open(r["filepath"]).convert("RGB")
        return self.tf(x), self.c2i[r["label"]]

train_tf = transforms.Compose([
    transforms.RandomResizedCrop(224, scale=(0.7, 1.0)),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(0.3, 0.3, 0.2),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
val_tf = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])

train_ds  = ManifestDataset("data/train_manifest_v2.csv", train_tf)
val_ds    = datasets.ImageFolder(CFG["val_dir"], transform=val_tf)
cc        = Counter(train_ds.df["label"])
sw        = [1.0 / cc[train_ds.df.iloc[i]["label"]] for i in range(len(train_ds))]
sampler   = WeightedRandomSampler(sw, len(sw), replacement=True)
train_ldr = DataLoader(train_ds, batch_size=CFG["batch_size"], sampler=sampler,
                        num_workers=4, pin_memory=True)
val_ldr   = DataLoader(val_ds, batch_size=CFG["batch_size"], shuffle=False, num_workers=4)

model_ft  = torch.load(CFG["checkpoint_in"], map_location=CFG["device"]).to(CFG["device"])
crit      = nn.CrossEntropyLoss()
opt       = optim.AdamW(model_ft.parameters(), lr=CFG["lr"], weight_decay=1e-4)
sched     = optim.lr_scheduler.CosineAnnealingLR(opt, T_max=CFG["epochs"])
best      = 0.0

for ep in range(1, CFG["epochs"] + 1):
    model_ft.train()
    for x, y in tqdm(train_ldr, desc=f"E{ep:02d}"):
        x, y = x.to(CFG["device"]), y.to(CFG["device"])
        opt.zero_grad()
        crit(model_ft(x), y).backward()
        opt.step()
    sched.step()

    model_ft.eval()
    c = t = 0
    with torch.no_grad():
        for x, y in val_ldr:
            p  = model_ft(x.to(CFG["device"])).argmax(1).cpu()
            c += (p == y).sum().item()
            t += y.size(0)
    acc = c / t
    print(f"E{ep:02d} val_acc={acc:.4f}")
    if acc > best:
        best = acc
        torch.save(model_ft, CFG["checkpoint_out"])

print(f"Done. Best val acc: {best:.4f} -> {CFG['checkpoint_out']}")

# ----------------------------------------------------------------
# 9. RE-INDEX WITH NEW MODEL
# Re-run section 2 with:
#   checkpoint_in = CFG["checkpoint_out"]
#   loop_iter     = CFG["loop_iter"] + 1
# Then compare error counts between loop_iter 0 and 1.
# ----------------------------------------------------------------

Summary

The Milvus-backed pipeline replaces flat CSV files with a queryable, searchable vector store that scales to millions of images. Inference inserts embeddings and prediction metadata into Milvus in batches. Scalar filters find wrong predictions and high-confidence errors with simple expressions. ANN search from the error-cluster centroid combined with greedy diversity selection surfaces the most informative and non-redundant hard examples. Delete plus re-insert writes human-corrected labels back to Milvus, keeping annotation state consistent across iterations. Dataset assembly pulls training records and annotated hard examples from Milvus directly into a manifest CSV which feeds a WeightedRandomSampler-backed DataLoader. Loop iteration tags let you compare model performance before and after fine-tuning directly in Milvus without any external tracking tool.

Overview and Architecture

Why Milvus for This Pipeline?

Full Data Flow

Loop Iteration Flow

Prerequisites and Setup

Required Packages

Start Milvus with Docker Compose

Verify the Connection

Step 1 — Milvus Schema Design

Field Definitions

Schema Lifecycle

Create the Collection

Step 2 — Run Inference and Index Embeddings

Hook-Based Embedding Extraction

Step 3 — Query Mispredicted Samples

Query Decision Tree

All Wrong Predictions

High-Confidence Wrong Predictions

Per-Class Confusion Breakdown

Targeted Query for a Specific Error Pair

Step 4 — Rank and Diversify via ANN Search

Diversity Selection Strategy

Compute the Error-Cluster Centroid

ANN Search from the Centroid

Greedy Furthest-Point Sampling

Per-Class Quota with Hybrid Search

Step 5 — Export and Prepare for Annotation

Annotation Workflow

Fetch Full Records and Copy Images

Generate Label Studio Import Manifest

Launch Label Studio

Step 6 — Re-annotate and Write Back to Milvus

Milvus Update Pattern

Parse Label Studio Export

Delete and Re-Insert with Corrected Labels

Step 7 — Build the Fine-Tuning Dataset from Milvus

Dataset Assembly Flow

Query Original and Annotated Records

Merge into a Unified Manifest

PyTorch Dataset Backed by the Manifest

Class-Balanced Sampler from Milvus Counts

Step 8 — Fine-Tune the Model

Training Phase Diagram

Fine-Tuning Script

Focal Loss for Hard Examples

Step 9 — Evaluate and Update Milvus with New Predictions

Cross-Iteration Comparison

Error Count Comparison Across Iterations

Per-Class Delta

Verify Previously Hard Examples Were Fixed

Step 10 — Closing the Active Learning Loop

Loop Controller

Advanced Techniques

Milvus Partitions Per Loop Iteration

Pseudo-Labelling High-Confidence Unlabelled Images

Semantic Near-Duplicate Detection via ANN

Tracking Label Noise Across Iterations

Paginating Large Result Sets

Common Pitfalls

Schema Immutability

Collection Must Be Loaded Before Querying

Always Flush After Insert

Avoid Catastrophic Forgetting

End-to-End Script

Summary

Related posts