Repository: pytorch/vision
Pull request: [WIP] Allow autocast for 1.6
This PR introduces a minimal set of diffs that enable torchvision models to work with `torch.cuda.amp.autocast`.  Custom C++ ops (`roi_align` and `nms`) require the most attention.

In later modifications to Pytorch, I'll allow external libs to use Pytorch's internal autocast utilities, after which this code can be made cleaner.  For 1.6, however, copy pasting some utilities is the best we can do.

Should close https://github.com/pytorch/pytorch/issues/37735.

Problem statement:
torch.cuda.amp.autocast not working with torchvision.models.detection.maskrcnn
## 🐛 Bug

I am converting the model into FP16.

Using `torch.cuda.amp.autocast`. But it throws me an error:

```
~/Documents/test/seg/models/archs/mask_rcnn.py in mixed_precision_one_batch(self, i, b)
    186         with autocast():
    187             self.model.train()
--> 188             loss_dict = self.model(images,targets)
    189             if not self.training:
    190                 self.model.eval()

~/anaconda3/envs/seg/lib/python3.7/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    556             result = self._slow_forward(*input, **kwargs)
    557         else:
--> 558             result = self.forward(*input, **kwargs)
    559         for hook in self._forward_hooks.values():
    560             hook_result = hook(self, input, result)

~/anaconda3/envs/seg/lib/python3.7/site-packages/torchvision/models/detection/generalized_rcnn.py in forward(self, images, targets)
     68         if isinstance(features, torch.Tensor):
     69             features = OrderedDict([('0', features)])
---> 70         proposals, proposal_losses = self.rpn(images, features, targets)
     71         detections, detector_losses = self.roi_heads(features, proposals, images.image_sizes, targets)
     72         detections = self.transform.postprocess(detections, images.image_sizes, original_image_sizes)

~/anaconda3/envs/seg/lib/python3.7/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    556             result = self._slow_forward(*input, **kwargs)
    557         else:
--> 558             result = self.forward(*input, **kwargs)
    559         for hook in self._forward_hooks.values():
    560             hook_result = hook(self, input, result)

~/anaconda3/envs/seg/lib/python3.7/site-packages/torchvision/models/detection/rpn.py in forward(self, images, features, targets)
    486         proposals = self.box_coder.decode(pred_bbox_deltas.detach(), anchors)
    487         proposals = proposals.view(num_images, -1, 4)
--> 488         boxes, scores = self.filter_proposals(proposals, objectness, images.image_sizes, num_anchors_per_level)
    489 
    490         losses = {}

~/anaconda3/envs/seg/lib/python3.7/site-packages/torchvision/models/detection/rpn.py in filter_proposals(self, proposals, objectness, image_shapes, num_anchors_per_level)
    408             boxes, scores, lvl = boxes[keep], scores[keep], lvl[keep]
    409             # non-maximum suppression, independently done per level
--> 410             keep = box_ops.batched_nms(boxes, scores, lvl, self.nms_thresh)
    411             # keep only topk scoring predictions
    412             keep = keep[:self.post_nms_top_n()]

~/anaconda3/envs/seg/lib/python3.7/site-packages/torchvision/ops/boxes.py in batched_nms(boxes, scores, idxs, iou_threshold)
     73     offsets = idxs.to(boxes) * (max_coordinate + 1)
     74     boxes_for_nms = boxes + offsets[:, None]
---> 75     keep = nms(boxes_for_nms, scores, iou_threshold)
     76     return keep
     77 

~/anaconda3/envs/seg/lib/python3.7/site-packages/torchvision/ops/boxes.py in nms(boxes, scores, iou_threshold)
     33         by NMS, sorted in decreasing order of scores
     34     """
---> 35     return torch.ops.torchvision.nms(boxes, scores, iou_threshold)
     36 
     37 

RuntimeError: Unrecognized tensor type ID: Autocast
```

cc @ngimel

Hints from issue discussion:
cc @mcarilli 
> RuntimeError: Unrecognized tensor type ID: Autocast

That’s a strange error. If you’re using a Pytorch build new enough to support autocast, AutocastTensorId should be present in the backend. If you’re not, I’d expect a different error, sooner (torch.cuda.amp.autocast should not be available in Python to begin with).

What version of Pytorch are you using?
@mcarilli This is the version:

```
import torchvision
print(torch.__version__)
print(torchvision.__version__)

1.6.0.dev20200421
0.7.0a0+6e47842
```
Looking at source of torch.nn.Module, it appears [_slow_forward](https://github.com/pytorch/pytorch/blob/5ec87a3c1a6acfb0eb4704873649ffcd7e65c221/torch/nn/modules/module.py#L533) path is only taken when jit tracing is enabled.  Are you using jit tracing explicitly in your script?
> 
> 
> Looking at source of torch.nn.Module, it appears [_slow_forward](https://github.com/pytorch/pytorch/blob/5ec87a3c1a6acfb0eb4704873649ffcd7e65c221/torch/nn/modules/module.py#L533) path is only taken when jit tracing is enabled. Are you using jit tracing explicitly in your script?

@mcarilli I don't know if it is inside the model. I think that it saw in source code.

However, i don't know exactly what is jit tracing. In my trining loop i don't see it. I use the model as it is in torchvision

I am using the Mask-RCNN implementation from torchvision. 
Here is the [docs](https://pytorch.org/docs/stable/torchvision/models.html#mask-r-cnn).

And the source code is in [here](https://github.com/pytorch/vision/blob/master/torchvision/models/detection/mask_rcnn.py)



@mcarilli Any update about your progress?
I tried some repro here https://github.com/xwang233/code-snippet/tree/master/mask-rcnn-37735. I get the same error message in 4/21 nightly build of pytorch and torchvision. 

The error appears to be gone in 6/22 nightly build of pytorch and torchvision, but it yields a different error message. Here is the backtrace
```
$ python repro.py
1.6.0a0+18122fa
0.7.0a0+42aa9b2
/home/xwang/Developer/pytorch/torch/nn/functional.py:3000: UserWarning: The default behavior for interpolate/upsample with float scale_factor will change in 1.6.0 to align with other frameworks/libraries, and use scale_factor directly, instead of relying on the computed output size. If you wish to keep the old behavior, please set recompute_scale_factor=True. See the documentation of nn.Upsample for details. 
  warnings.warn("The default behavior for interpolate/upsample with float scale_factor will change "
/home/xwang/Developer/vision/torchvision/ops/boxes.py:102: UserWarning: This overload of nonzero is deprecated:
        nonzero()
Consider using one of the following signatures instead:
        nonzero(*, bool as_tuple) (Triggered internally at  ../torch/csrc/utils/python_arg_parser.cpp:766.)
  keep = keep.nonzero().squeeze(1)
Traceback (most recent call last):
  File "repro.py", line 44, in <module>
    loss_dict = model(d, tds)
  File "/home/xwang/Developer/pytorch/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/xwang/Developer/vision/torchvision/models/detection/generalized_rcnn.py", line 99, in forward
    detections, detector_losses = self.roi_heads(features, proposals, images.image_sizes, targets)
  File "/home/xwang/Developer/pytorch/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/xwang/Developer/vision/torchvision/models/detection/roi_heads.py", line 752, in forward
    box_features = self.box_roi_pool(features, proposals, image_shapes)
  File "/home/xwang/Developer/pytorch/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/xwang/Developer/vision/torchvision/ops/poolers.py", line 222, in forward
    spatial_scale=scale, sampling_ratio=self.sampling_ratio)
  File "/home/xwang/Developer/vision/torchvision/ops/roi_align.py", line 45, in roi_align
    sampling_ratio, aligned)
RuntimeError: Expected tensor for argument #1 'input' to have the same type as tensor for argument #2 'rois'; but type torch.cuda.HalfTensor does not equal torch.cuda.FloatTensor (while checking arguments for ROIAlign_forward_cuda)
```

@mcarilli has more information. 
The new backtrace makes sense:  some C++ function is receiving mismatched arguments.  I'll track down the implementation.  If it's a custom autograd function, we'll need to apply [torch.cuda.amp.custom_fwd/bwd](https://pytorch.org/docs/master/amp.html#torch.cuda.amp.custom_fwd).  If it's a torch backend function, I'll need to add it to the Amp promote list, or possibly the FP32 list or FP16 list, if the op has a preferred precision.  @fmassa do you know if this op has a preferred precision?

update:  looks like torchvision [attaches custom ops to `torch`](https://github.com/pytorch/vision/blob/1a04d3c265679e1a508e7cd627006aaa9ef1ccfb/torchvision/csrc/vision.cpp#L45-L55) via a newish extension API that registers external C++ functions with the dispatcher.  In other words it does _not_ use conventional Python-side custom autograd functions, so torch.cuda.amp.custom_fwd/bwd decorators likely won't help.  Will talk to @fmassa about how best to cover this case in torchvision.
After discussion with @ezyang we have a tentative plan to fix. Ed is laying the foundation in https://github.com/pytorch/vision/pull/2366, which updates `roi_align` to allow Pytorch's internal dispatcher to select different code paths.  The dispatcher knows about autocast, so once Ed's diffs are merged I can register a code path that casts inputs if autocast is active, if @fmassa confirms this is a sensible approach.

If approved, the fix will be in torchvision, not pytorch, therefore torchvision-specific.  Other external libs that use the new-style C++ registration API may still break with autocast, and it's unclear what a generally autocast-safe registration API might look like.  However, a quick-and-dirty torchvision-specific fix is probably the best we can do for 1.6 and it's not a bad plan:
- I expect few external libs use the new-style registrations yet.
- the work gives us (meaning mostly Ed) a clear example of what external libs need for native-like interaction with the dispatcher.  Hacking on it helps inform registration API improvements.
This approach seems good to me

Commit message: Fixes Xiao's repro

Patch to review:
diff --git a/torchvision/csrc/ROIAlign.h b/torchvision/csrc/ROIAlign.h
index 7a856f34d63..a40833e29df 100644
--- a/torchvision/csrc/ROIAlign.h
+++ b/torchvision/csrc/ROIAlign.h
@@ -4,6 +4,7 @@
 
 #ifdef WITH_CUDA
 #include "cuda/vision_cuda.h"
+#include "autocast.h"
 #endif
 #ifdef WITH_HIP
 #include "hip/vision_cuda.h"
@@ -35,6 +36,25 @@ at::Tensor roi_align(
       aligned);
 }
 
+at::Tensor ROIAlign_autocast(
+    const at::Tensor& input,
+    const at::Tensor& rois,
+    const double spatial_scale,
+    const int64_t pooled_height,
+    const int64_t pooled_width,
+    const int64_t sampling_ratio,
+    const bool aligned) {
+  c10::impl::ExcludeDispatchKeyGuard no_autocast(c10::DispatchKey::Autocast);
+  return roi_align(
+      autocast::_cast(at::kFloat, input),
+      autocast::_cast(at::kFloat, rois),
+      spatial_scale,
+      pooled_height,
+      pooled_width,
+      sampling_ratio,
+      aligned);
+}
+
 at::Tensor _roi_align_backward(
     const at::Tensor& grad,
     const at::Tensor& rois,
diff --git a/torchvision/csrc/autocast.h b/torchvision/csrc/autocast.h
new file mode 100644
index 00000000000..4189f6386a0
--- /dev/null
+++ b/torchvision/csrc/autocast.h
@@ -0,0 +1,26 @@
+#pragma once
+
+#ifdef WITH_CUDA
+namespace autocast {
+
+inline bool is_eligible(const at::Tensor& arg) {
+  return (arg.is_cuda() && arg.is_floating_point() && (arg.scalar_type() != at::kDouble));
+}
+
+// Overload to catch Tensor args
+inline at::Tensor _cast(at::ScalarType to_type, const at::Tensor& arg) {
+  if (is_eligible(arg) && (arg.scalar_type() != to_type)) {
+    return arg.to(to_type);
+  } else {
+    return arg;
+  }
+}
+
+// Template to catch non-Tensor args
+template<typename T>
+inline T _cast(at::ScalarType to_type, T arg) {
+  return arg;
+}
+
+}
+#endif
diff --git a/torchvision/csrc/vision.cpp b/torchvision/csrc/vision.cpp
index 7f56bdb51a1..4af6a4cd4b0 100644
--- a/torchvision/csrc/vision.cpp
+++ b/torchvision/csrc/vision.cpp
@@ -69,6 +69,13 @@ TORCH_LIBRARY_IMPL(torchvision, CUDA, m) {
 }
 #endif
 
+// Autocast only needs to wrap forward pass ops.
+#if defined(WITH_CUDA)
+TORCH_LIBRARY_IMPL(torchvision, Autocast, m) {
+  m.impl("roi_align", ROIAlign_autocast);
+}
+#endif
+
 TORCH_LIBRARY_IMPL(torchvision, Autograd, m) {
   m.impl("roi_align", ROIAlign_autograd);
   m.impl("_roi_align_backward", ROIAlign_backward_autograd);
diff --git a/torchvision/ops/poolers.py b/torchvision/ops/poolers.py
index 06bbc86a93c..0c9578007a7 100644
--- a/torchvision/ops/poolers.py
+++ b/torchvision/ops/poolers.py
@@ -224,7 +224,7 @@ def forward(self, x, boxes, image_shapes):
             if torchvision._is_tracing():
                 tracing_results.append(result_idx_in_level.to(dtype))
             else:
-                result[idx_in_level] = result_idx_in_level
+                result[idx_in_level] = result_idx_in_level.to(result.dtype)
 
         if torchvision._is_tracing():
             result = _onnx_merge_levels(levels, tracing_results)


Code review comments:
diff_hunk: @@ -4,6 +4,7 @@
 
 #ifdef WITH_CUDA
 #include "cuda/vision_cuda.h"
+#include "autocast.h"
path: torchvision/csrc/ROIAlign.h
text: @user1:
Tests are failing because this is only included on CUDA builds, but the `autocast` is used unconditionally.

Maybe you could add another `ifdef WITH_CUDA` around `ROIALIGN_autocast` ?

@author:
Good catch, thanks!