ToTensor should support conversion to tensors of any dtype and device

🚀 Feature

torchvision.transforms.ToTensor should support parameters to specify the dtype and target device of the tensors produced.

Motivation

Currently, ToTensor always converts images to FP32 tensors in the range [0, 1] if they have 8-bit values. This behavior is inefficient for use cases where the images need to be converted to another data type (such as FP16 or bfloat16) and/or transferred to an accelerator such as a GPU. Two cases come to mind:

• Case 1: We want to load images with 8-bit values into PyTorch and convert them to HalfTensors.
  • First, ToTensor converts the images into FloatTensors, quadrupling the memory they take up. Then, the calling code converts the images from FloatTensors to HalfTensors, which are half as large as the FloatTensors.
  • It would be better not to have to allocate the memory used to store the FloatTensors, and go straight from the original image format to HalfTensors.
• Case 2: We want to load images with 8-bit values onto a GPU as FloatTensors.
  • First, ToTensor converts the images into FloatTensors in the range [0, 1], so they take up 4 times as much memory as they did before. Then, this data must be transferred to the GPU.
  • It would be better to convert the images into ByteTensors, then move them to the GPU and convert them to FloatTensors in the range [0, 1] on the GPU. This would use 1/4 of the data transfer bandwidth as the approach outlined above, and because the division by 255 is done in parallel on the GPU, this step is much faster. I've written code that implements this and have noticed slight improvements in performance as a result.

Pitch

At the least, we should add dtype and device parameters to the ToTensor constructor, like those for Tensor.to(), which would determine the dtype and device of the tensors produced by this transform. If the dtype is an integer or boolean data type, we don't divide by 255. We can copy other parameters from Tensor.to() if they would be helpful.

Alternatives

I've written custom code using transforms.Lambda that converts images to FP32 tensors after moving them to the GPU, but I think this functionality would be more widely available if it were built into ToTensor.

Additional context

cc @vfdev-5

Start with torchvision.transforms.ToTensor and compare the requested constructor options with the linked Tensor.to behavior. Define and test dtype and device handling, including integer and boolean dtypes that should not be divided by 255; done means the transform produces tensors with the requested dtype and device.