DETR Set Loss

Problem

DETR-style set loss. Find the assignment between queries and targets that minimizes the total matching cost, then average cost_bbox * l1 + cost_class * ce over the matched pairs. Assume num_queries >= num_targets and search over all row/column permutations with itertools.permutations.

The signature is detr_set_loss(cost_matrix, l1_matrix, ce_matrix, cost_bbox, cost_class). Return [assignment, loss], where assignment is a list of [row, col] pairs.

Hint

DETR matches predictions to targets with a Hungarian-style assignment before computing the loss.

Test cases

case 1

in: [[[1,2],[3,0.5]],[[0.5,1],[2,0.1]],[[0.2,0.8],[0.9,0.1]],5,1]

expected: [[[0,0],[1,1]],1.65]

case 2

in: [[[0.1,0.2,0.3],[0.4,0.5,0.6]],[[1,1,1],[2,2,2]],[[0.1,0.2,0.3],[0.3,0.2,0.1]],2,1]

expected: [[[0,0],[1,1]],3.15]

case 3

in: [[[1,4],[2,1],[3,2]],[[0.5,0.2],[0.1,0.7],[0.3,0.4]],[[0.1,0.9],[0.8,0.2],[0.5,0.5]],1,2]

expected: [[[0,0],[1,1]],0.9]

case 4

in: [[[1,1],[1,1]],[[2,2],[2,2]],[[0,0],[0,0]],1,1]

expected: [[[0,0],[1,1]],2]

Starter codePython

import itertools
def detr_set_loss(cost_matrix, l1_matrix, ce_matrix, cost_bbox, cost_class):
    # Your code here
    pass

4 test cases · run your code against them in the editor