Spaces:
Running
Running
File size: 2,567 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 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | from collections import OrderedDict
import numpy as np
import torch
import torch.nn as nn
import torchvision
class VGGInputNormalization(torch.nn.Module):
def __init__(self, inplace=True):
super().__init__()
self.inplace = inplace
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]
self.register_buffer('mean', torch.tensor(mean))
self.register_buffer('std', torch.tensor(std))
def forward(self, tensor):
if self.inplace:
tensor /= 255.0
else:
tensor = tensor / 255.0
tensor -= self.mean
tensor /= self.std
return tensor
class VGG19BNNamedFeatures(torch.nn.Sequential):
def __init__(self):
names = []
for block in range(5):
block_size = 2 if block < 2 else 4
for layer in range(block_size):
names.append(f'conv{block+1}_{layer+1}')
names.append(f'bn{block+1}_{layer+1}')
names.append(f'relu{block+1}_{layer+1}')
names.append(f'pool{block+1}')
vgg = torchvision.models.vgg19_bn(pretrained=True)
vgg_features = vgg.features
vgg.classifier = torch.nn.Sequential()
assert len(names) == len(vgg_features)
named_features = OrderedDict({'normalize': VGGInputNormalization()})
for name, feature in zip(names, vgg_features):
if isinstance(feature, nn.MaxPool2d):
feature.ceil_mode = True
named_features[name] = feature
super().__init__(named_features)
class VGG19NamedFeatures(torch.nn.Sequential):
def __init__(self):
names = []
for block in range(5):
block_size = 2 if block < 2 else 4
for layer in range(block_size):
names.append(f'conv{block+1}_{layer+1}')
names.append(f'relu{block+1}_{layer+1}')
names.append(f'pool{block+1}')
vgg = torchvision.models.vgg19(pretrained=True)
vgg_features = vgg.features
vgg.classifier = torch.nn.Sequential()
assert len(names) == len(vgg_features)
named_features = OrderedDict({'normalize': VGGInputNormalization()})
for name, feature in zip(names, vgg_features):
if isinstance(feature, nn.MaxPool2d):
feature.ceil_mode = True
named_features[name] = feature
super().__init__(named_features)
|