#!/usr/bin/env python
# coding=utf-8
"""A network model that prefixes a z-normalization step to any other module"""
import torch
import torch.nn
[docs]class TorchVisionNormalizer(torch.nn.Module):
"""A simple normalizer that applies the standard torchvision normalization
This module does not learn.
The values applied in this "prefix" operator are defined at
https://pytorch.org/docs/stable/torchvision/models.html, and are as
follows:
* ``mean``: ``[0.485, 0.456, 0.406]``,
* ``std``: ``[0.229, 0.224, 0.225]``
"""
def __init__(self):
super(TorchVisionNormalizer, self).__init__()
mean = torch.as_tensor([0.485, 0.456, 0.406])[None, :, None, None]
std = torch.as_tensor([0.229, 0.224, 0.225])[None, :, None, None]
self.register_buffer("mean", mean)
self.register_buffer("std", std)
self.name = "torchvision-normalizer"
[docs] def forward(self, inputs):
return inputs.sub(self.mean).div(self.std)