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
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.