Spaces:
Running
Running
File size: 741 Bytes
e25cfe0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 |
from collections import OrderedDict
import numpy as np
import torch
import torch.nn as nn
import torchvision
class Normalizer(nn.Module):
def __init__(self):
super(Normalizer, self).__init__()
mean = np.array([0.485, 0.456, 0.406])
mean = mean[:, np.newaxis, np.newaxis]
std = np.array([0.229, 0.224, 0.225])
std = std[:, np.newaxis, np.newaxis]
# don't persist to keep old checkpoints working
self.register_buffer('mean', torch.tensor(mean), persistent=False)
self.register_buffer('std', torch.tensor(std), persistent=False)
def forward(self, tensor):
tensor = tensor / 255.0
tensor -= self.mean
tensor /= self.std
return tensor |