def _resnet(arch, block, layers, pretrained, progress, **kwargs):
model = ResNet(block, layers, **kwargs)
if pretrained:
state_dict = load_state_dict_from_url(model_urls[arch],
progress=progress)
model.load_state_dict(state_dict, strict=False)
url = 'http://www.cs.cornell.edu/projects/megadepth/dataset/models/best_generalization_net_G.pth'
DE_weights = load_state_dict_from_url(url, progress=progress)
model.DE.load_state_dict(DE_weights)
return model
def _load_pretrained(model_name, model, progress):
if model_name not in _MODEL_URLS or _MODEL_URLS[model_name] is None:
raise ValueError(
"No checkpoint is available for model type {}".format(model_name))
checkpoint_url = _MODEL_URLS[model_name]
model.load_state_dict(
load_state_dict_from_url(checkpoint_url, progress=progress))
def alexnet(pretrained=False, progress=True, **kwargs):
model = AlexNet(**kwargs)
if pretrained:
state_dict = load_state_dict_from_url(model_urls['alexnet'],
progress=progress)
model.load_state_dict(state_dict)
return model
def googlenet(pretrained=False, progress=True, **kwargs):
r"""GoogLeNet (Inception V1) model architecture from
going deeper with convolutional
Args:
pretrained (bool): If True,returns a model pretrained on ImageNet
progress (bool): If True, displays a progress bar of the download tp stderr
"""
if pretrained:
if "transform_input" not in kwargs:
kwargs["transform_input"] = True
if "aux_logits" not in kwargs:
kwargs["aux_logits"] = False
if kwargs["aux_logits"]:
warnings.warn(
"auxiliary heads in the pretrained googlenet model are Not pretrained, So make sure to train them."
)
original_aux_logits = kwargs["aux_logits"]
kwargs["aux_logits"] = True
kwargs["init_weights"] = False
model = GoogLeNet(**kwargs)
state_dict = load_state_dict_from_url(model_urls["googlenet"],
progress=progress)
model.load_state_dict(state_dict)
if not original_aux_logits:
model.aux_logits = False
del model.aux1, model.aux2
return model
return GoogLeNet(**kwargs)
def mobilenet_v2(output_layers=None,
pretrained=False,
progress=True,
**kwargs):
"""
Constructs a MobileNetV2 architecture from
`"MobileNetV2: Inverted Residuals and Linear Bottlenecks" <https://arxiv.org/abs/1801.04381>`_.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
if output_layers is None:
output_layers = ['default']
else:
for l in output_layers:
if l not in ['conv1', 'layer1', 'layer2', 'layer3', 'layer4']:
raise ValueError('Unknown layer: {}'.format(l))
model = MobileNetV2(output_layers=output_layers, **kwargs)
if pretrained:
state_dict = load_state_dict_from_url(model_urls['mobilenet_v2'],
progress=progress)
new_dict = dict()
for i, (key, val) in enumerate(state_dict.items()):
new_key = list(model.state_dict().keys())[i]
new_dict[new_key] = val
model.load_state_dict(new_dict, strict=True)
return model
def _resnet(arch,
block,
layers,
pretrained,
progress,
with_ibn,
replace_stride_with_dilation=[False, False, False],
freeze_bn=False,
**kwargs):
model = ResNet(block,
layers,
with_ibn=with_ibn,
replace_stride_with_dilation=replace_stride_with_dilation,
freeze_bn=freeze_bn,
**kwargs)
if pretrained:
state_dict = load_state_dict_from_url(model_urls[arch],
progress=progress)
model_dict = model.state_dict()
pretrained_dict = {
k: v
for k, v in state_dict.items() if k in model_dict
}
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
return model
def _shufflenetv2(arch, pretrained, progress, *args, **kwargs):
model = ShuffleNetV2(*args, **kwargs)
if pretrained:
model_url = model_urls[arch]
if model_url is None:
raise NotImplementedError(
'pretrained {} is not supported as of now'.format(arch))
else:
state_dict = load_state_dict_from_url(model_url, progress=progress)
_pretrained_dict = OrderedDict()
for idx, (k, v) in enumerate(state_dict.items()):
splitted_k = k.split('.')
# special for 1.0x
if 29 < idx < 280:
splitted_k[-2] = str(int(splitted_k[-2]) + 1)
_pretrained_dict['.'.join(splitted_k)] = v
else:
_pretrained_dict[k] = v
model.load_state_dict(_pretrained_dict)
# release
del _pretrained_dict
del state_dict
return model
def vgg(cfg, progress, **kwargs):
pretrained = kwargs['pretrained']
batch_norm = kwargs['batch_norm']
model = VGG(
make_layers(cfgs[cfg], conv_kernel_size=3, batch_norm=batch_norm),
**kwargs)
if pretrained:
if batch_norm:
model_url = 'https://download.pytorch.org/models/vgg16_bn-6c64b313.pth'
else:
model_url = 'https://download.pytorch.org/models/vgg16-397923af.pth'
state_dict = load_state_dict_from_url(model_url, progress=progress)
state_dict['features.0.weight'] = state_dict[
'features.0.weight'][:, 0, :, :].unsqueeze(1)
state_dict = {
param_name: param
for param_name, param in state_dict.items()
if not ('classifier' in param_name)
}
model.load_state_dict(state_dict, strict=False)
return model
def _resnet(arch, block, layers, pretrained, progress, **kwargs):
model = ResNet(block, layers, **kwargs)
if pretrained:
state_dict = load_state_dict_from_url(model_urls[arch],
progress=progress)
model.load_state_dict(state_dict)
return model
def inception_v3(pretrained=False, progress=True, **kwargs):
r"""Inception v3 model architecture from
`"Rethinking the Inception Architecture for Computer Vision" <http://arxiv.org/abs/1512.00567>`_.
.. note::
**Important**: In contrast to the other models the inception_v3 expects tensors with a size of
N x 3 x 299 x 299, so ensure your images are sized accordingly.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
model = Inception3(**kwargs)
if pretrained:
state_dict = load_state_dict_from_url(
model_urls['inception_v3_google'], progress=progress)
helper_state_dict = state_dict.copy()
for key in helper_state_dict.keys():
if re.match(r"Aux*", key):
del state_dict[key]
# state_dict["Conv2d_1a_3x3.conv.weight"]=state_dict["Conv2d_1a_3x3.conv.weight"][:,:2,:,:]
# state_dict["fc.weight"]=state_dict["fc.weight"][:2,:]
# state_dict["fc.bias"]=state_dict["fc.bias"][:2]
model.load_state_dict(state_dict)
return model
def mbnetv2_dropout(pretrained=False, progress=True, **kwargs):
"""
Constructs a MobileNetV2 architecture from
`"MobileNetV2: Inverted Residuals and Linear Bottlenecks" <https://arxiv.org/abs/1801.04381>`_.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
"""
model = MobileNetV2(haslinear=True, **kwargs)
if pretrained:
state_dict = load_state_dict_from_url(model_urls['mobilenet_v2'],
progress=progress)
del state_dict['classifier.1.weight']
del state_dict['classifier.1.bias']
new_dict = {}
cumulative_layer = 0
for li, layers in enumerate([4, 3, 4, 6, 2]):
for l in range(layers):
key = 'features.{}.'.format(l + cumulative_layer)
new_key = 'layer{}.{}.'.format(li + 1, l)
for k in state_dict:
if key in k:
new_k = k.replace(key, new_key)
new_dict[new_k] = state_dict[k]
cumulative_layer += layers
new_params = model.state_dict()
new_params.update(new_dict)
model.load_state_dict(new_params)
return model
def _vgg(arch, cfg, batch_norm, pretrained, progress, **kwargs):
if pretrained:
kwargs['init_weights'] = False
model = VGG(make_layers(cfgs[cfg], batch_norm=batch_norm), **kwargs)
for name, param in model.named_parameters():
param.requires_grad = True
if pretrained:
state_dict = load_state_dict_from_url(model_urls[arch],
progress=progress)
# print(state_dict.keys())
new_state_dict = dict()
for name, param in state_dict.items():
if name in [
"classifier.0.weight", "classifier.0.bias",
"classifier.3.weight", "classifier.3.bias",
"classifier.6.weight", "classifier.6.bias"
]:
continue
else:
new_state_dict[name] = param
# print(new_state_dict.keys())
# state_dict = torch.load('/network/tmp1/bakhtias/PreTrainedModels/vgg16-82412952.pth')
model.load_state_dict(new_state_dict)
return model
def vgg16(in_channels, pretrained, **kwargs):
'''
Function to load vgg16 features model
'''
if pretrained:
kwargs['init_weights'] = False
model = VGG(in_channels, cfg, **kwargs)
#download model from url
if pretrained:
if kwargs['batch_norm']:
url = model_urls['vgg16_bn']
else:
url = model_urls['vgg16']
state_dict = load_state_dict_from_url(url, progress=True)
state_dict = remove_classifier_keys(state_dict)
if in_channels == 1:
state_dict = change_state_dict_shape(
state_dict, layer_name="features.0.weight")
model.load_state_dict(state_dict)
return model
def _resnet(arch,
block,
layers,
pretrained,
progress,
layer_4,
fixed,
tanh=False):
model = ResNet(block, layers, tanh=tanh)
if pretrained:
model_dict = load_state_dict_from_url(model_urls[arch],
progress=progress)
# model_dict = model_zoo.load_url("https://download.pytorch.org/models/resnet18-5c106cde.pth")
model_dict.pop('fc.weight', None)
model_dict.pop('fc.bias', None)
model.load_state_dict(load_my_state_dict(model, model_dict))
if layer_4:
for name, child in model.named_children():
if name != 'layer4' and name != 'fc':
for param in child.parameters():
param.requires_grad = False
if fixed:
for name, child in model.named_children():
if name != 'layer4' and name != 'fc':
for param in child.parameters():
param.requires_grad = False
# model.eval()
# model.load_state_dict(state_dict)
return model
def resnet152(pretrained=False, model_root=None, **kwargs):
model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs)
if pretrained:
state_dict = load_state_dict_from_url(model_urls['resnet152'],
model_root)
model.load_state_dict(state_dict)
return model
def resnet34(pretrained=False, model_root=None, **kwargs):
model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)
if pretrained:
state_dict = load_state_dict_from_url(model_urls['resnet34'],
model_root)
model.load_state_dict(state_dict)
return model
def _resnet(arch, block, layers, pretrained, progress, fixed, tanh=False):
model = ResNet(block, layers)
if pretrained:
model_dict = load_state_dict_from_url(model_urls[arch],
progress=progress)
def remove_layer4(model_dict):
layer4_params = []
for k, _ in model_dict.items():
if k.startswith('layer4'):
layer4_params.append(k)
for p in layer4_params:
model_dict.pop(p, None)
return model_dict
model_dict = remove_layer4(model_dict)
model_dict.pop('fc.weight', None)
model_dict.pop('fc.bias', None)
model.load_state_dict(model_dict)
if fixed:
for name, child in model.named_children():
for param in child.parameters():
param.requires_grad = False
return model
def fid_inception_v3():
"""Build pretrained Inception model for FID computation
The Inception model for FID computation uses a different set of weights
and has a slightly different structure than torchvision's Inception.
This method first constructs torchvision's Inception and then patches the
necessary parts that are different in the FID Inception model.
"""
inception = _inception_v3(num_classes=1008,
aux_logits=False,
pretrained=False)
inception.Mixed_5b = FIDInceptionA(192, pool_features=32)
inception.Mixed_5c = FIDInceptionA(256, pool_features=64)
inception.Mixed_5d = FIDInceptionA(288, pool_features=64)
inception.Mixed_6b = FIDInceptionC(768, channels_7x7=128)
inception.Mixed_6c = FIDInceptionC(768, channels_7x7=160)
inception.Mixed_6d = FIDInceptionC(768, channels_7x7=160)
inception.Mixed_6e = FIDInceptionC(768, channels_7x7=192)
inception.Mixed_7b = FIDInceptionE_1(1280)
inception.Mixed_7c = FIDInceptionE_2(2048)
state_dict = load_state_dict_from_url(FID_WEIGHTS_URL, progress=True)
inception.load_state_dict(state_dict)
return inception
def googlenet(pretrained=False, progress=True, **kwargs):
r"""GoogLeNet (Inception v1) model architecture from
`"Going Deeper with Convolutions" <http://arxiv.org/abs/1409.4842>`_.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
aux_logits (bool): If True, adds two auxiliary branches that can improve training.
Default: *False* when pretrained is True otherwise *True*
transform_input (bool): If True, preprocesses the input according to the method with which it
was trained on ImageNet. Default: *False*
"""
if pretrained:
if 'transform_input' not in kwargs:
kwargs['transform_input'] = True
if 'aux_logits' not in kwargs:
kwargs['aux_logits'] = False
if kwargs['aux_logits']:
warnings.warn(
'auxiliary heads in the pretrained googlenet model are NOT pretrained, '
'so make sure to train them')
original_aux_logits = kwargs['aux_logits']
kwargs['aux_logits'] = True
kwargs['init_weights'] = False
model = GoogLeNet(**kwargs)
state_dict = load_state_dict_from_url(
googlenet_model_urls['googlenet'], progress=progress)
model.load_state_dict(state_dict)
if not original_aux_logits:
model.aux_logits = False
del model.aux1, model.aux2
return model
return GoogLeNet(**kwargs)
def load_pretrained_network(self, dataset_numclasses):
print('Loading pretrained model from {}..', model_url)
state_dict = load_state_dict_from_url(model_url, progress=True)
self.load_state_dict(state_dict)
if self.num_classes != dataset_numclasses:
self.num_classes = dataset_numclasses
self.last_linear = nn.Linear(self.last_linear.in_features, dataset_numclasses)
def inception_v3(pretrained=False, progress=True, **kwargs):
r"""Inception v3 model architecture from
`"Rethinking the Inception Architecture for Computer Vision" <http://arxiv.org/abs/1512.00567>`_.
.. note::
**Important**: In contrast to the other models the inception_v3 expects tensors with a size of
N x 3 x 299 x 299, so ensure your images are sized accordingly.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
progress (bool): If True, displays a progress bar of the download to stderr
aux_logits (bool): If True, add an auxiliary branch that can improve training.
Default: *True*
transform_input (bool): If True, preprocesses the input according to the method with which it
was trained on ImageNet. Default: *False*
"""
if pretrained:
if 'transform_input' not in kwargs:
kwargs['transform_input'] = True
if 'aux_logits' in kwargs:
original_aux_logits = kwargs['aux_logits']
kwargs['aux_logits'] = True
else:
original_aux_logits = True
model = Inception3(**kwargs)
state_dict = load_state_dict_from_url(
model_urls['inception_v3_google'], progress=progress)
model.load_state_dict(state_dict)
if not original_aux_logits:
model.aux_logits = False
del model.AuxLogits
return model
return Inception3(**kwargs)
def _vgg(arch, cfg, batch_norm, pretrained, progress):
model = VGG(make_layers(cfgs[cfg], batch_norm=batch_norm)).to(device)
if pretrained:
state_dict = load_state_dict_from_url(model_urls[arch],
progress=progress)
model.load_state_dict(state_dict, strict=False)
return model
def _resnet(arch, block, layers, pretrained, progress, **kwargs):
model = ResNet(block, layers, **kwargs)
if pretrained:
state_dict = load_state_dict_from_url(model_urls[arch], progress=progress)
if not model.norm_layer_kwargs.get('track_running_stats', True):
state_dict = keyfilter(lambda k: 'running' not in k, state_dict)
model.load_state_dict(state_dict)
return model
def dla60x(**kwargs): # DLA-X-60
BottleneckX.expansion = 2
model = DLA([1, 1, 1, 2, 3, 1], [16, 32, 128, 256, 512, 1024],
block=BottleneckX,
**kwargs)
state_dict = load_state_dict_from_url(model_urls['dla60x'])
model.load_state_dict(state_dict)
return model
def dla46x_c(**kwargs): # DLA-X-46-C
BottleneckX.expansion = 2
model = DLA([1, 1, 1, 2, 2, 1], [16, 32, 64, 64, 128, 256],
block=BottleneckX,
**kwargs)
state_dict = load_state_dict_from_url(model_urls['dla46x_c'])
model.load_state_dict(state_dict)
return model
def dla34(**kwargs): # DLA-34
model = DLA([1, 1, 1, 2, 2, 1], [16, 32, 64, 128, 256, 512],
block=BasicBlock,
**kwargs)
state_dict = load_state_dict_from_url(model_urls['dla34'])
model.load_state_dict(state_dict)
return model
def _resnet(arch, block, layers, pretrained, progress, **kwargs):
model = CondResNet(block, layers, **kwargs)
if pretrained:
state_dict = load_state_dict_from_url(model_urls[arch],
progress=progress)
missing_parameter = model.load_state_dict(state_dict, strict=False)
print(missing_parameter)
return model
def _vovnet(arch, config_stage_ch, config_concat_ch, block_per_stage,
layer_per_block, pretrained, progress, **kwargs):
model = VoVNet(config_stage_ch, config_concat_ch, block_per_stage,
layer_per_block, **kwargs)
if pretrained:
state_dict = load_state_dict_from_url(model_urls[arch])
model.load_state_dict(state_dict, strict=False)
return model
def _squeezenet(version, pretrained, progress, **kwargs):
model = SqueezeNet(version, **kwargs)
if pretrained:
arch = 'squeezenet' + version
state_dict = load_state_dict_from_url(model_urls[arch],
progress=progress)
model.load_state_dict(state_dict)
return model
def _resnet(arch, block, layers, pretrained, progress, **kwargs):
model = ResNet(block, layers, **kwargs)
if pretrained:
state_dict = load_state_dict_from_url(model_urls[arch],
progress=progress)
model.load_state_dict(state_dict,
strict=False) # only load existing encoder layer
return model
