Detr: Possible bug in MHAttention for mask prediction

[Ali2500]

Regarding this code snippet in the MHAttention class in segmentation.py:

weights = torch.einsum("bqnc,bnchw->bqnhw", qh * self.normalize_fact, kh)

if mask is not None:
     weights.masked_fill_(mask.unsqueeze(1).unsqueeze(1), float("-inf"))
weights = F.softmax(weights.flatten(2), dim=-1).view(weights.size())

In the last line, I'm assuming weights initially has shape [b, q, n, h, w]. It is then flattened to [b, q, n * h * w] before applying a softmax over the last dimension. Correct me if I'm wrong, but since this is multi-head attention, isn't the softmax supposed to be applied over each head separately? If so then then it should be weights.flatten(3) instead of weights.flatten(2).

[fmassa]

Hi @Ali2500

Thanks a lot for opening this issue.

You are right -- there is a problem with the softmax computation in the MHA for panoptic segmentation.

I tried running a training to see if this affects results, and it seems that with the proposed fix we get better panoptic results.

Here are the results I got after the standard 25 epochs:

Control run (current master)

          |    PQ     SQ     RQ     N
--------------------------------------
All       |  43.9   79.6   54.3   133
Things    |  48.8   80.1   60.1    80
Stuff     |  36.7   78.8   45.7    53

segm AP: 31.8

With softmax fix

          |    PQ     SQ     RQ     N
--------------------------------------
All       |  44.3   80.0   54.5   133
Things    |  49.2   80.6   60.3    80
Stuff     |  37.0   79.1   45.9    53

segm AP: 32.9

---

For some reason the results on master are also better than the reported results we have. @szagoruyko run the original experiments on panoptic and his environment might have been slightly different than mine.

One thing to note is that with the softmax fix training of the panoptic head started with a much higher PQ (38.0 vs 3.8), although over the course of the training they reached results on a similar ballpark.