Xla: Support in roi_align for MaskRCNN

[ofekp]

🚀 Feature

Today there is no way to train MaskRCNN from torchvision on TPU since there is no support in roi_align in XLA.

Motivation

MaskRCNN is widely used in segmentation models, we recently used it in our project and we could not run on TPU purely because we were lacking the implementation of roi_align. OK, there can be more methods we did not get to reach, but this one surely needs to be implemented.

Pitch

Implementation of roi_align in pytorch/xml.
This is the torchvision equivalent (which does not support XLA devices):
https://pytorch.org/docs/stable/torchvision/ops.html#torchvision.ops.roi_align
CPP implementation from torchvision:
https://github.com/pytorch/vision/blob/131ba1320b8208f10eb58d5feb7416c90ed839bb/torchvision/csrc/cuda/ROIAlign_cuda.cu#L64

Alternatives

An alternative might be to implement this using other methods in pytorch/xla which are already implemented, not sure it is feasible.

Additional context

I managed to switch box_ops.batched_nms with nms from torch_xla but then ran into the roi_align method which also has no support in XLA devices and is also not implemented in torch_xla

Thank you!

[oke-aditya]

It would be great to have TPU support for all the object detection, instance segmentation required modules in torchvision.

E.g. NMS, roi align etc. I'm unsure which are supported as of now.

[asuhan]

Such functionality can be either supported through interop (convert the XLA tensors to CPU, run the kernel, convert the result back to XLA) or by rewriting it in terms of tensors. Have a look at https://github.com/asuhan/maskrcnn-benchmark/blob/static_shapes/maskrcnn_benchmark/layers/tensor_roi_align.py.

[ofekp]

@asuhan I don't have much experience with training on TPU, but I feel like converting to CPU and then back to tensor is missing the point of speeding up training. If the model is trained fast enough even with such modification then please let me know, but I think maybe the conversion and network latency may cause it to be very slow. Did you test this? Do you know what overhead it introduces to the training time (of course it is model dependent but any comparison is better than nothing)?

[asuhan]

@ofekp It all depends on how expensive that particular operation is. In the case of ROI align, yes, it'd be very expensive - a lot of data to transfer, which is why I re-implemented it in terms of tensors. In other cases, though, it might basically be fast enough. It all depends on the concrete workload.

[ofekp]

@asuhan Can you please explain how I can use your implementation while replacing the torchvision implementation which does not support XLA? 🙏

[asuhan]

This is what I did: https://github.com/asuhan/maskrcnn-benchmark/blob/753e5d2d0285a6175340e38ca203d899cc07916c/maskrcnn_benchmark/layers/roi_align.py#L59, AFAICT you could make the same change here: https://github.com/pytorch/vision/blob/a4736ea6e8ff25a1265ca5adf9e6e244d78500e8/torchvision/ops/roi_align.py#L71.

[stale[bot]]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.

[maramir6]

Has been any improvement in this topic?