`transforms.v2.Normalize` caused small `H2D` and cpu waits gpu if process image on gpu

[NewUserHa]

🚀 The feature

vision/torchvision/transforms/v2/_misc.py

Lines 142 to 175 in 10f68db

	class Normalize(Transform):
	"""Normalize a tensor image or video with mean and standard deviation.
	This transform does not support PIL Image.
	Given mean: ``(mean[1],...,mean[n])`` and std: ``(std[1],..,std[n])`` for ``n``
	channels, this transform will normalize each channel of the input
	``torch.*Tensor`` i.e.,
	``output[channel] = (input[channel] - mean[channel]) / std[channel]``
	.. note::
	This transform acts out of place, i.e., it does not mutate the input tensor.
	Args:
	mean (sequence): Sequence of means for each channel.
	std (sequence): Sequence of standard deviations for each channel.
	inplace(bool,optional): Bool to make this operation in-place.
	"""
	_v1_transform_cls = _transforms.Normalize
	def __init__(self, mean: Sequence[float], std: Sequence[float], inplace: bool = False):
	super().__init__()
	self.mean = list(mean)
	self.std = list(std)
	self.inplace = inplace
	def check_inputs(self, sample: Any) -> Any:
	if has_any(sample, PIL.Image.Image):
	raise TypeError(f"{type(self).__name__}() does not support PIL images.")
	def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
	return self._call_kernel(F.normalize, inpt, mean=self.mean, std=self.std, inplace=self.inplace)

This Normalize class can support tensors as init arguments just like transforms.v2.functional.normalize_image, and save it for later use instread of create it everytime, becausse it will degrade performance especially for on gpu.

Motivation, pitch

I encountered a performance issue which slow down the program. When process images on gpu, the transforms.v2.Normalize will creates tensor on the fly and causes small H2D pageable-> devices and cpu stop to wait gpu with CudaStreamSynchronize before send following kernels.

It not only cause H2D in the middle but also create tensor repeatedly which is unnesscerry. Plus the transforms.v2.functional.normalize_image support tensor as input at first.

I can make a PR

Alternatives

No response

Additional context

No response