def mobilenet_v2(pretrained=False, checkpoints=None, 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(**kwargs)
if pretrained:
if checkpoints is not None:
model.load_state_dict(torch.load(checkpoints), True)
return model
state_dict = load_state_dict_from_url(model_urls['mobilenet_v2'],
progress=progress)
model.load_state_dict(state_dict)
return model
return model
def _vgg(arch,
cfg,
batch_norm,
pretrained=False,
checkpoints=None,
progress=True,
**kwargs):
if pretrained:
kwargs['init_weights'] = False
model = VGG(make_layers(cfgs[cfg], batch_norm=batch_norm), **kwargs)
if pretrained:
if checkpoints is not None:
model.load_state_dict(torch.load(checkpoints), True)
return model
state_dict = load_state_dict_from_url(model_urls[arch],
progress=progress)
model.load_state_dict(state_dict)
return model
return model
def vgg(depth, batch_norm, num_classes, pretrained):
model = VGG(make_layers(cfgs[depth], batch_norm=batch_norm),
num_classes,
init_weights=True)
arch = 'vgg' + str(depth)
if batch_norm == True: arch += '_bn'
if pretrained and (num_classes == 1000) and (arch
in pretrained_model_urls):
state_dict = load_state_dict_from_url(pretrained_model_urls[arch],
progress=True)
model.load_state_dict(state_dict)
elif pretrained:
raise ValueError(
'No pretrained model in vggnet {} model with class number {}'.
format(depth, num_classes))
return model
def dpn131(pretrained=False, test_time_pool=False, **kwargs):
"""Constructs a DPN-131 model.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet-1K
test_time_pool (bool): If True, pools features for input resolution beyond
standard 224x224 input with avg+max at inference/validation time
**kwargs : Keyword args passed to model __init__
num_classes (int): Number of classes for classifier linear layer, default=1000
"""
model = DPN(num_init_features=128,
k_r=160,
groups=40,
k_sec=(4, 8, 28, 3),
inc_sec=(16, 32, 32, 128),
test_time_pool=test_time_pool,
**kwargs)
if pretrained:
model.load_state_dict(load_state_dict_from_url(model_urls['dpn131']))
return model
def _build_vgg_loss(self, avg_pool, feature_norm, weights, device):
self.content_losses, self.style_losses = {}, {}
self.vgg_loss = nn.Sequential()
vgg = models.vgg19(pretrained=False).features
if weights in ('original', 'normalized'):
state_dict = load_state_dict_from_url(
'https://storage.googleapis'
f'.com/prism-weights/vgg19-{weights}.pth')
else:
state_dict = torch.load(weights)
vgg.load_state_dict(state_dict)
vgg = vgg.eval()
for param in vgg.parameters():
param.requires_grad_(False)
i_pool, i_conv = 1, 0
for layer in vgg.children():
if isinstance(layer, nn.Conv2d):
i_conv += 1
name = f'conv_{i_pool}_{i_conv}'
elif isinstance(layer, nn.ReLU):
name = f'relu_{i_pool}_{i_conv}'
layer = nn.ReLU(inplace=False)
elif isinstance(layer, nn.MaxPool2d):
name = f'pool_{i_pool}'
if avg_pool:
layer = nn.AvgPool2d(kernel_size=2, stride=2, padding=0)
i_pool += 1
i_conv = 0
self.vgg_loss.add_module(name, layer)
if name in self.content_weights:
content_loss = ContentLoss('none')
self.vgg_loss.add_module(f'content_loss_{i_pool}_{i_conv}',
content_loss)
self.content_losses[name] = content_loss
if name in self.style_weights:
style_loss = StyleLoss('none', feature_norm)
self.vgg_loss.add_module(f'style_loss_{i_pool}_{i_conv}',
style_loss)
self.style_losses[name] = style_loss
if (len(self.style_weights) == len(self.style_losses)
and len(self.content_weights) == len(self.content_losses)):
break
self.vgg_loss.to(device)
def mobilenet_v2(num_classes,
in_channels=3,
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 pretrained and in_channels != 3:
# raise ValueError('ImageNet pretrained models only support 3 input channels, but got {}'.format(in_channels))
if pretrained:
model = MobileNetV2(**kwargs)
state_dict = load_state_dict_from_url(model_urls['mobilenet_v2'],
progress=progress)
model.load_state_dict(state_dict)
conv0 = model.features[0][0]
model.features[0][0] = nn.Conv2d(in_channels=in_channels,
out_channels=conv0.out_channels,
kernel_size=conv0.kernel_size,
stride=conv0.stride,
padding=conv0.padding,
bias=conv0.bias)
if in_channels <= 3:
model.features[0][
0].weight.data = conv0.weight[:, 0:in_channels, :, :]
else:
multi = in_channels // 3
last = in_channels % 3
model.features[0][0].weight.data = torch.cat(
[conv0.weight for x in range(multi)], dim=1)
model.features[0][0].weight.data = conv0.weight[:, :last, :, :]
model.classifier[1] = nn.Linear(model.classifier[1].in_features,
num_classes)
else:
model = MobileNetV2(num_classes=num_classes,
in_channels=in_channels,
**kwargs)
return model
def AlexNetDANN(pretrained=False, progress=True, **kwargs):
r"""AlexNet model architecture from the
`"One weird trick..." <https://arxiv.org/abs/1404.5997>`_ paper.
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 = DANN(**kwargs)
if pretrained:
state_dict = load_state_dict_from_url(model_urls['alexnet'],
progress=progress)
# removing unused params
state_dict.popitem("classifier.6.bias")
state_dict.popitem("classifier.6.weight")
model.load_state_dict(state_dict,strict=False)
model.update_weigth()
return model
def __init__(self, num_classes=1000, pretrained=False, **kwargs):
super().__init__(block=models.resnet.BasicBlock,
layers=[2, 2, 2, 2],
num_classes=num_classes,
**kwargs)
if pretrained:
state_dict = load_state_dict_from_url(
models.resnet.model_urls["resnet18"], progress=True)
self.load_state_dict(state_dict)
self.avgpool = nn.AvgPool2d((7, 7))
self.last_conv = torch.nn.Conv2d(in_channels=self.fc.in_features,
out_channels=num_classes,
kernel_size=1)
self.last_conv.weight.data.copy_(
self.fc.weight.data.view(*self.fc.weight.data.shape, 1, 1))
self.last_conv.bias.data.copy_(self.fc.bias.data)
def _darknet(arch: str, pretrained: bool, progress: bool, *args: Any,
**kwargs: Any) -> DarkNet:
"""
Constructs a DarkNet architecture from
# TODO
"""
model = DarkNet(*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)
model.load_state_dict(state_dict)
return model
def _get_model_by_name(model_name, classes=1000, pretrained=False):
block_args_list, global_params = get_efficientnet_params(model_name, override_params={'num_classes': classes})
model = EfficientNet(block_args_list, global_params)
try:
if pretrained:
pretrained_state_dict = load_state_dict_from_url(IMAGENET_WEIGHTS[model_name])
if classes != 1000:
random_state_dict = model.state_dict()
pretrained_state_dict['_fc.weight'] = random_state_dict['_fc.weight']
pretrained_state_dict['_fc.bias'] = random_state_dict['_fc.bias']
model.load_state_dict(pretrained_state_dict)
except KeyError as e:
print(f"NOTE: Currently model {e} doesn't have pretrained weights, therefore a model with randomly initialized"
" weights is returned.")
return model
def ResNextFPN50(is_pretrained = True, use_se=False):
fpn = ResFPN(Bottleneck, [3, 4, 6, 3], groups=32, width_per_group=4, dilation=1, use_se=use_se)
if is_pretrained is False:
for m in fpn.modules():
if isinstance(m, nn.Conv2d):
nn.init.normal_(m.weight, mean=0, std=0.01)
if m.bias is not None:
nn.init.constant_(m.bias, 0)
else:
print('Loading pretrained ResNeXt50 model with ImageNet..')
state_dict = load_state_dict_from_url(model_urls['resnext50_32x4d'], progress=True)
del state_dict['fc.weight']
del state_dict['fc.bias']
missing_keys = fpn.load_state_dict(state_dict, strict=False)
print(missing_keys)
return fpn
def _efficientdet(arch, pretrained=None, **kwargs):
cfgs = deepcopy(CFGS)
cfg_settings = cfgs[arch]["default"]
cfg_params = cfg_settings.pop("params")
kwargs.update(cfg_params)
model = EfficientDet(**kwargs)
if pretrained:
state_dict = load_state_dict_from_url(cfgs[arch][pretrained]["url"])
kwargs_cls = kwargs.get("num_classes", None)
if kwargs_cls and kwargs_cls != cfg_settings["num_classes"]:
logging.warning(
f"Using model pretrained for {cfg_settings['num_classes']} classes with {kwargs_cls} classes. Last layer is initialized randomly"
)
state_dict["cls_head_conv.1.weight"] = model.state_dict()[f"cls_head_conv.1.weight"]
state_dict["cls_head_conv.1.bias"] = model.state_dict()["cls_head_conv.1.bias"]
# strict=False to avoid error on extra bias in BiFPN
model.load_state_dict(state_dict, strict=False)
setattr(model, "pretrained_settings", cfg_settings)
return model
def _load_state_dict(model, model_url, progress):
# '.'s are no longer allowed in module names, but previous _DenseLayer
# has keys 'norm.1', 'relu.1', 'conv.1', 'norm.2', 'relu.2', 'conv.2'.
# They are also in the checkpoints in model_urls. This pattern is used
# to find such keys.
pattern = re.compile(
r'^(.*denselayer\d+\.(?:norm|relu|conv))\.((?:[12])\.(?:weight|bias|running_mean|running_var))$'
)
state_dict = load_state_dict_from_url(model_url, progress=progress)
# print(state_dict)
# exit(-1)
for key in list(state_dict.keys()):
res = pattern.match(key)
if res:
new_key = res.group(1) + res.group(2)
state_dict[new_key] = state_dict[key]
del state_dict[key]
model.load_state_dict(state_dict)
def inception_v3(pretrained=False, checkpoints=None, 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
kwargs[
'init_weights'] = False # we are loading weights from a pretrained model
model = Inception3(**kwargs)
if checkpoints is not None:
model.load_state_dict(torch.load(checkpoints), True)
if not original_aux_logits:
model.aux_logits = False
del model.AuxLogits
return model
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 load_pretrained_weights(self):
vgg16_weights = hub.load_state_dict_from_url(
"https://download.pytorch.org/models/vgg16_bn-6c64b313.pth")
count_vgg = 0
count_this = 0
vggkeys = list(vgg16_weights.keys())
thiskeys = list(self.state_dict().keys())
corresp_map = []
while (True):
vggkey = vggkeys[count_vgg]
thiskey = thiskeys[count_this]
if "classifier" in vggkey:
break
while vggkey.split(".")[-1] not in thiskey:
count_this += 1
thiskey = thiskeys[count_this]
corresp_map.append([vggkey, thiskey])
count_vgg += 1
count_this += 1
mapped_weights = self.state_dict()
for k_vgg, k_segnet in corresp_map:
if (self.in_channels !=
3) and "features" in k_vgg and "conv1_1." not in k_segnet:
mapped_weights[k_segnet] = vgg16_weights[k_vgg]
elif (self.in_channels == 3) and "features" in k_vgg:
mapped_weights[k_segnet] = vgg16_weights[k_vgg]
try:
self.load_state_dict(mapped_weights)
print("Loaded VGG-16 weights in Segnet !")
except:
print("Error VGG-16 weights in Segnet !")
raise
def fetch_model(name, map_location=None, **kwargs):
"""Fetch model from URL.
Loads model or state dict from URL.
Args:
name: Model name hosted on `celldetection.org` or url. Urls must start with 'http'.
map_location: A function, `torch.device`, string or a dict specifying how to remap storage locations.
**kwargs: From the doc of `torch.models.utils.load_state_dict_from_url`.
"""
url = name if name.startswith(
'http') else f'https://celldetection.org/torch/models/{name}.pt'
m = load_state_dict_from_url(url, map_location=map_location, **kwargs)
if isinstance(m, dict) and 'cd.models' in m.keys():
from .. import models
conf = m['cd.models']
m = getattr(models, conf['model'])(*conf['a'], **conf['kw'])
m.load_state_dict(conf['state_dict'])
return m
def create(arch, image_size, channels, pretrained, progress):
""" Creates a specified GAN model
Args:
arch (str): Arch name of model.
image_size (int): Number of image size.
channels (int): Number of input channels.
pretrained (bool): Load pretrained weights into the model.
progress (bool): Show progress bar when downloading weights.
Returns:
PyTorch model.
"""
model = Generator(image_size, channels)
if pretrained:
state_dict = load_state_dict_from_url(model_urls[arch],
progress=progress,
map_location=torch.device("cpu"))
model.load_state_dict(state_dict)
return model
def _efficient_x3d(
pretrained: bool = False,
progress: bool = True,
checkpoint_path: str = None,
# Model params
expansion: str = "XS",
**kwargs: Any,
) -> nn.Module:
model = create_x3d(
expansion=expansion,
**kwargs,
)
if pretrained and checkpoint_path is not None:
state_dict = load_state_dict_from_url(checkpoint_path,
progress=progress)
model.load_state_dict(state_dict, strict=True)
return model
def _resnext(arch, block, layers, pretrained, progress, **kwargs):
model = ResNet(block, layers, **kwargs)
state_dict = load_state_dict_from_url(model_urls[arch], progress=progress)
# new_state_dict = {}
new_state_dict = model.state_dict()
print("11111", len(state_dict))
for k, v in state_dict.items():
# print(k)
if k in new_state_dict.keys() and (
not k.startswith('fc')
or not k.startswith('layer4')): # 不加载全连接层
new_state_dict[k] = v
# # print("{}: {}".format(k,v.requires_grad))
# # elif k.startswith('fc') or k.startswith('layer4'):
# # print("yes, ",k)
# new_state_dict = model.state_dict()
# new_state_dict.update(state_dict)
print("22222", len(new_state_dict))
model.load_state_dict(new_state_dict)
return model
def resnext(depth, num_classes, pretrained):
model = _resnext(mode=cfgs[depth][0],
block=cfgs[depth][1],
cardinality=cfgs[depth][2],
layers=cfgs[depth][3],
num_classes=num_classes)
arch = 'resnext' + str(depth)
if pretrained and (num_classes == 1000) and (arch
in pretrained_model_urls):
state_dict = load_state_dict_from_url(pretrained_model_urls[arch],
progress=True)
model.load_state_dict(state_dict)
elif pretrained:
raise ValueError(
'No pretrained model in resnext {} model with class number {}'.
format(depth, num_classes))
return model
def alexnetDANN(pretrained=False, progress=True, **kwargs):
"""
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 = DANNModel(num_classes = 1000, **kwargs)
if pretrained:
state_dict = load_state_dict_from_url(model_urls['alexnet'],
progress=progress)
model.load_state_dict(state_dict, strict = False)
#we copy the pretrained weights of the classifier to the domain classifier
model.discriminator[1].weight.data = model.classifier[1].weight.data.clone()
model.discriminator[1].bias.data = model.classifier[1].bias.data.clone()
model.discriminator[4].weight.data = model.classifier[4].weight.data.clone()
model.discriminator[4].bias.data = model.classifier[4].bias.data.clone()
return model
def dann_net(pretrained=False, progress=True, **kwargs):
r"""AlexNet model architecture from the
`"One weird trick..." <https://arxiv.org/abs/1404.5997>`_ paper.
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
"""
net = DannNet(**kwargs)
if pretrained:
state_dict = load_state_dict_from_url(model_urls['alexnet'],
progress=progress)
net.load_state_dict(state_dict, strict=False)
net.dann_classifier[1].weight.data = net.classifier[
1].weight.data.clone()
net.dann_classifier[1].bias.data = net.classifier[1].bias.data.clone()
net.dann_classifier[4].weight.data = net.classifier[
4].weight.data.clone()
net.dann_classifier[4].bias.data = net.classifier[4].bias.data.clone()
return net
def _vgg(architecture, config, batch_norm, pretrained, progress, **kwargs):
if pretrained:
kwargs['init_weights'] = False
model = VGG(make_layers(cfgs[config], batch_norm=batch_norm), **kwargs)
if pretrained:
state_dict = load_state_dict_from_url(model_urls[architecture],
progress=progress)
del state_dict["classifier.0.weight"]
del state_dict["classifier.0.bias"]
del state_dict["classifier.3.weight"]
del state_dict["classifier.3.bias"]
del state_dict["classifier.6.weight"]
del state_dict["classifier.6.bias"]
model.load_state_dict(state_dict)
return model
def _hrnet(pretrained, checkpoints, progress, **kwargs):
update_config(config, model_cfg[kwargs['model_name']])
model = get_cls_net(config)
if pretrained:
if checkpoints is not None:
model.load_state_dict(torch.load(checkpoints), True)
return model
state_dict = load_state_dict_from_url(model_cfg[kwargs['model_name'] +
'_url'])
model.load_state_dict(state_dict)
return model
for m in model.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight,
mode='fan_out',
nonlinearity='relu')
elif isinstance(m, nn.BatchNorm2d):
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)
return model
def _resnet(arch,
block,
layers,
pretrained,
progress,
use_cbam=False,
use_mixpool=False,
**kwargs):
model = ResNet(block,
layers,
use_cbam=use_cbam,
use_mixpool=use_mixpool,
**kwargs)
if pretrained and use_cbam:
state_dict = load_state_dict_from_url(model_urls[arch],
progress=progress)
new_state_dict = model.state_dict()
new_state_dict.update(state_dict)
model.load_state_dict(new_state_dict)
return model
def _resnet(
arch: str,
block: Type[Union[BasicBlock, Bottleneck]],
layers: List[int],
pretrained: bool,
progress: bool,
custom_class_num: bool = False,
**kwargs: Any
) -> ResNet:
model = ResNet(block, layers, custom_class_num=custom_class_num, **kwargs)
if pretrained:
state_dict = load_state_dict_from_url(model_urls[arch],
progress=progress)
model.load_state_dict(state_dict)
if custom_class_num:
return ResNetFlex(model)
return model
def mobilenet_v2(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(**kwargs)
if pretrained:
org_state_dict = load_state_dict_from_url(model_urls['mobilenet_v2'],
progress=progress)
state_dict = {}
for k1, k2 in zip(org_state_dict.keys(), model.state_dict().keys()):
state_dict[k2] = org_state_dict[k1]
model.load_state_dict(state_dict)
return model
def mobilenet_v2(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(**kwargs)
if pretrained:
try:
from torch.hub import load_state_dict_from_url
except ImportError:
from torch.utils.model_zoo import load_url as load_state_dict_from_url
state_dict = load_state_dict_from_url(
'https://download.pytorch.org/models/mobilenet_v2-b0353104.pth',
progress=True)
model.load_state_dict(state_dict)
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)
if pretrained:
state_dict = load_state_dict_from_url(model_urls[arch],
progress=progress)
snn_state_dict = {}
# compared to the ANN parameters, the modules are lifted
# we modify the state dict accordingly here
for key in state_dict:
l = key.split(".")
l.insert(-1, "lifted_module")
new_key = ".".join(l)
snn_state_dict[new_key] = state_dict[key]
model.load_state_dict(snn_state_dict)
return model
def _cvt(arch, pretrained, progress,
num_layers, num_heads, mlp_ratio, embedding_dim,
kernel_size=4, positional_embedding='learnable',
*args, **kwargs):
model = CVT(num_layers=num_layers,
num_heads=num_heads,
mlp_ratio=mlp_ratio,
embedding_dim=embedding_dim,
kernel_size=kernel_size,
*args, **kwargs)
if pretrained and arch in model_urls:
state_dict = load_state_dict_from_url(model_urls[arch],
progress=progress)
if positional_embedding == 'learnable':
state_dict = pe_check(model, state_dict)
elif positional_embedding == 'sine':
state_dict['classifier.positional_emb'] = model.state_dict()['classifier.positional_emb']
model.load_state_dict(state_dict)
return model
