Initialize MeanStdNormalizer.
(
self,
mean: Union[Sequence[float], torch.Tensor],
std: Union[Sequence[float], torch.Tensor],
)
| 19 | std_inv: torch.Tensor |
| 20 | |
| 21 | def __init__( |
| 22 | self, |
| 23 | mean: Union[Sequence[float], torch.Tensor], |
| 24 | std: Union[Sequence[float], torch.Tensor], |
| 25 | ): |
| 26 | """Initialize MeanStdNormalizer.""" |
| 27 | super(MeanStdNormalizer, self).__init__() |
| 28 | if not isinstance(mean, torch.Tensor): |
| 29 | mean = torch.as_tensor(mean).view(-1, 1, 1) |
| 30 | if not isinstance(std, torch.Tensor): |
| 31 | std = torch.as_tensor(std).view(-1, 1, 1) |
| 32 | self.register_buffer("mean", mean) |
| 33 | # We use inverse std to use a multiplication which is better supported by the hardware |
| 34 | self.register_buffer("std_inv", 1.0 / std) |
| 35 | |
| 36 | def forward(self, image: torch.Tensor) -> torch.Tensor: |
| 37 | """Apply mean and std normalization over input image.""" |