use iou scaling to avoid errors with torch ddp

Rustem Galiullin · Rustem Galiullin · commit 4d1144e88bee · 2023-05-08T16:41:11.000+04:00
diff --git a/segmentation_models_pytorch/decoders/sam/model.py b/segmentation_models_pytorch/decoders/sam/model.py
@@ -194,5 +194,7 @@ def forward(self, x):
             multimask_output=self._decoder_multiclass_output,
         )
         masks = self.postprocess_masks(low_res_masks, input_size=img_size, original_size=img_size)
+        # use scaling below in order to make it work with torch DDP
+        masks = masks * iou_predictions.view(-1, masks.size(1), 1, 1)
         output = self.segmentation_head(masks)
         return output

Original file line number	Diff line number	Diff line change
`@@ -194,5 +194,7 @@ def forward(self, x):`
`194`	`194`	`multimask_output=self._decoder_multiclass_output,`
`195`	`195`	`)`
`196`	`196`	`masks = self.postprocess_masks(low_res_masks, input_size=img_size, original_size=img_size)`
	`197`	`+ # use scaling below in order to make it work with torch DDP`
	`198`	`+ masks = masks * iou_predictions.view(-1, masks.size(1), 1, 1)`
`197`	`199`	`output = self.segmentation_head(masks)`
`198`	`200`	`return output`