matching

Set-prediction matching between predictions and ground-truth objects.

Functions:

hungarian_match –

One-to-one assignment of predictions to targets minimizing the total cost (Hungarian algorithm).
hungarian_match_batched –

One-to-one assignment of every scene's predictions to its targets, with a single device transfer.

hungarian_match [source] ¶

hungarian_match(cost: Tensor) -> Tuple[Tensor, Tensor]

One-to-one assignment of predictions to targets minimizing the total cost (Hungarian algorithm).

The matcher of the set-prediction detectors (3DETR, TransFusion): every target is matched to exactly one prediction and the sum of the matched costs is minimal. The cost matrix is built by the caller (e.g. a weighted sum of a classification, a center-distance and an IoU term) and must be finite.

Parameters:

cost (Tensor) –

Pairwise cost of every prediction / target pair, shape \((Q, M)\).

Returns:

Tensor –

A tuple (pred_indices, target_indices) of long tensors on cost.device, each of length
Tensor –

\(\min(Q, M)\), such that pred_indices[i] is matched to target_indices[i]; both are empty when
Tuple[Tensor, Tensor] –

either side is.

Shape

cost: \((Q, M)\)
pred_indices: \((\min(Q, M),)\)
target_indices: \((\min(Q, M),)\)

Example

>>> cost = torch.tensor([[0.9, 0.1], [0.2, 0.8], [0.5, 0.5]])
>>> pred, target = hungarian_match(cost)
>>> pred.tolist(), target.tolist()
([0, 1], [1, 0])

hungarian_match_batched [source] ¶

hungarian_match_batched(
    cost: Tensor, num_targets: Sequence[int]
) -> Tensor

One-to-one assignment of every scene's predictions to its targets, with a single device transfer.

The batched form of hungarian_match for a padded cost tensor: scene \(i\) uses the first num_targets[i] target columns of cost[i]. The whole tensor is moved to the CPU once and every scene is solved there, so a loop over scenes (and decoder layers, when they are stacked along the batch dimension) costs one synchronization instead of one per scene.

Parameters:

cost (Tensor) –

Pairwise costs, shape \((N, Q, M_\max)\); the columns beyond a scene's target count are ignored and may hold anything finite or not.
num_targets (Sequence[int]) –

Number of valid target columns of each scene, length \(N\).

Returns:

Tensor –

The matched target index of every prediction, shape \((N, Q)\) long on cost.device, \(-1\) where the
Tensor –

prediction is unmatched.

Shape

cost: \((N, Q, M_\max)\)
output: \((N, Q)\)

Example

>>> cost = torch.tensor([[[0.9, 0.1], [0.2, 0.8], [0.5, 0.5]]])
>>> hungarian_match_batched(cost, [2]).tolist()
[[1, 0, -1]]