66 lines
2.1 KiB
Python
Executable File
66 lines
2.1 KiB
Python
Executable File
import torch.nn as nn
|
|
import utils.vgg_util as vgg_util
|
|
from torchvision import models
|
|
|
|
|
|
def conv_to_relu(layer_names):
|
|
out_layernames = []
|
|
for name in layer_names:
|
|
if name.startswith("conv"):
|
|
out_layernames.append("relu" + name[4:])
|
|
else:
|
|
out_layernames.append(name)
|
|
return out_layernames
|
|
|
|
|
|
class FeatVGG(nn.Module):
|
|
def __init__(self, content_layers=["relu3_1"]):
|
|
super(FeatVGG, self).__init__()
|
|
self.content_layers = conv_to_relu(content_layers)
|
|
self.vgg19 = vgg_util.get_renamed_vgg()
|
|
self.last_c_layer = self.content_layers[-1]
|
|
is_last_content = False
|
|
replace_layers, del_layers = [], []
|
|
for name, mod in self.vgg19.named_children():
|
|
if is_last_content:
|
|
del_layers.append(name)
|
|
else:
|
|
if name == self.last_c_layer:
|
|
is_last_content = True
|
|
|
|
for name in del_layers:
|
|
delattr(self.vgg19, name)
|
|
|
|
# no need for gradweight vgg19
|
|
for param in self.vgg19.parameters():
|
|
param.requires_grad = False
|
|
|
|
def forward(self, input_img):
|
|
# input image is BGR image
|
|
# each channel ranges in [0,255]
|
|
# should be normalized with mean = [0.406*255, 0.456*255, 0.485*255] = [103,116,123]
|
|
# out = {}
|
|
return self.vgg19(input_img)
|
|
|
|
|
|
class VGGNet_multilayer(nn.Module):
|
|
def __init__(self):
|
|
"""Select conv1_1 ~ conv5_1 activation maps."""
|
|
super(VGGNet_multilayer, self).__init__()
|
|
self.select = ["0", "5", "10", "19", "28"]
|
|
self.vgg = models.vgg19(pretrained=True).features
|
|
|
|
def forward(self, x):
|
|
"""Extract multiple convolutional feature maps.
|
|
x: rgb image
|
|
ranges in [0,1]
|
|
should be normalzied with mean = [0.485, 0.456, 0.406]
|
|
and variance = [0.229, 0.224, 0.225]
|
|
"""
|
|
features = []
|
|
for name, layer in self.vgg._modules.items():
|
|
x = layer(x)
|
|
if name in self.select:
|
|
features.append(x)
|
|
return features
|