mask2former head fails during training because of amp (ValueError: cost matrix is infeasible)

[we11as22]

Thanks for your error report and we appreciate it a lot.

Checklist

1. I have searched related issues but cannot get the expected help.
2. I have read the FAQ documentation but cannot get the expected help.
3. The bug has not been fixed in the latest version.

Describe the bug
A clear and concise description of what the bug is.

Reproduction

1. What command or script did you run?

A placeholder for the command.

2. Did you make any modifications on the code or config? Did you understand what you have modified?
3. What dataset did you use?

Environment

1. Please run python mmdet/utils/collect_env.py to collect necessary environment information and paste it here.
2. You may add addition that may be helpful for locating the problem, such as
   • How you installed PyTorch [e.g., pip, conda, source]
   • Other environment variables that may be related (such as $PATH, $LD_LIBRARY_PATH, $PYTHONPATH, etc.)

Error traceback
If applicable, paste the error trackback here.

A placeholder for trackback.

Bug fix
If you have already identified the reason, you can provide the information here. If you are willing to create a PR to fix it, please also leave a comment here and that would be much appreciated!

[hhaAndroid]

@AlexeySudakovB01-109 Thank you for the feedback. We will work on fixing it as soon as possible.

[liuxingbin]

any update?

[leo-q8]

Looking forward to further work！

[leo-q8]

@we11as22 @hhaAndroid I am facing the same issue, when training Mask2Former with --amp. And I have found a solution here huggingface/transformers#21644. As it proposed, this is caused by how scipy.optimize.linear_sum_assignment handles infinite values. Replacing these with very large numbers seems to fix the issue. And I added the following code before line 125.

cost = torch.minimum(cost, torch.tensor(1e10))
cost = torch.maximum(cost, torch.tensor(-1e10))

mmdetection/mmdet/models/task_modules/assigners/hungarian_assigner.py

Lines 123 to 131 in cfd5d3a

	cost = torch.stack(cost_list).sum(dim=0)
	# 3. do Hungarian matching on CPU using linear_sum_assignment
	cost = cost.detach().cpu()
	if linear_sum_assignment is None:
	raise ImportError('Please run "pip install scipy" '
	'to install scipy first.')
	matched_row_inds, matched_col_inds = linear_sum_assignment(cost)

But when trying to train with amp again, I encountered the NaN mask loss problem as follows:

Further debugging revealed issues in the calculation of loss_mask in this line. In amp mode, multiplying the FP16 variable num_total_masks by a large integer variable self.num_points (it is 12544 in my configs)results in an overflowed FP16 variable, further leading to the occurrence of NaN mask loss.

By casting num_total_masks to the fp32 type with avg_factor=num_total_masks.float() * self.num_points, I managed to eliminate this issue. The model can now be successfully trained in amp mode, but the accuracy is still under validation 😂 .
Hope this is helpful to you！

[FreeWillThorn]

@leo-q8 Or you can do like this:

replace the optimizer in the config file by :

optim_wrapper = dict( type='AmpOptimWrapper', optimizer=dict( type='AdamW', lr=0.00001, weight_decay=0.05, eps=1e-8, betas=(0.9, 0.999) ), clip_grad=dict(max_norm=0.01, norm_type=2), dtype='float32' # Use float32 precision )

[andrewcaunes]

@FreeWillThorn Isn't the whole point of AmpOptimWrapper to use FP16 precision ?