Torchvision vs Detectron2 MaskRCNN

[isaaccorley]

Hi! I see that the pretrained MaskRCNN models use the detectron2 library.

For various reasons, I prefer not to use detectron2 and would rather use torchvision's MaskRCNN implementation.

I see that you have training code here for torchvision's MaskRCNN. Did you find detectron2 to perform better for the paper results?

[jveitchmichaelis]

I wanted to use Torchvision's implementation, but I ran into some significant difficulty training it due to the way the loss function is calculated for huge numbers of objects. I kept getting out of memory errors on the CPU side. I'm not sure how Detectron2 avoids this, but I think it's related to how the prediction masks are represented and the RAM involved in manipulating them. The same happened withtransformers actually, the Hungarian matching loss blows up RAM consumption.

Similarly to the other issue, if you'd be willing to try training it I can support. This is also relevant for the current project I'm working on at Weecology (DeepForest) because we'll probably hit the same issues.

That file is a bit of a monster 😅 let me have a look.

[isaaccorley]

Yes, I think you're right that for torchvision instance segmentation you have to construct a mask of shape (N, H,W), where N=num_objects, for every image in your batch and if you have many objects per in a single image this consumes a ton of memory. I'll take a stab at training since I do have a machine with 256GB RAM and a 48GB A100 maybe this won't be an issue.