Vectorize the equalize transformation

[datumbox]

🚀 Feature

The current implementation of the equalize transformation is not vectorized. It processes each channel of each image of the batch using a for loop. This is done because the method uses internally the torch.histc which currently can't produce a histogram across a dimension.

vision/torchvision/transforms/functional_tensor.py

Lines 1281 to 1282 in 83171d6

	def _scale_channel(img_chan):
	hist = torch.histc(img_chan.to(torch.float32), bins=256, min=0, max=255)

@fmassa proposed a workaround that trades memory for speed and achieves vectorization. Check #3123 (comment) for the overview/context of the proposal.

To adopt his, a few addtional changes need to be made:

• The rest of the _scale_channel() needs to be adapted to vectorize the remaining operations.
• The _equalize_single_image() needs to be removed and the stacking operation in equalize() need to be adapted.

cc @vfdev-5 @fmassa

[avijit9]

@datumbox Can I try to solve this one, if nobody is working on it?

[datumbox]

Go for it @avijit9 !

[NicolasHug]

According to #3334, vectorizing the histogram computation seems to lead to a slow-down. I'll close the issue