OSNet Person ReID in Python: Choosing a Variant, Exporting to ONNX, and Measuring CPU Latency

In Real Time Human Detection and ReID with Yolo11 + OSNet we ran person re-identification on an Android phone with Kotlin, C++ and TensorFlow Lite. Since then, most of the searches that reach that post ask a simpler question: which OSNet model should I use, and how do I run it? This post is the desktop companion. Everything below runs in Python on a laptop CPU, uses only public code and public weights, and every number in it comes from running the scripts on this page.

By the end you will have:

• a 60-line helper that loads any OSNet checkpoint from the torchreid model zoo and turns person crops into 512-d embeddings,
• YOLO11 person crops matched by cosine similarity,
• an ONNX export that needs neither PyTorch nor torchreid at inference time,
• CPU latency for four OSNet variants, PyTorch vs ONNX Runtime, plus one result that surprised us about batching.

What OSNet gives you

OSNet (Omni-Scale Network, Zhou et al., ICCV 2019) is a small CNN built for person re-identification. Each block mixes several receptive-field sizes, so one feature captures both local details (a logo, a bag strap) and whole-body cues (overall colour, silhouette). In torchreid it takes a 256 × 128 RGB crop and, in eval mode, returns a 512-d feature vector. Two crops of the same person should produce vectors that point in nearly the same direction.

The torchreid model zoo publishes several widths. Numbers below are copied from the model zoo page, shown as Rank-1 (mAP) in %:

Two things stand out. First, the same-domain column is the one most READMEs quote, but it is the least useful for you: it measures a model tested on the same cameras it was trained on. Your cameras are not in Market-1501. The cross-domain column (train on MSMT17, test on Market-1501 without fine-tuning) is a better guess at what happens on your footage, and there the gap between widths is larger. Second, osnet_ain_x1_0 (OSNet with instance normalisation, from the follow-up TPAMI paper) is the best cross-domain model in the zoo at the same size as osnet_x1_0.

We use the MSMT17 checkpoints trained with combineall=True from that table: they were trained on the largest of the three datasets and are the ones most trackers ship by default.

Setup

python -m venv venv && source venv/bin/activate
pip install torch torchvision onnx onnxruntime ultralytics gdown tensorboard
# torchreid: either the official repo (github.com/KaiyangZhou/deep-person-reid, install from source)
# or the community wheel on PyPI, which is what we used here:
pip install torchreid==0.2.5

The PyPI wheel is a repackaging of KaiyangZhou's deep-person-reid by a third party, and it puts the models under torchreid.reid.models instead of torchreid.models. The helper below imports from either. It also imports torch.utils.tensorboard on load, which is why tensorboard is in the list.

Download the weights (Google Drive IDs from the model zoo page):

mkdir -p weights && cd weights
gdown 1Kkx2zW89jq_NETu4u42CFZTMVD5Hwm6e -O osnet_x0_25_msmt17.pth
gdown 1DHgmb6XV4fwG3n-CnCM0zdL9nMsZ9_RF -O osnet_x0_5_msmt17.pth
gdown 1IosIFlLiulGIjwW3H8uMRmx3MzPwf86x -O osnet_x1_0_msmt17.pth
gdown 1SigwBE6mPdqiJMqhuIY4aqC7--5CsMal -O osnet_ain_x1_0_msmt17.pth

The osnet_x0_25 file is 9.3 MB even though the network has 0.2 M parameters. The rest is the training classifier: one row per MSMT17 identity (4,101 × 512 floats). You don't need it at inference.

Load a checkpoint and embed crops

osnet_reid.py:

import numpy as np
import torch
import torch.nn.functional as F

try:  # official repo
    from torchreid.models import build_model
except ImportError:  # PyPI wheel torchreid==0.2.5
    from torchreid.reid.models import build_model

MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
H, W = 256, 128  # what every model in the torchreid zoo was trained on


def load_osnet(name: str, weights: str) -> torch.nn.Module:
    """Build `name` (e.g. 'osnet_x0_25') and load a model-zoo checkpoint.
    The checkpoint's classifier (one row per training identity) is dropped:
    in eval mode OSNet returns the 512-d feature that sits before it."""
    model = build_model(name, num_classes=1, pretrained=False)
    ckpt = torch.load(weights, map_location="cpu", weights_only=True)
    state = ckpt.get("state_dict", ckpt)
    state = {k.removeprefix("module."): v for k, v in state.items()}
    state = {k: v for k, v in state.items() if not k.startswith("classifier.")}
    missing, unexpected = model.load_state_dict(state, strict=False)
    assert not unexpected and all(k.startswith("classifier.") for k in missing), (missing, unexpected)
    return model.eval()


def preprocess(crops_rgb: list[np.ndarray]) -> torch.Tensor:
    """uint8 HxWx3 RGB crops -> float tensor [N, 3, 256, 128], ImageNet-normalised."""
    batch = []
    for crop in crops_rgb:
        t = torch.tensor(crop).permute(2, 0, 1).float().unsqueeze(0) / 255.0
        t = F.interpolate(t, size=(H, W), mode="bilinear", align_corners=False)
        batch.append(t)
    x = torch.cat(batch)
    return (x - torch.from_numpy(MEAN)[:, None, None]) / torch.from_numpy(STD)[:, None, None]


@torch.inference_mode()
def embed(model: torch.nn.Module, crops_rgb: list[np.ndarray]) -> np.ndarray:
    """L2-normalised [N, 512] embeddings: a dot product is the cosine similarity."""
    return F.normalize(model(preprocess(crops_rgb)), dim=1).numpy()


def export_onnx(model: torch.nn.Module, path: str) -> None:
    """Dynamic batch axis, so one session handles 1..N crops per frame."""
    torch.onnx.export(
        model, torch.randn(1, 3, H, W), path,
        input_names=["images"], output_names=["features"],
        dynamic_axes={"images": {0: "n"}, "features": {0: "n"}},
        opset_version=17, dynamo=False,
    )

Three details in there are where most "ReID doesn't work" bug reports come from:

• Strip module. and drop classifier.*. Of the four checkpoints we used, the AIN one was saved from DataParallel (every key starts with module.), and all four carry the training head. The assert makes sure everything else loaded; a silent partial load gives you a randomly initialised network that still returns 512 numbers.
• model.eval(). In train mode OSNet returns classifier logits, not features (with num_classes=1 you get one number per crop), and batch norm uses batch statistics.
• RGB, ImageNet mean/std, 256 × 128. The same preprocessing the zoo models were trained with.

weights_only=True keeps torch.load from unpickling arbitrary objects. These checkpoints are plain tensors, so it costs nothing.

Detect, crop, match

For detection we use the public YOLO11n weights from Ultralytics and keep class 0 (person). One thing to watch: Ultralytics gives you the original image in OpenCV's BGR order, and OSNet expects RGB.

import numpy as np
from ultralytics import YOLO
from osnet_reid import embed, load_osnet


def person_crops(detector, sources, min_conf=0.5):
    crops = []
    for res in detector(sources, classes=[0], conf=min_conf, verbose=False):
        rgb = res.orig_img[:, :, ::-1]  # BGR -> RGB
        for x1, y1, x2, y2 in res.boxes.xyxy.round().int().tolist():
            crops.append(np.ascontiguousarray(rgb[y1:y2, x1:x2]))
    return crops


detector = YOLO("yolo11n.pt")
crops = person_crops(detector, ["https://ultralytics.com/images/bus.jpg",
                                "https://ultralytics.com/images/zidane.jpg"])
model = load_osnet("osnet_x0_25", "weights/osnet_x0_25_msmt17.pth")
gallery = embed(model, crops)          # [N, 512]
sim = gallery @ gallery.T              # cosine similarity matrix

Matching is a nearest-neighbour search: compare a new crop's embedding with the embeddings you have already seen and take the best one, if it is similar enough. Because the vectors are L2-normalised, cosine similarity is a dot product, and a whole gallery is one matrix multiply.

A quick sanity check (not a benchmark)

Two sample images do not contain anyone twice, so to see the separation we made a "second sighting" of each person: the crop is mirrored, darkened by 30%, downscaled 3× and back, and its box shifted by 8% horizontally and 5% vertically. YOLO11n found five person boxes at confidence ≥ 0.5. For each variant we compared each altered crop with the five originals:

Five people is far too few to rank models; the zoo table above is the real comparison. What this does show is how a threshold behaves. Every model put the right person first, but the gap between "least similar same person" and "most similar different person" is only 0.18 for osnet_x0_25 and 0.29 for osnet_x1_0. Pick your acceptance threshold from pairs on your own cameras, lighting and crop sizes. A number copied from a blog post or a GitHub issue was tuned on someone else's footage.

On real footage: one person, four angles

Altered copies of the same crop are the easy case. The hard case is the one ReID exists for: the same person seen from another side, in another pose, in another clip. So we took three short phone clips recorded in our office (walking away from and back toward the camera, standing still, sitting on a chair), each with a colleague standing in the background. YOLO11n found exactly two people in every frame. We used the author walking away from the camera as the query and ranked seven crops from the three clips with osnet_x1_0:

All four sightings of the author rank above all three of the colleague, so nearest-neighbour matching gets this right. The numbers are the more useful part:

• The same view stays close: 0.90 for the author from behind again, even though that crop is only 59×177 pixels.
• A change of view roughly halves the score: 0.48–0.58 from the front, standing or sitting. That is below the 0.585 that the most similar pair of different people reached in the synthetic check above. A threshold tuned on altered copies would have rejected every one of these real matches.
• Different people stay low here (0.17–0.19), but that is partly luck: dark jacket against grey jumper, and the colleague's crops are only about 56×129 pixels. Two people in similar clothes would sit much closer.

Two people are not a benchmark either. What this shows is why the threshold has to come from your own footage: ranking held up, but a fixed cut-off taken from easy pairs would not.

Export to ONNX and drop PyTorch

from osnet_reid import export_onnx, load_osnet

model = load_osnet("osnet_x1_0", "weights/osnet_x1_0_msmt17.pth")
export_onnx(model, "osnet_x1_0.onnx")

We used the TorchScript-based exporter (dynamo=False), which takes dynamic_axes directly; since PyTorch 2.9 the default is the newer torch.export-based one, so expect a deprecation warning. Exporting the AIN variant also prints a warning that instance norm is exported in training mode. That is expected: instance norm always normalises with the statistics of the input itself, in training and in eval. The parity check below shows the exported graph gives the same output.

On the deployment side you only need ONNX Runtime and OpenCV:

import cv2
import numpy as np
import onnxruntime as ort

MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)


class OSNetONNX:
    def __init__(self, path, threads=4):
        opts = ort.SessionOptions()
        opts.intra_op_num_threads = threads
        self.sess = ort.InferenceSession(path, opts, providers=["CPUExecutionProvider"])

    def __call__(self, crops_rgb):
        x = np.stack([cv2.resize(c, (128, 256), interpolation=cv2.INTER_LINEAR) for c in crops_rgb])
        x = ((x.astype(np.float32) / 255.0 - MEAN) / STD).transpose(0, 3, 1, 2)  # NHWC -> NCHW
        f = self.sess.run(None, {"images": np.ascontiguousarray(x)})[0]
        return f / np.linalg.norm(f, axis=1, keepdims=True)

Note that cv2.resize takes (width, height), so it is (128, 256) here. We checked both paths:

• Same input tensor, PyTorch vs ONNX Runtime: largest absolute difference across the four models was 3.5 × 10^-4, and the lowest cosine similarity between the two outputs was 0.9999998.
• Full deployment path (OpenCV resize + ONNX Runtime) vs PyTorch path (F.interpolate) on the five real crops: lowest cosine 0.99994 for osnet_x1_0 and 0.99992 for osnet_x0_25. The resize implementations differ slightly; the embeddings don't move in any way that matters.

If you resize with Pillow instead, be aware that Pillow's resize antialiases when downscaling while F.interpolate and cv2.INTER_LINEAR do not. Use the same resize function at gallery-building time and at query time.

How fast is it on a CPU?

bench.py exports each variant, feeds random input of shape [1, 3, 256, 128] and [8, 3, 256, 128], and takes the median of 50 timed runs after 10 warm-up runs, with 4 threads for both PyTorch and ONNX Runtime. Machine: Apple M4 (10-core) laptop, Python 3.12, PyTorch 2.15.0 eager, ONNX Runtime 1.31.1 CPU provider. The laptop was running other jobs during some runs, so we ran the benchmark eight times (five times for the last column) and report the lowest median for each cell. The full script is at the end of the post.

What we took from it:

• ONNX Runtime was 2.2 to 3.3 times faster than eager PyTorch for a single crop. The export is close to free, so do it before you reach for a GPU.
• Batching did not pay off on this CPU for the wider models. For osnet_x1_0, one batch of 8 took 50 ms while 8 separate single-crop calls took 34 ms. Only osnet_x0_25 got faster with a batch. We have not profiled why (cache size is a plausible guess), so don't assume this holds on your hardware. Time both on your own machine before you build a batching queue.
• Even the largest model is cheap for a handful of people. Eight crops with osnet_x1_0 took about 34 ms on 4 CPU threads, before any quantisation or GPU.

Which variant should you pick?

• New cameras, no fine-tuning: osnet_ain_x1_0. It has the best cross-domain numbers in the zoo, and on CPU it costs about the same as osnet_x1_0.
• Tight CPU or mobile budget: osnet_x0_25 is 0.9 MB and 1.3 ms per crop here, but it scores about 10 Rank-1 points lower cross-domain than osnet_ain_x1_0 and had the narrowest same/different margin in our sanity check. Plan to use it together with motion cues (box overlap, a Kalman filter) rather than on its own.
• In between: osnet_x0_5, which we also used on Android in the earlier post.
• If you can label even a few hundred crops from your own cameras, fine-tuning any of these with torchreid will usually do more than switching variants.

Checklist

• Check that the checkpoint really loaded: compare the expected keys, not just "no error".
• Call model.eval(); in train mode you get logits.
• Feed RGB, 256 × 128, ImageNet mean/std, and use the same resize everywhere.
• L2-normalise, then compare with a dot product.
• Export with a dynamic batch axis, then check PyTorch vs ONNX output on real crops.
• Calibrate the match threshold on your own footage.
• Time batch-of-N vs N single calls on the hardware you will deploy to.

Benchmark script

"""usage: python bench.py WEIGHTS_DIR OUT_DIR"""
import json, os, platform, statistics, sys, time

import numpy as np
import onnxruntime as ort
import torch

from osnet_reid import H, W, export_onnx, load_osnet

VARIANTS = ["osnet_x0_25", "osnet_x0_5", "osnet_x1_0", "osnet_ain_x1_0"]
THREADS, RUNS, WARMUP = 4, 50, 10


def median_ms(fn):
    for _ in range(WARMUP):
        fn()
    times = []
    for _ in range(RUNS):
        t0 = time.perf_counter()
        fn()
        times.append((time.perf_counter() - t0) * 1000)
    return round(statistics.median(times), 2)


def main(weights_dir, out_dir):
    os.makedirs(out_dir, exist_ok=True)
    torch.set_num_threads(THREADS)
    opts = ort.SessionOptions()
    opts.intra_op_num_threads = THREADS
    rng = np.random.default_rng(0)
    for name in VARIANTS:
        model = load_osnet(name, os.path.join(weights_dir, f"{name}_msmt17.pth"))
        onnx_path = os.path.join(out_dir, f"{name}.onnx")
        export_onnx(model, onnx_path)
        sess = ort.InferenceSession(onnx_path, opts, providers=["CPUExecutionProvider"])
        row = {"model": name, "onnx_MB": round(os.path.getsize(onnx_path) / 1e6, 1)}
        for n in (1, 8):
            x = rng.standard_normal((n, 3, H, W), dtype=np.float32)
            xt = torch.from_numpy(x)
            with torch.inference_mode():
                ref = model(xt).numpy()
                row[f"torch_b{n}_ms"] = median_ms(lambda: model(xt))
            out = sess.run(None, {"images": x})[0]
            row[f"ort_b{n}_ms"] = median_ms(lambda: sess.run(None, {"images": x}))
            row[f"b{n}_max_abs_diff"] = float(np.abs(ref - out).max())
            if n == 8:
                row["ort_8x1_ms"] = median_ms(lambda: [sess.run(None, {"images": x[i:i + 1]}) for i in range(n)])
        print(json.dumps(row))
    print(platform.machine(), platform.python_version(), torch.__version__, ort.__version__)


if __name__ == "__main__":
    main(*sys.argv[1:3])

References

1. K. Zhou, Y. Yang, A. Cavallaro, T. Xiang. Omni-Scale Feature Learning for Person Re-Identification. ICCV 2019.
2. K. Zhou, Y. Yang, A. Cavallaro, T. Xiang. Learning Generalisable Omni-Scale Representations for Person Re-Identification. TPAMI 2021 (OSNet-AIN).
3. torchreid model zoo (accuracy, parameters, FLOPs and weights used here) and source.
4. Ultralytics YOLO11 documentation.
5. ONNX Runtime documentation.