Batched tracking workflow
Detection-to-track association on a batch of frames using IoU costs, with CPU and CUDA backends and one-pass unpacking.
A multi-object-tracking scenario: for each frame in a batch, match the current set of detected objects (detections, e.g., bounding boxes from a detector) to the objects the tracker is already following (tracks). The goal is to decide which detection corresponds to which track — assigning each detection to at most one track at minimum cost.
This tutorial covers three steps:
- Build batched IoU costs on the chosen device.
- Pick the right batched op for the data shape.
- Recover
(matches, unmatched_rows, unmatched_cols)in one pass.
Synthesizing realistic costs
import torch
import torchmatch # extensions load eagerly at import
def boxes(n, device, generator):
"""Random axis-aligned boxes in [0, 1]² with side length ~ U(0.1, 0.4)."""
xy = torch.rand(n, 2, generator=generator, device=device)
wh = torch.rand(n, 2, generator=generator, device=device) * 0.3 + 0.1
return torch.cat([xy, xy + wh], dim=1) # (n, 4) = x1, y1, x2, y2
def iou_cost(boxes_a, boxes_b):
"""Pairwise (1 - IoU) cost matrix."""
lo = torch.maximum(boxes_a[:, None, :2], boxes_b[None, :, :2])
hi = torch.minimum(boxes_a[:, None, 2:], boxes_b[None, :, 2:])
inter = (hi - lo).clamp_min(0).prod(-1)
area_a = (boxes_a[:, 2:] - boxes_a[:, :2]).prod(-1)[:, None]
area_b = (boxes_b[:, 2:] - boxes_b[:, :2]).prod(-1)[None, :]
iou = inter / (area_a + area_b - inter + 1e-9)
return 1.0 - iou
def build_batch(B, N, device):
g = torch.Generator(device=device).manual_seed(0)
return torch.stack([
iou_cost(boxes(N, device, g), boxes(N, device, g))
for _ in range(B)
]).contiguous()
Picking the right batched op
There are only a few batched solver functions to choose from. For a batch
of B problems of size N×N:
def solve_batch(costs: torch.Tensor) -> torch.Tensor:
"""Dispatch to the fastest backend for this batch shape."""
B, N, _ = costs.shape
if costs.is_cuda and N <= 64:
# CUDA tiled backend is CUDA-graph-safe; pick it when it applies.
return torchmatch.assignment.ops.jonker_dense_batch(costs)
if costs.is_cuda:
# Tiled CUDA kernel rejects K > 64, so route oversized problems via CPU.
out_cpu = torchmatch.assignment.ops.jonker_dense_batch(costs.cpu())
return out_cpu.to(costs.device)
return torchmatch.assignment.ops.jonker_dense_batch(costs)
The CUDA tiled kernel is fastest for typical tracking costs when N ≤ 64.
For larger N, routing through CPU is faster because the CPU op distributes
each problem to a separate thread in parallel — CUDA's advantage diminishes
when each problem is too large to fit in on-chip shared memory.
Running it
device = "cuda" if torch.cuda.is_available() else "cpu"
# B=16 frames, N=32 boxes per frame
costs = build_batch(B=16, N=32, device=device)
row_to_col = solve_batch(costs)
print(row_to_col.shape)
# (16, 32). For each frame, the row→col mapping.
# Total cost across the batch
batch_idx = torch.arange(16, device=device)[:, None]
row_idx = torch.arange(32, device=device)[None, :]
matched = row_to_col >= 0
totals = (costs[batch_idx, row_idx, row_to_col.clamp_min(0)] * matched).sum(-1)
print(totals)
Getting the unmatched sets in one pass
When you need to separately access the matched pairs and the unmatched
detections or tracks, the _unpacked variants return those as three tensors
directly. Without them, you would have to loop over the row-to-column index
tensor in Python and compute the unmatched sets yourself.
The _unpacked variants cost about the same as the packed ones (within
about 5 %). When the unpacked output is what you need, use them instead
of a post-hoc Python loop.
Adding feasibility gating (+inf edges)
Real trackers exclude implausible pairs before solving: any detection–track
pair whose center-point distance exceeds a threshold is forbidden by setting
its cost to +inf. This prevents the solver from ever matching a detection
to a track that is too far away.
def gated_iou_cost(boxes_a, boxes_b, gate=0.3):
"""IoU cost with centroid-distance gating."""
cost = iou_cost(boxes_a, boxes_b)
centers_a = (boxes_a[:, None, :2] + boxes_a[:, None, 2:]) / 2
centers_b = (boxes_b[None, :, :2] + boxes_b[None, :, 2:]) / 2
dist = (centers_a - centers_b).norm(dim=-1)
return cost.masked_fill(dist > gate, float("inf"))
Use it the same way; solve_batch(costs) handles the +inf entries
internally. The Jonker-Volgenant (JV) solvers handle this automatically:
they replace +inf internally with a large finite value (a sentinel) so
the underlying algorithm never sees infinity, but the forbidden-pair
constraint is still respected. Leave at least one feasible matching per
problem after gating; the JV ops mark unmatchable rows with -1 rather
than failing.
When to use a CUDA graph
The CUDA backend of jonker_dense_batch is the only assignment solver that
works inside a CUDA graph. A CUDA graph records a sequence of GPU operations
once and can replay them repeatedly with minimal CPU overhead — useful when
the same solver call runs on every frame of an inference loop:
g = torch.cuda.CUDAGraph()
costs_in = torch.empty(16, 32, 32, device="cuda")
with torch.cuda.graph(g):
out = torchmatch.assignment.ops.jonker_dense_batch(costs_in)
# Per-frame:
costs_in.copy_(new_costs)
g.replay()
torch.cuda.synchronize()
# out is now populated
The other CUDA solvers (munkres, hybrid, lawler) cannot be recorded
into a CUDA graph — they synchronize with the CPU mid-execution, which
breaks graph capture. Do not use them inside a torch.cuda.graph block.
Unlike the other CUDA solvers, jonker_dense_batch never pauses to
communicate with the CPU during execution, so it can be captured into a
CUDA graph without restriction.
What's next
- Choosing the right op: when to pick which variant.
- Operations reference: exact constraints and output shapes.