Feature Proposal: -
Torchvision Object detection models should support TPU.
TPUs will greatly lower training time for models like Faster RCNN and coming soon Retina Net. Torch natively supports TPU through xla library. It is possible to train the classification models on TPU as well. But when we try object detection models on TPU it gives the following error
File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 550, in __call__
result = self.forward(*input, **kwargs)
File "/usr/local/lib/python3.6/dist-packages/torchvision/models/detection/generalized_rcnn.py", line 70, in forward
proposals, proposal_losses = self.rpn(images, features, targets)
File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 550, in __call__
result = self.forward(*input, **kwargs)
File "/usr/local/lib/python3.6/dist-packages/torchvision/models/detection/rpn.py", line 493, in forward
boxes, scores = self.filter_proposals(proposals, objectness, images.image_sizes, num_anchors_per_level)
File "/usr/local/lib/python3.6/dist-packages/torchvision/models/detection/rpn.py", line 416, in filter_proposals
keep = box_ops.batched_nms(boxes, scores, lvl, self.nms_thresh)
RuntimeError: Cannot access data pointer of Tensor that doesn't have storage
This error is because some features such rpn.py , ops.batched_nms is not TPU compatible.
Will we get any support for TPU soon for these models ?
I created an issue on GitHub as well but I got no response. So I thought maybe I will add here to discuss.