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