Mask2Former - ValueError: cost matrix is infeasible
Author: asgerius
Pull request: no

### System Info

```
- `transformers` version: 4.26.0
- Platform: Linux-5.4.0-137-generic-x86_64-with-glibc2.27
- Python version: 3.10.9
- Huggingface_hub version: 0.12.0
- PyTorch version (GPU?): 1.13.1+cu117 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: 4 x RTX 2080 Ti
- Using distributed or parallel set-up in script?: Single node DistributedDataParallel setup
```

### Who can help?

@amyeroberts 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

I am fine-tuning Mask2Former for a semantic segmentation task. I sometimes get the error `ValueError: cost matrix is infeasible`. With the learning rate of `1e-4` that I use, it usually takes many thousands of batches with the loss happily dropping to get this error. From my experience, the higher the learning rate, the more often it will happen. The error happens during the forward pass, which roughly looks like this:
```py
image_processor = AutoImageProcessor.from_pretrained("facebook/mask2former-swin-small-ade-semantic")
inputs = image_processor.preprocess(batch, mask_labels, return_tensors="pt")

batch, mask, class_labels = inputs["pixel_values"], inputs["mask_labels"], inputs["class_labels"]
batch = batch.to(device)
mask = [x.to(device) for x in mask]
class_labels = [x.to(device) for x in class_labels]

with torch.cuda.amp.autocast():
    out = model(
        pixel_values = batch,
        mask_labels  = mask,
        class_labels = class_labels,
    )
```

Unfortunately, I don't have the time to set up a full minimally reproducible example, but I have tracked the error down to `cost_matrix` containing `torch.inf` [here](https://github.com/huggingface/transformers/blob/48327c57182fdade7f7797d1eaad2d166de5c55b/src/transformers/models/mask2former/modeling_mask2former.py#L491) (or see stack trace below).

Full stack trace
```
Traceback (most recent call last):
  File "/path/to/my/amazing/virtual/environment/lib/python3.10/site-packages/torch/multiprocessing/spawn.py", line 69, in _wrap
    fn(i, *args)
  File "/path/to/my/code/train/run.py", line 36, in _train_wrapper
    train(rank, world_size, job, dgpm)
  File "/path/to/my/code/train/train.py", line 219, in train
    out = model(
  File "/path/to/my/amazing/virtual/environment/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/path/to/my/amazing/virtual/environment/lib/python3.10/site-packages/torch/nn/parallel/distributed.py", line 1040, in forward
    output = self._run_ddp_forward(*inputs, **kwargs)
  File "/path/to/my/amazing/virtual/environment/lib/python3.10/site-packages/torch/nn/parallel/distributed.py", line 1000, in _run_ddp_forward
    return module_to_run(*inputs[0], **kwargs[0])
  File "/path/to/my/amazing/virtual/environment/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/path/to/my/code/model/mask2former.py", line 35, in forward
    return self.model(pixel_values=pixel_values, mask_labels=mask_labels, class_labels=class_labels)
  File "/path/to/my/amazing/virtual/environment/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/path/to/my/amazing/virtual/environment/lib/python3.10/site-packages/transformers/models/mask2former/modeling_mask2former.py", line 2464, in forward
    loss_dict = self.get_loss_dict(
  File "/path/to/my/amazing/virtual/environment/lib/python3.10/site-packages/transformers/models/mask2former/modeling_mask2former.py", line 2351, in get_loss_dict
    loss_dict: Dict[str, Tensor] = self.criterion(
  File "/path/to/my/amazing/virtual/environment/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/path/to/my/amazing/virtual/environment/lib/python3.10/site-packages/transformers/models/mask2former/modeling_mask2former.py", line 792, in forward
    indices = self.matcher(masks_queries_logits, class_queries_logits, mask_labels, class_labels)
  File "/path/to/my/amazing/virtual/environment/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/path/to/my/amazing/virtual/environment/lib/python3.10/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/path/to/my/amazing/virtual/environment/lib/python3.10/site-packages/transformers/models/mask2former/modeling_mask2former.py", line 496, in forward
    assigned_indices: Tuple[np.array] = linear_sum_assignment(cost_matrix.cpu())
ValueError: cost matrix is infeasible
```

### Expected behavior

From what I can tell, this is not expected behavior and is caused by how `scipy.optimize.linear_sum_assignment` handles infinite values. Replacing these with very large numbers seems to fix the issue, as proposed [here](https://stackoverflow.com/questions/42035999/why-does-linear-sum-assignment-in-scipy-optimize-never-return-if-one-of-the-assi) (though for a slightly different issue). This is achieved by adding the following two lines above the call to `linear_sum_assignment` in the line linked earlier.
```py
cost_matrix = torch.minimum(cost_matrix, torch.tensor(1e10))
cost_matrix = torch.maximum(cost_matrix, torch.tensor(-1e10))
```

I have found that the error becomes more common as the learning rate is increased, so it could be related to diverging loss. However, I first discovered the error using the same learning rate as in the Mask2Former paper, `1e-4`, so I would not expect this to be too high, especially since it had happily chugged along for over 11k batches with dropping or steady loss before throwing the error.

If this is expected behavior, I at least think the error message should be improved.

Edit: I have fixed this locally by overwriting the `Mask2FormerHungarianMatcher` class with the aforementioned fix. I have not seen any diverging loss since then over many runs of thousands of epochs.

Comments:
Cc @alaradirik 

Hi @asgerius, thanks for opening the issue, I'm looking into this and will try to replicate the error first

Friendly ping @alaradirik :) 

Hi @asgerius, I tried fine-tuning Mask2Former on the semantic segmentation subset of the Scene Parsing dataset and couldn't replicate the issue. 

Is it possible that you are using a buggy version of scipy (the bug in scipy.optimize.linear_sum_assignment is fixed in this [PR](https://github.com/scipy/scipy/pull/7031)) or there are issues with the data preprocessing? 

You can refer to [Fine-tuning MaskFormer](https://pyimagesearch.com/2023/03/13/train-a-maskformer-segmentation-model-with-hugging-face-transformers/) blog post (exactly the same steps for Mask2Former) on PyImageSearch. I can take another look if the issue still persists but I'd need a minimal reproducible example to pinpoint the exact issue.

Hi
I put together a small example (see attached), and the results are somewhat contradictory to what I wrote in the original post. The error does indeed seem to be caused by diverging loss, which often recovers after a few batches. If I implement the described fix, the code does not crash, but instead just produces a high loss. However, I have never seen this behavior in my actual project, where the loss remains well-behaved when the fix is implemented. The only major difference is the data source, as the data in this example is simply random noise. I should also mention that my use case is two-class semantic segmentation.

Further, amp seems to be another important factor in addition to the learning rate. All my trainings have been run with it enabled. If disable in the example (controlled by the `use_amp` variable), the error becomes significantly harder to reproduce, indicating to me that the error is caused by overflowing floats.

To run the code, just put the files in the same directory and run `python example.py`. My versions of the dependencies are `torch==1.13.1 numpy==1.24.2 scipy==1.10.0 transformers==4.26.0`. With `use_amp = True` and `lr = 1e-4`, I usually get the error within the first 10-20 batches.

I changed the file types to .txt, as github does not allow .py and .json as attachments, so you'll have to change them back.
[example.txt](https://github.com/huggingface/transformers/files/11040427/example.txt)
[facebook.mask2former-swin-small-ade-semantic.config.txt](https://github.com/huggingface/transformers/files/11040428/facebook.mask2former-swin-small-ade-semantic.config.txt)



This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

@alaradirik I am facing the same issue, where I am using MaskFormer 
```

    914     cost_matrix = self.cost_mask * cost_mask + self.cost_class * cost_class + self.cost_dice * cost_dice
    915     # do the assigmented using the hungarian algorithm in scipy
--> 916     assigned_indices: Tuple[np.array] = linear_sum_assignment(cost_matrix.cpu())
    917     indices.append(assigned_indices)
    919 # It could be stacked in one tensor

ValueError: cost matrix is infeasible
```
I was following @NielsRogge tutorial for fine tuning on semantic masks where I only changed training to :

```
 with torch.cuda.amp.autocast():
            outputs = model(
                pixel_values=data["pixel_values"].to(device),
                mask_labels=[labels.to(device) for labels in data["mask_labels"]],
                class_labels=[labels.to(device) for labels in data["class_labels"]],
            )
            loss = outputs.loss

optimizer.zero_grad()
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
```
and my ground truths are only 0s and 1s i.e. binary masks
@alaradirik  I don't know how to replicate the issue as it don't know the real cause just occurs sometimes.

Small update: I have also seen the error during inference (amp enabled) when running on my trained model. However, this seems to be incredibly rare, as I have only ever experienced it once. I have not seen it after I implemented the fix described above in the inference code.