Skip to content

GitLab

Unverified Commit c0a354d8 authored by Chenhao Xu's avatar Chenhao Xu Committed by GitHub
Browse files

fix bug in mask2former: cost matrix is infeasible (#27897) 

fix bug: cost matrix is infeasible
parent 7e35f370
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 0 deletions
src/transformers/models/mask2former/modeling_mask2former.py
View file @ c0a354d8
...@@ -471,6 +471,9 @@ class Mask2FormerHungarianMatcher(nn.Module): ...@@ -471,6 +471,9 @@ class Mask2FormerHungarianMatcher(nn.Module):
cost_dice = pair_wise_dice_loss(pred_mask, target_mask) cost_dice = pair_wise_dice_loss(pred_mask, target_mask)
# final cost matrix # final cost matrix
cost_matrix = self.cost_mask * cost_mask + self.cost_class * cost_class + self.cost_dice * cost_dice cost_matrix = self.cost_mask * cost_mask + self.cost_class * cost_class + self.cost_dice * cost_dice
# eliminate infinite values in cost_matrix to avoid the error ``ValueError: cost matrix is infeasible``
cost_matrix = torch.minimum(cost_matrix, torch.tensor(1e10))
cost_matrix = torch.maximum(cost_matrix, torch.tensor(-1e10))
# do the assigmented using the hungarian algorithm in scipy # do the assigmented using the hungarian algorithm in scipy
assigned_indices: Tuple[np.array] = linear_sum_assignment(cost_matrix.cpu()) assigned_indices: Tuple[np.array] = linear_sum_assignment(cost_matrix.cpu())
indices.append(assigned_indices) indices.append(assigned_indices)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment