diff --git a/test/test_transforms_v2.py b/test/test_transforms_v2.py index 5f91ee066d4..b1ba76632af 100644 --- a/test/test_transforms_v2.py +++ b/test/test_transforms_v2.py @@ -5670,6 +5670,18 @@ def test_transform(self, make_input): check_sample_input=self._sample_input_adapter, ) + def test_transform_accepts_tensor_statistics(self): + mean = torch.tensor(self.MEAN) + std = torch.tensor(self.STD) + transform = transforms.Normalize(mean=mean, std=std) + + assert transform.mean is mean + assert transform.std is std + image = make_image(dtype=torch.float32) + actual = transform(image) + expected = F.normalize_image(image, mean=mean, std=std) + assert_equal(actual, expected) + def _reference_normalize_image(self, image, *, mean, std): image = image.numpy() mean, std = (np.array(stat, dtype=image.dtype).reshape((-1, 1, 1)) for stat in [mean, std]) diff --git a/torchvision/transforms/v2/_misc.py b/torchvision/transforms/v2/_misc.py index 305149c87b1..38df7394d5c 100644 --- a/torchvision/transforms/v2/_misc.py +++ b/torchvision/transforms/v2/_misc.py @@ -152,18 +152,19 @@ class Normalize(Transform): 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. + mean (sequence or Tensor): Sequence of means for each channel. A tensor can + be provided to keep the statistics on the same device as the input. + std (sequence or Tensor): 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): + def __init__(self, mean: Sequence[float] | torch.Tensor, std: Sequence[float] | torch.Tensor, inplace: bool = False): super().__init__() - self.mean = list(mean) - self.std = list(std) + self.mean = mean if isinstance(mean, torch.Tensor) else list(mean) + self.std = std if isinstance(std, torch.Tensor) else list(std) self.inplace = inplace def check_inputs(self, sample: Any) -> Any: