[WIP] Allow autocast for 1.6

[mcarilli]

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 pytorch/pytorch#37735.

[fmassa]

Thanks for the PR!

Can you add some full model-forward test in a similar location as

vision/test/test_models.py

Lines 273 to 293 in 5247f7b

	@unittest.skipIf(not torch.cuda.is_available(), 'needs GPU')
	def test_fasterrcnn_switch_devices(self):
	model = models.detection.fasterrcnn_resnet50_fpn(num_classes=50, pretrained_backbone=False)
	model.cuda()
	model.eval()
	input_shape = (3, 300, 300)
	x = torch.rand(input_shape, device='cuda')
	model_input = [x]
	out = model(model_input)
	self.assertIs(model_input[0], x)
	self.assertEqual(len(out), 1)
	self.assertTrue("boxes" in out[0])
	self.assertTrue("scores" in out[0])
	self.assertTrue("labels" in out[0])
	# now switch to cpu and make sure it works
	model.cpu()
	x = x.cpu()
	out_cpu = model([x])
	self.assertTrue("boxes" in out_cpu[0])
	self.assertTrue("scores" in out_cpu[0])
	self.assertTrue("labels" in out_cpu[0])

and also add a forward-backward test for roi_align after

vision/test/test_ops.py

Lines 290 to 291 in 5247f7b

	def _test_boxes_shape(self):
	self._helper_boxes_shape(ops.roi_align)

(and we will clean it up later on when adding support for the other ops to autocast)

[codecov]

Codecov Report

Merging #2384 into master will increase coverage by 0.70%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##           master    #2384      +/-   ##
==========================================
+ Coverage   70.65%   71.36%   +0.70%     
==========================================
  Files          94       94              
  Lines        7897     8328     +431     
  Branches     1241     1385     +144     
==========================================
+ Hits         5580     5943     +363     
- Misses       1934     1972      +38     
- Partials      383      413      +30     

Impacted Files	Coverage Δ
torchvision/ops/poolers.py	97.02% <100.00%> (ø)
torchvision/io/image.py	71.73% <0.00%> (+0.77%)	⬆️
torchvision/transforms/functional_tensor.py	65.41% <0.00%> (+0.83%)	⬆️
torchvision/transforms/functional.py	81.91% <0.00%> (+2.09%)	⬆️
torchvision/transforms/transforms.py	78.25% <0.00%> (+2.23%)	⬆️

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 86b6c3e...f3451c9. Read the comment docs.

[fmassa]

Test failures seem related

[fmassa]

Thanks a lot Michael!