def vgg(**config):
dataset = config.pop('dataset', 'imagenet')
depth = config.pop('depth', 16)
bn = config.pop('bn', True)
if dataset == 'imagenet':
config.setdefault('num_classes', 1000)
if depth == 11:
if bn is False:
return vgg11(pretrained=False, **config)
else:
return vgg11_bn(pretrained=False, **config)
if depth == 13:
if bn is False:
return vgg13(pretrained=False, **config)
else:
return vgg13_bn(pretrained=False, **config)
if depth == 16:
if bn is False:
return vgg16(pretrained=False, **config)
else:
return vgg16_bn(pretrained=False, **config)
if depth == 19:
if bn is False:
return vgg19(pretrained=False, **config)
else:
return vgg19_bn(pretrained=False, **config)
elif dataset == 'cifar10':
config.setdefault('num_classes', 10)
elif dataset == 'cifar100':
config.setdefault('num_classes', 100)
config.setdefault('batch_norm', bn)
return VGG(model_name[depth], **config)
def __init__(self, batchSize):
super(GeneratorLoss, self).__init__()
vgg = vgg13(pretrained=True)
loss_network = nn.Sequential(*list(vgg.features)[:25]).eval()
for param in loss_network.parameters():
param.requires_grad = False
self.loss_network = loss_network
self.mse_loss = nn.MSELoss()
self.tv_loss = TVLoss()
self.laplace = LaplacianLoss()
self.blur_kernel = get_gaussian_kernel(kernel_size=3)
self.blur_kernel2 = get_gaussian_kernel(kernel_size=5)
self.blur_kernel3 = get_gaussian_kernel(kernel_size=7)
self.adversarial_criterion = nn.BCELoss()
def vgg_13(batch_norm=True, pretrained=False, fixed_feature=True):
""" VGG 13-layer model from torchvision's vgg model.
:param batch_norm: train model with batch normalization
:param pretrained: if true, return a model pretrained on ImageNet
:param fixed_feature: if true and pretrained is true, model features are fixed while training.
"""
if batch_norm:
from torchvision.models.vgg import vgg13_bn
model = vgg13_bn(pretrained)
else:
from torchvision.models.vgg import vgg13
model = vgg13(pretrained)
ff = True if pretrained and fixed_feature else False
return _VGG(model, model.features, ff)
def recordVGG(info):
global SKIP
import torchvision.models.vgg as vggGen
if not (SKIP and 'vgg11' in info['name_list']):
INFO("proceeding for VGG11...")
net = vggGen.vgg11(pretrained=True).cuda()
sum = __summary(net, [3, 224, 224], verbose=True)
__writeInfoJSON(sum, 'vgg11')
else:
INFO("Skip VGG11")
if not (SKIP and 'vgg13' in info['name_list']):
INFO("proceeding for VGG13...")
net = vggGen.vgg13(pretrained=True).cuda()
sum = __summary(net, [3, 224, 224], verbose=True)
__writeInfoJSON(sum, 'vgg13')
else:
INFO("Skip VGG13")
if not (SKIP and 'vgg16' in info['name_list']):
INFO("proceeding for VGG16...")
net = vggGen.vgg16(pretrained=True).cuda()
sum = __summary(net, [3, 224, 224], verbose=True)
__writeInfoJSON(sum, 'vgg16')
else:
INFO("Skip VGG16")
if not (SKIP and 'vgg19' in info['name_list']):
INFO("proceeding for VGG19...")
net = vggGen.vgg19(pretrained=True).cuda()
sum = __summary(net, [3, 224, 224], verbose=True)
__writeInfoJSON(sum, 'vgg19')
else:
INFO("Skip VGG19")
if not (SKIP and 'vgg11_bn' in info['name_list']):
INFO("proceeding for VGG11_bn...")
net = vggGen.vgg11_bn(pretrained=True).cuda()
sum = __summary(net, [3, 224, 224], verbose=True)
__writeInfoJSON(sum, 'vgg11_bn')
else:
INFO("Skip VGG11_bn")
if not (SKIP and 'vgg13_bn' in info['name_list']):
INFO("proceeding for VGG13_bn...")
net = vggGen.vgg13_bn(pretrained=True).cuda()
sum = __summary(net, [3, 224, 224], verbose=True)
__writeInfoJSON(sum, 'vgg13_bn')
else:
INFO("Skip VGG13_bn")
if not (SKIP and 'vgg16_bn' in info['name_list']):
INFO("proceeding for VGG16_bn...")
net = vggGen.vgg16_bn(pretrained=True).cuda()
sum = __summary(net, [3, 224, 224], verbose=True)
__writeInfoJSON(sum, 'vgg16_bn')
else:
INFO("Skip VGG16_bn")
if not (SKIP and 'vgg19_bn' in info['name_list']):
INFO("proceeding for VGG19_bn...")
net = vggGen.vgg19_bn(pretrained=True).cuda()
sum = __summary(net, [3, 224, 224], verbose=True)
__writeInfoJSON(sum, 'vgg19_bn')
else:
INFO("Skip VGG19_bn")
def load_model(model_name, classes=1000, pretrained=True, in_channels=3):
"""Load the specified VGG architecture for ImageNet
Args:
model_name: VGG architecture type
classes: number of predicted classes
pretrained: load pretrained network on ImageNet
"""
if pretrained:
assert classes == 1000, "Pretrained models are provided only for Imagenet."
kwargs = {'num_classes': classes}
if model_name == 'vgg11':
net = VGG.vgg11(pretrained=pretrained, **kwargs)
if in_channels != 3:
input_layer = nn.Conv2d(in_channels, 64, kernel_size=3, padding=1)
nn.init.kaiming_normal_(input_layer.weight,
mode='fan_out',
nonlinearity='relu')
input_layer.bias.data.zero_()
net.features[0] = input_layer
elif model_name == 'vgg13':
net = VGG.vgg13(pretrained=pretrained, **kwargs)
if in_channels != 3:
input_layer = nn.Conv2d(in_channels, 64, kernel_size=3, padding=1)
nn.init.kaiming_normal_(input_layer.weight,
mode='fan_out',
nonlinearity='relu')
input_layer.bias.data.zero_()
net.features[0] = input_layer
elif model_name == 'vgg16':
net = VGG.vgg16(pretrained=pretrained, **kwargs)
if in_channels != 3:
input_layer = nn.Conv2d(in_channels, 64, kernel_size=3, padding=1)
nn.init.kaiming_normal_(input_layer.weight,
mode='fan_out',
nonlinearity='relu')
input_layer.bias.data.zero_()
net.features[0] = input_layer
elif model_name == 'vgg19':
net = VGG.vgg19(pretrained=pretrained, **kwargs)
if in_channels != 3:
input_layer = nn.Conv2d(in_channels, 64, kernel_size=3, padding=1)
nn.init.kaiming_normal_(input_layer.weight,
mode='fan_out',
nonlinearity='relu')
input_layer.bias.data.zero_()
net.features[0] = input_layer
elif model_name == 'vgg11bn':
net = VGG.vgg11_bn(pretrained=pretrained, **kwargs)
if in_channels != 3:
input_layer = nn.Conv2d(in_channels, 64, kernel_size=3, padding=1)
nn.init.kaiming_normal_(input_layer.weight,
mode='fan_out',
nonlinearity='relu')
input_layer.bias.data.zero_()
net.features[0] = input_layer
elif model_name == 'vgg13bn':
net = VGG.vgg13_bn(pretrained=pretrained, **kwargs)
if in_channels != 3:
input_layer = nn.Conv2d(in_channels, 64, kernel_size=3, padding=1)
nn.init.kaiming_normal_(input_layer.weight,
mode='fan_out',
nonlinearity='relu')
input_layer.bias.data.zero_()
net.features[0] = input_layer
elif model_name == 'vgg16bn':
net = VGG.vgg16_bn(pretrained=pretrained, **kwargs)
if in_channels != 3:
input_layer = nn.Conv2d(in_channels, 64, kernel_size=3, padding=1)
nn.init.kaiming_normal_(input_layer.weight,
mode='fan_out',
nonlinearity='relu')
input_layer.bias.data.zero_()
net.features[0] = input_layer
elif model_name == 'vgg19bn':
net = VGG.vgg19_bn(pretrained=pretrained, **kwargs)
if in_channels != 3:
input_layer = nn.Conv2d(in_channels, 64, kernel_size=3, padding=1)
nn.init.kaiming_normal_(input_layer.weight,
mode='fan_out',
nonlinearity='relu')
input_layer.bias.data.zero_()
net.features[0] = input_layer
elif model_name == 'vgg19_orig':
net = VGG.vgg19(pretrained=False, **kwargs)
if in_channels != 3:
input_layer = nn.Conv2d(in_channels, 64, kernel_size=3, padding=1)
net.features[0] = input_layer
init_weights_vgg_orig(net)
elif model_name == 'alexnet':
net = AlexNet(pretrained=pretrained, **kwargs)
if in_channels != 3:
input_layer = nn.Conv2d(in_channels,
64,
kernel_size=11,
stride=4,
padding=2)
nn.init.kaiming_normal_(input_layer.weight,
mode='fan_out',
nonlinearity='relu')
input_layer.bias.data.zero_()
net.features[0] = input_layer
elif model_name == 'lenet':
kwargs['in_channels'] = in_channels
net = lenet(**kwargs)
else:
raise ValueError("Unsupported model architecture.")
return net