Specify channel dim for transforms.Normalize

🚀 The feature

Specify channel dim for transforms.Normalize, transforms.functional.normalize, transforms.functional_tensor.normalize, To enable transforms.Normalize to normalize according mean and std by specified channel.

A solution is adding a new argument dim_channel to the classes and functions above and

# in transforms.functional_tensor.normalize
broadcast_ch_shape = [1 for _ in range(tensor.ndim)]
broadcast_ch_shape[dim_channel] = -1
if mean.ndim == 1:
    mean = mean.view(*broadcast_ch_shape)
if std.ndim == 1:
    std = std.view(*broadcast_ch_shape)
return tensor.sub_(mean).div_(std)

Motivation, pitch

Recent torchvision deprecated transforms._transforms_video and added features in many transforms to process [..., H, W] shaped tensors. For video transforming, it is a great improvement, meanwhile, transforms.Normalize is not lucky enough to be among these transforms. This means that the users either resort to other transforms such as pytorchvideo.transforms.Normalize or normalize each frame seperately. The requested feature will relieve this pain, and video transforms can be more nice and neat.

Alternatives

No response

Additional context

No response

cc @vfdev-5 @datumbox

Start with the mentioned transforms.Normalize, transforms.functional.normalize, and transforms.functional_tensor.normalize entry points, then trace how mean and std are currently broadcast. Check the existing transform behavior for image tensors and determine how a specified channel dimension should work for video-shaped tensors; done means all three APIs support the requested channel dimension consistently.