def get_resnext(num_layers,
cardinality=32,
bottleneck_width=4,
use_se=False,
pretrained=False,
root=os.path.expanduser('~/.torch/models'),
**kwargs):
r"""ResNext model from `"Aggregated Residual Transformations for Deep Neural Network"
<http://arxiv.org/abs/1611.05431>`_ paper.
Parameters
----------
num_layers : int
Numbers of layers. Options are 50, 101.
cardinality: int
Number of groups
bottleneck_width: int
Width of bottleneck block
pretrained : bool or str
Boolean value controls whether to load the default pretrained weights for model.
root : str, default '~/.torch/models'
Location for keeping the model parameters.
norm_layer : object
Normalization layer used (default: :class:`nn.BatchNorm`)
Can be :class:`nn.BatchNorm` or :class:`other normalization`.
norm_kwargs : dict
Additional `norm_layer` arguments, for example `num_devices=4`
"""
assert num_layers in resnext_spec, \
"Invalid number of layers: %d. Options are %s" % (
num_layers, str(resnext_spec.keys()))
layers = resnext_spec[num_layers]
net = ResNext(layers,
cardinality,
bottleneck_width,
use_se=use_se,
**kwargs)
if pretrained:
from model.model_store import get_model_file
if not use_se:
net.load_state_dict(
torch.load(
get_model_file('resnext%d_%dx%dd' %
(num_layers, cardinality, bottleneck_width),
root=root)))
else:
net.load_state_dict(
torch.load(
get_model_file('se_resnext%d_%dx%dd' %
(num_layers, cardinality, bottleneck_width),
root=root)))
from data.imagenet import ImageNetAttr
attrib = ImageNetAttr()
net.synset = attrib.synset
net.classes = attrib.classes
net.classes_long = attrib.classes_long
return net
def get_squeezenet(version,
pretrained=False,
root=os.path.join(os.path.expanduser('~'), '.torch/models'),
**kwargs):
r"""SqueezeNet model from the `"SqueezeNet: AlexNet-level accuracy with 50x fewer parameters
and <0.5MB model size" <https://arxiv.org/abs/1602.07360>`_ paper.
SqueezeNet 1.1 model from the `official SqueezeNet repo
<https://github.com/DeepScale/SqueezeNet/tree/master/SqueezeNet_v1.1>`_.
SqueezeNet 1.1 has 2.4x less computation and slightly fewer parameters
than SqueezeNet 1.0, without sacrificing accuracy.
Parameters
----------
version : str
Version of squeezenet. Options are '1.0', '1.1'.
pretrained : bool or str
Boolean value controls whether to load the default pretrained weights for model.
root : str, default '~/.torch/models'
Location for keeping the model parameters.
"""
net = SqueezeNet(version, **kwargs)
if pretrained:
import torch
from model.model_store import get_model_file
net.load_state_dict(
torch.load(get_model_file('squeezenet%s' % version, root=root)))
from data.imagenet import ImageNetAttr
attrib = ImageNetAttr()
net.synset = attrib.synset
net.classes = attrib.classes
net.classes_long = attrib.classes_long
return net
def resnet50_v1d_86(pretrained=False, root=os.path.expanduser('~/.torch/models'), **kwargs):
"""Constructs a ResNetV1d-50_1.8x model. Uses resnet50_v1d construction from resnetv1b.py
Parameters
----------
pretrained : bool or str
Boolean value controls whether to load the default pretrained weights for model.
String value represents the hashtag for a certain version of pretrained weights.
root : str, default '~/.torch/models'
Location for keeping the model parameters.
"""
model = ResNetV1b(BottleneckV1b, [3, 4, 6, 3], deep_stem=True, avg_down=True,
**kwargs)
dirname = os.path.dirname(__file__)
json_filename = os.path.join(dirname, 'resnet%d_v%dd_%.1fx' % (50, 1, 1.8) + ".json")
with open(json_filename, "r") as jsonFile:
params_items = json.load(jsonFile, object_pairs_hook=OrderedDict)
prune_torch_block.idx = 0
if pretrained:
from model.model_store import get_model_file
params_file = get_model_file('resnet%d_v%dd_%.2f' % (50, 1, 0.86), root=root)
prune_torch_block(model, list(params_items.keys()), list(params_items.values()),
params=torch.load(params_file), pretrained=True)
else:
prune_torch_block(model, list(params_items.keys()), list(params_items.values()),
params=None, pretrained=False)
if pretrained:
from data.imagenet import ImageNetAttr
attrib = ImageNetAttr()
model.synset = attrib.synset
model.classes = attrib.classes
model.classes_long = attrib.classes_long
return model
def resnet101_v1s(pretrained=False,
root=os.path.expanduser('~/.torch/models'),
**kwargs):
"""Constructs a ResNetV1s-101 model.
Parameters
----------
pretrained : bool or str
Boolean value controls whether to load the default pretrained weights for model.
root : str, default '~/.torch/models'
Location for keeping the model parameters.
dilated: bool, default False
Whether to apply dilation strategy to ResNetV1b, yielding a stride 8 model.
norm_layer : object
Normalization layer used (default: :class:`nn.BatchNorm`).
Can be :class:`nn.BatchNorm` or :class:`other normalization`.
"""
model = ResNetV1b(BottleneckV1b, [3, 4, 23, 3],
deep_stem=True,
stem_width=64,
**kwargs)
if pretrained:
import torch
from model.model_store import get_model_file
model.load_state_dict(
torch.load(get_model_file('resnet%d_v%ds' % (101, 1), root=root)))
from data.imagenet import ImageNetAttr
attrib = ImageNetAttr()
model.synset = attrib.synset
model.classes = attrib.classes
model.classes_long = attrib.classes_long
return model
def get_bisenet(dataset='pascal_paper',
backbone='resnet18',
pretrained_base=True,
pretrained=False,
root=os.path.expanduser('~/.torch/models'),
**kwargs):
acronyms = {
'pascal_voc': 'voc',
'pascal_paper': 'voc',
'pascal_aug': 'voc',
'ade20k': 'ade',
'coco': 'coco',
'citys': 'citys',
}
from data import datasets
config = bisenet_spec[backbone]
if config[2] is not None:
feat = nn.ModuleList(
_parse_network(config[1],
outputs=config[2],
pretrained=pretrained_base))
else:
feat = config[1]
model = BiseNet(config[0],
datasets[dataset].NUM_CLASS,
backbone=feat,
**kwargs)
if pretrained:
from model.model_store import get_model_file
model.load_state_dict(
torch.load(
get_model_file('bisenet_%s_%s' % (backbone, acronyms[dataset]),
root=root)))
return model
def get_vgg(num_layers,
pretrained=False,
root=os.path.join(os.path.expanduser('~'), '.torch/models'),
**kwargs):
r"""VGG model from the `"Very Deep Convolutional Networks for Large-Scale Image Recognition"
<https://arxiv.org/abs/1409.1556>`_ paper.
Parameters
----------
num_layers : int
Number of layers for the variant of densenet. Options are 11, 13, 16, 19.
pretrained : bool or str
Boolean value controls whether to load the default pretrained weights for model.
String value represents the hashtag for a certain version of pretrained weights.
root : str, default ~/.torch/models
Location for keeping the model parameters.
"""
layers, filters = vgg_spec[num_layers]
net = VGG(layers, filters, img_size=224, **kwargs)
if pretrained:
import torch
from model.model_store import get_model_file
batch_norm_suffix = '_bn' if kwargs.get('batch_norm') else ''
net.load_state_dict(
torch.load(
get_model_file('vgg%d%s' % (num_layers, batch_norm_suffix),
root=root)))
from data.imagenet import ImageNetAttr
attrib = ImageNetAttr()
net.synset = attrib.synset
net.classes = attrib.classes
net.classes_long = attrib.classes_long
return net
def inception_v3(pretrained=False,
root=os.path.expanduser('~/.torch/models'),
**kwargs):
r"""Inception v3 model from
`"Rethinking the Inception Architecture for Computer Vision"
<http://arxiv.org/abs/1512.00567>`_ paper.
Parameters
----------
pretrained : bool or str
Boolean value controls whether to load the default pretrained weights for model.
root : str, default $TORCH_HOME/models
Location for keeping the model parameters.
norm_layer : object
Normalization layer used (default: :class:`nn.BatchNorm`)
Can be :class:`nn.BatchNorm` or :class:`other normalization`.
norm_kwargs : dict
Additional `norm_layer` arguments.
"""
net = Inception3(**kwargs)
if pretrained:
import torch
from model.model_store import get_model_file
net.load_state_dict(
torch.load(get_model_file('inceptionv3', root=root)))
from data.imagenet import ImageNetAttr
attrib = ImageNetAttr()
net.synset = attrib.synset
net.classes = attrib.classes
net.classes_long = attrib.classes_long
return net
def get_cifar_resnet(version, num_layers, pretrained=False,
root=os.path.expanduser('~/.torch/models'), **kwargs):
r"""ResNet V1 model from `"Deep Residual Learning for Image Recognition"
<http://arxiv.org/abs/1512.03385>`_ paper.
ResNet V2 model from `"Identity Mappings in Deep Residual Networks"
<https://arxiv.org/abs/1603.05027>`_ paper.
Parameters
----------
version : int
Version of ResNet. Options are 1, 2.
num_layers : int
Numbers of layers. Needs to be an integer in the form of 6*n+2, e.g. 20, 56, 110, 164.
pretrained : bool, default False
Whether to load the pretrained weights for model.
root : str, default '~/.torch/models'
Location for keeping the model parameters.
norm_layer : object
Normalization layer used (default: :class:`nn.BatchNorm`)
Can be :class:`nn.BatchNorm` or :class:`other normalization`.
norm_kwargs : dict
Additional `norm_layer` arguments
"""
layers, channels = _get_resnet_spec(num_layers)
resnet_class = resnet_net_versions[version - 1]
block_class = resnet_block_versions[version - 1]
net = resnet_class(block_class, layers, channels, **kwargs)
if pretrained:
import torch
from model.model_store import get_model_file
net.load_state_dict(torch.load(get_model_file('cifar_resnet%d_v%d' % (num_layers, version), root=root),
map_location=lambda storage, loc: storage))
return net
def resnet152_v1b(pretrained=False,
root=os.path.expanduser('~/.torch/models'),
**kwargs):
"""Constructs a ResNetV1b-152 model.
Parameters
----------
pretrained : bool or str
Boolean value controls whether to load the default pretrained weights for model.
root : str, default '~/.torch/models'
Location for keeping the model parameters.
dilated: bool, default False
Whether to apply dilation strategy to ResNetV1b, yielding a stride 8 model.
norm_layer : object
Normalization layer used (default: :class:`nn.BatchNorm`)
Can be :class:`nn.BatchNorm` or :class:`other normalization`.
last_gamma : bool, default False
Whether to initialize the gamma of the last BatchNorm layer in each bottleneck to zero.
use_global_stats : bool, default False
Whether forcing BatchNorm to use global statistics instead of minibatch statistics;
optionally set to True if finetuning using ImageNet classification pretrained models.
"""
model = ResNetV1b(BottleneckV1b, [3, 8, 36, 3], **kwargs)
if pretrained:
import torch
from model.model_store import get_model_file
model.load_state_dict(
torch.load(get_model_file('resnet%d_v%db' % (152, 1), root=root)))
from data.imagenet import ImageNetAttr
attrib = ImageNetAttr()
model.synset = attrib.synset
model.classes = attrib.classes
model.classes_long = attrib.classes_long
return model
def get_mask_rcnn(name, dataset, pretrained=False,
root=os.path.expanduser('~/.torch/models'), **kwargs):
r"""Utility function to return mask rcnn networks.
Parameters
----------
name : str
Model name.
dataset : str
The name of dataset.
pretrained : bool or str
Boolean value controls whether to load the default pretrained weights for model.
String value represents the hashtag for a certain version of pretrained weights.
ctx : mxnet.Context
Context such as mx.cpu(), mx.gpu(0).
root : str
Model weights storing path.
Returns
-------
mxnet.gluon.HybridBlock
The Mask RCNN network.
"""
net = MaskRCNN(**kwargs)
if pretrained:
from model.model_store import get_model_file
full_name = '_'.join(('mask_rcnn', name, dataset))
net.load_state_dict(torch.load(get_model_file(full_name, root=root)))
return net
def get_ccnet(dataset='pascal_voc',
backbone='resnet50',
pretrained=False,
pretrained_base=True,
root=os.path.expanduser('~/.torch/models'),
**kwargs):
acronyms = {
'pascal_voc': 'voc',
'pascal_paper': 'voc',
'pascal_aug': 'voc',
'ade20k': 'ade',
'coco': 'coco',
'citys': 'citys',
}
from data import datasets
# infer number of classes
model = CCNet(datasets[dataset].NUM_CLASS,
backbone=backbone,
pretrained_base=pretrained_base,
**kwargs)
if pretrained:
from model.model_store import get_model_file
model.load_state_dict(
torch.load(
get_model_file('ccnet_%s_%s' % (backbone, acronyms[dataset]),
root=root)))
return model
def get_corner_squeeze(name, pretrained=False, root=os.path.expanduser('~/.torch/models'), **kwargs):
net = CornerSqueeze(**kwargs)
if pretrained:
from model.model_store import get_model_file
full_name = name
net.load_state_dict(torch.load(get_model_file(full_name, root=root)))
return net
def get_yolact(name,
dataset,
pretrained=False,
root=os.path.expanduser('~/.torch/models'),
**kwargs):
net = Yolact(**kwargs)
if pretrained:
from model.model_store import get_model_file
full_name = '_'.join(('yolact', name, dataset))
net.load_state_dict(torch.load(get_model_file(full_name, root=root)))
return net
def get_dla(name,
levels,
channels,
block,
pretrained=False,
root=os.path.expanduser('~/.torch/models'),
**kwargs):
net = DLA(levels, channels, block=block, **kwargs)
if pretrained:
from model.model_store import get_model_file
net.load_state_dict(torch.load(get_model_file(name, root=root)))
return net
def get_simple_pose_resnet(base_name,
filters=(2048, 256, 256, 256),
pretrained=False,
root=os.path.expanduser('~/.torch/models'),
**kwargs):
net = SimplePoseResNet(base_name, num_deconv_filters=filters, **kwargs)
if pretrained:
from model.model_store import get_model_file
net.load_state_dict(
torch.load(get_model_file('simple_pose_%s' % base_name,
root=root)))
return net
def get_centernet(name,
features,
deconv_layers,
heads,
head_conv,
pretrained=False,
root=os.path.expanduser('~/.torch/models'),
**kwargs):
net = CenterNet(features, deconv_layers, heads, head_conv, **kwargs)
if pretrained:
import torch
from model.model_store import get_model_file
net.load_state_dict(torch.load(get_model_file(name, root=root)))
return net
def get_darknet(darknet_version, num_layers, pretrained=False,
root=os.path.join(os.path.expanduser('~'), '.torch/models'), **kwargs):
"""Get darknet by `version` and `num_layers` info.
Parameters
----------
darknet_version : str
Darknet version, choices are ['v3'].
num_layers : int
Number of layers.
pretrained : bool or str
Boolean value controls whether to load the default pretrained weights for model.
root : str, default '~/.torch/models'
Location for keeping the model parameters.
norm_layer : object
Normalization layer used (default: :class:`nn.BatchNorm`)
Can be :class:`nn.BatchNorm` or :class:`other normalization`.
norm_kwargs : dict
Additional `norm_layer` arguments.
Returns
-------
nn.Module
Darknet network.
Examples
--------
>>> model = get_darknet('v3', 53, pretrained=True)
>>> print(model)
"""
assert darknet_version in darknet_versions and darknet_version in darknet_spec, (
"Invalid darknet version: {}. Options are {}".format(
darknet_version, str(darknet_versions.keys())))
specs = darknet_spec[darknet_version]
assert num_layers in specs, (
"Invalid number of layers: {}. Options are {}".format(num_layers, str(specs.keys())))
layers, channels = specs[num_layers]
darknet_class = darknet_versions[darknet_version]
net = darknet_class(layers, channels, **kwargs)
if pretrained:
import torch
from model.model_store import get_model_file
from data.imagenet import ImageNetAttr
net.load_state_dict(torch.load(get_model_file('darknet%d' % num_layers, root=root)))
attrib = ImageNetAttr()
net.synset = attrib.synset
net.classes = attrib.classes
net.classes_long = attrib.classes_long
return net
def get_fcn(dataset='pascal_voc',
backbone='resnet50',
pretrained=False,
root=os.path.expanduser('~/.torch/models'),
pretrained_base=True,
**kwargs):
r"""FCN model from the paper `"Fully Convolutional Network for semantic segmentation"
<https://people.eecs.berkeley.edu/~jonlong/long_shelhamer_fcn.pdf>`_
Parameters
----------
dataset : str, default pascal_voc
The dataset that model pretrained on. (pascal_voc, ade20k)
pretrained : bool or str
Boolean value controls whether to load the default pretrained weights for model.
String value represents the hashtag for a certain version of pretrained weights.
ctx : Context, default CPU
The context in which to load the pretrained weights.
root : str, default '~/.mxnet/models'
Location for keeping the model parameters.
pretrained_base : bool or str, default True
This will load pretrained backbone network, that was trained on ImageNet.
Examples
--------
>>> model = get_fcn(dataset='pascal_voc', backbone='resnet50', pretrained=False)
>>> print(model)
"""
acronyms = {
'pascal_voc': 'voc',
'pascal_paper': 'voc',
'pascal_aug': 'voc',
'ade20k': 'ade',
'coco': 'coco',
'citys': 'citys'
}
from data import datasets
# infer number of classes
model = FCN(datasets[dataset].NUM_CLASS,
backbone=backbone,
pretrained_base=pretrained_base,
**kwargs)
if pretrained:
from model.model_store import get_model_file
model.load_state_dict(
torch.load(
get_model_file('fcn_%s_%s' % (backbone, acronyms[dataset]),
root=root)))
return model
def get_resnet(version,
num_layers,
pretrained=False,
root=os.path.expanduser('~/.torch/models'),
**kwargs):
r"""ResNet V1 model from `"Deep Residual Learning for Image Recognition"
<http://arxiv.org/abs/1512.03385>`_ paper.
ResNet V2 model from `"Identity Mappings in Deep Residual Networks"
<https://arxiv.org/abs/1603.05027>`_ paper.
Parameters
----------
version : int
Version of ResNet. Options are 1, 2.
num_layers : int
Numbers of layers. Options are 18, 34, 50, 101, 152.
pretrained : bool or str
Boolean value controls whether to load the default pretrained weights for model.
root : str, default $~/.torch/models
Location for keeping the model parameters.
norm_layer : object
Normalization layer used (default: :class:`nn.BatchNorm`)
Can be :class:`nn.BatchNorm` or :class:`other normalization`.
norm_kwargs : dict
Additional `norm_layer` arguments
"""
assert num_layers in resnet_spec, \
"Invalid number of layers: %d. Options are %s" % (
num_layers, str(resnet_spec.keys()))
block_type, layers, channels = resnet_spec[num_layers]
assert 1 <= version <= 2, \
"Invalid resnet version: %d. Options are 1 and 2." % version
resnet_class = resnet_net_versions[version - 1]
block_class = resnet_block_versions[version - 1][block_type]
net = resnet_class(block_class, layers, channels, **kwargs)
if pretrained:
import torch
from model.model_store import get_model_file
net.load_state_dict(
torch.load(get_model_file('resnet%d_v%d' % (num_layers, version),
root=root),
map_location=lambda storage, loc: storage))
from data.imagenet import ImageNetAttr
attrib = ImageNetAttr()
net.synset = attrib.synset
net.classes = attrib.classes
net.classes_long = attrib.classes_long
return net
def get_ocnet(dataset='pascal_voc', backbone='resnet50', oc='base', pretrained=False, pretrained_base=True,
jpu=False, root=os.path.expanduser('~/.torch/models'), **kwargs):
acronyms = {
'pascal_voc': 'voc',
'citys': 'citys',
}
from data import datasets
# infer number of classes
model = OCNet(datasets[dataset].NUM_CLASS, oc=oc, backbone=backbone,
jpu=jpu, pretrained_base=pretrained_base, **kwargs)
if pretrained:
from model.model_store import get_model_file
name = 'ocnet_%s_%s_%s' % (oc, backbone, acronyms[dataset])
name = name + '_jpu' if jpu else name
model.load_state_dict(torch.load(get_model_file(name, root=root)))
return model
def get_cifar_wide_resnet(num_layers,
width_factor=1,
drop_rate=0.0,
pretrained=False,
root=os.path.expanduser('~/.torch/models'),
**kwargs):
r"""ResNet V1 model from `"Deep Residual Learning for Image Recognition"
<http://arxiv.org/abs/1512.03385>`_ paper.
ResNet V2 model from `"Identity Mappings in Deep Residual Networks"
<https://arxiv.org/abs/1603.05027>`_ paper.
Parameters
----------
num_layers : int
Numbers of layers. Needs to be an integer in the form of 6*n+2, e.g. 20, 56, 110, 164.
width_factor: int
The width factor to apply to the number of channels from the original resnet.
drop_rate: float
The rate of dropout.
pretrained : bool, default False
Whether to load the pretrained weights for model.
root : str, default '~/.torch/models'
Location for keeping the model parameters.
norm_layer : object
Normalization layer used (default: :class:`nn.BatchNorm`)
Can be :class:`nn.BatchNorm` or :class:`other normalization`.
norm_kwargs : dict
Additional `norm_layer` arguments.
"""
assert (num_layers - 4) % 6 == 0
n = (num_layers - 4) // 6
layers = [n] * 3
channels = [3, 16, 16 * width_factor, 32 * width_factor, 64 * width_factor]
net = CIFARWideResNet(CIFARBasicBlockV2, layers, channels, drop_rate,
**kwargs)
if pretrained:
import torch
from model.model_store import get_model_file
net.load_state_dict(
torch.load(get_model_file('cifar_wideresnet%d_%d' %
(num_layers, width_factor),
root=root),
map_location=lambda storage, loc: storage))
return net
def get_cifar_resnext(num_layers,
cardinality=16,
bottleneck_width=64,
pretrained=False,
root=os.path.expanduser('~/.torch/models'),
**kwargs):
r"""ResNext model from `"Aggregated Residual Transformations for Deep Neural Networks"
<http://arxiv.org/abs/1611.05431>`_ paper.
Parameters
----------
num_layers : int
Numbers of layers. Needs to be an integer in the form of 9*n+2, e.g. 29
cardinality: int
Number of groups
bottleneck_width: int
Width of bottleneck block
pretrained : bool, default False
Whether to load the pretrained weights for model.
root : str, default '~/.torch/models'
Location for keeping the model parameters.
norm_layer : object
Normalization layer used (default: :class:`nn.BatchNorm`)
Can be :class:`nn.BatchNorm` or :class:`other normalization`.
norm_kwargs : dict
Additional `norm_layer` arguments.
"""
assert (num_layers - 2) % 9 == 0
layer = (num_layers - 2) // 9
layers = [layer] * 3
net = CIFARResNext(layers, cardinality, bottleneck_width, **kwargs)
if pretrained:
import torch
from model.model_store import get_model_file
net.load_state_dict(
torch.load(
get_model_file('cifar_resnext%d_%dx%dd' %
(num_layers, cardinality, bottleneck_width),
root=root)))
return net
def get_mobilenet_v2(multiplier, pretrained=False, root=os.path.expanduser('~/.torch/models'),
**kwargs):
r"""MobileNetV2 model from the
`"Inverted Residuals and Linear Bottlenecks:
Mobile Networks for Classification, Detection and Segmentation"
<https://arxiv.org/abs/1801.04381>`_ paper.
Parameters
----------
multiplier : float
The width multiplier for controlling the model size. Only multipliers that are no
less than 0.25 are supported. The actual number of channels is equal to the original
channel size multiplied by this multiplier.
pretrained : bool or str
Boolean value controls whether to load the default pretrained weights for model.
String value represents the hashtag for a certain version of pretrained weights.
root : str, default ~/.torch/models
Location for keeping the model parameters.
norm_layer : object
Normalization layer used (default: :class:`nn.BatchNorm`)
Can be :class:`nn.BatchNorm` or :class:`other normalization`.
norm_kwargs : dict
Additional `norm_layer` arguments.
"""
net = MobileNetV2(multiplier, **kwargs)
if pretrained:
import torch
from model.model_store import get_model_file
version_suffix = '{0:.2f}'.format(multiplier)
if version_suffix in ('1.00', '0.50'):
version_suffix = version_suffix[:-1]
net.load_state_dict(torch.load(get_model_file('mobilenetv2_%s' % version_suffix, root=root)))
from data.imagenet import ImageNetAttr
attrib = ImageNetAttr()
net.synset = attrib.synset
net.classes = attrib.classes
net.classes_long = attrib.classes_long
return net
def get_psp(dataset='pascal_voc', backbone='resnet50', pretrained=False, pretrained_base=True,
root=os.path.join(os.path.expanduser('~'), '.torch/models'), **kwargs):
r"""Pyramid Scene Parsing Network
Parameters
----------
dataset : str, default pascal_voc
The dataset that model pretrained on. (pascal_voc, ade20k)
pretrained : bool or str
Boolean value controls whether to load the default pretrained weights for model.
String value represents the hashtag for a certain version of pretrained weights.
ctx : Context, default CPU
The context in which to load the pretrained weights.
root : str, default '~/.mxnet/models'
Location for keeping the model parameters.
pretrained_base : bool or str, default True
This will load pretrained backbone network, that was trained on ImageNet.
Examples
--------
>>> model = get_fcn(dataset='pascal_voc', backbone='resnet50', pretrained=False)
>>> print(model)
"""
acronyms = {
'pascal_voc': 'voc',
'pascal_paper': 'voc',
'pascal_aug': 'voc',
'ade20k': 'ade',
'coco': 'coco',
'citys': 'citys',
}
from data import datasets
# infer number of classes
model = PSPNet(datasets[dataset].NUM_CLASS, backbone=backbone,
pretrained_base=pretrained_base, **kwargs)
if pretrained:
from model.model_store import get_model_file
model.load_state_dict(torch.load(get_model_file('psp_%s_%s' % (backbone, acronyms[dataset]),
root=root)))
return model
def alexnet(pretrained=False,
root=os.path.expanduser('~/.torch/models'),
**kwargs):
r"""AlexNet model from the `"One weird trick..." <https://arxiv.org/abs/1404.5997>`_ paper.
Parameters
----------
pretrained : bool or str
Boolean value controls whether to load the default pretrained weights for model.
root : str, default '~/.torch/models'
Location for keeping the model parameters.
"""
net = AlexNet(**kwargs)
if pretrained:
import torch
from model.model_store import get_model_file
net.load_state_dict(torch.load(get_model_file('alexnet', root=root)))
from data.imagenet import ImageNetAttr
attrib = ImageNetAttr()
net.synset = attrib.synset
net.classes = attrib.classes
net.classes_long = attrib.classes_long
return net
def get_vgg_atrous_extractor(num_layers,
im_size,
pretrained=False,
root=os.path.expanduser('~/.torch/models'),
**kwargs):
"""Get VGG atrous feature extractor networks.
Parameters
----------
num_layers : int
VGG types, can be 11,13,16,19.
im_size : int
VGG detection input size, can be 300, 512.
pretrained : bool or str
Boolean value controls whether to load the default pretrained weights for model.
root : str
Model weights storing path.
Returns
-------
nn.Module
The returned network.
"""
layers, filters = vgg_spec[num_layers]
extras = extra_spec[im_size]
net = VGGAtrousExtractor(layers, filters, extras, **kwargs)
if pretrained:
from model.model_store import get_model_file
batch_norm_suffix = '_bn' if kwargs.get('batch_norm') else ''
net.load_state_dict(
torch.load(
get_model_file('vgg%d_atrous%s_%d' %
(num_layers, batch_norm_suffix, im_size),
root=root)))
return net
def get_yolov3(name,
stages,
out_channels,
block_channels,
filters,
anchors,
strides,
classes,
dataset,
pretrained=False,
root=os.path.expanduser('~/.torch/models'),
**kwargs):
"""Get YOLOV3 models.
Parameters
----------
name : str or None
Model name, if `None` is used, you must specify `features` to be a `nn.Module`.
stages : iterable of str or `nn.Module`
List of network internal output names, in order to specify which layers are
used for predicting bbox values.
If `name` is `None`, `features` must be a `nn.Module` which generate multiple
outputs for prediction.
filters : iterable of float or None
List of convolution layer channels which is going to be appended to the base
network feature extractor. If `name` is `None`, this is ignored.
sizes : iterable fo float
Sizes of anchor boxes, this should be a list of floats, in incremental order.
The length of `sizes` must be len(layers) + 1. For example, a two stage SSD
model can have ``sizes = [30, 60, 90]``, and it converts to `[30, 60]` and
`[60, 90]` for the two stages, respectively. For more details, please refer
to original paper.
ratios : iterable of list
Aspect ratios of anchors in each output layer. Its length must be equals
to the number of SSD output layers.
steps : list of int
Step size of anchor boxes in each output layer.
classes : iterable of str
Names of categories.
dataset : str
Name of dataset. This is used to identify model name because models trained on
different datasets are going to be very different.
pretrained : bool or str
Boolean value controls whether to load the default pretrained weights for model.
String value represents the hashtag for a certain version of pretrained weights.
pretrained_base : bool or str, optional, default is True
Load pretrained base network, the extra layers are randomized. Note that
if pretrained is `True`, this has no effect.
root : str
Model weights storing path.
norm_layer : object
Normalization layer used (default: :class:`mxnet.gluon.nn.BatchNorm`)
Can be :class:`mxnet.gluon.nn.BatchNorm` or :class:`mxnet.gluon.contrib.nn.SyncBatchNorm`.
norm_kwargs : dict
Additional `norm_layer` arguments, for example `num_devices=4`
for :class:`mxnet.gluon.contrib.nn.SyncBatchNorm`.
Returns
-------
HybridBlock
A YOLOV3 detection network.
"""
net = YOLOV3(stages,
out_channels,
block_channels,
filters,
anchors,
strides,
classes=classes,
**kwargs)
if pretrained:
from model.model_store import get_model_file
full_name = '_'.join(('yolo3', name, dataset))
net.load_state_dict(torch.load(get_model_file(full_name, root=root)))
return net
def get_ssd(name,
base_size,
features,
filters,
channels,
sizes,
ratios,
steps,
classes,
dataset,
pretrained=False,
pretrained_base=True,
root=os.path.expanduser('~/.torch/models'),
**kwargs):
"""Get SSD models.
Parameters
----------
name : str or None
Model name, if `None` is used, you must specify `features` to be a `nn.Module`.
base_size : int
Base image size for training, this is fixed once training is assigned.
A fixed base size still allows you to have variable input size during test.
features : iterable of str or `nn.Module`
List of network internal output names, in order to specify which layers are
used for predicting bbox values.
If `name` is `None`, `features` must be a `nn.Module` which generate multiple
outputs for prediction.
filters : iterable of float or None
List of convolution layer channels which is going to be appended to the base
network feature extractor. If `name` is `None`, this is ignored.
channels : iterable of float or None
List of convolution layer in channels
sizes : iterable fo float
Sizes of anchor boxes, this should be a list of floats, in incremental order.
The length of `sizes` must be len(layers) + 1. For example, a two stage SSD
model can have ``sizes = [30, 60, 90]``, and it converts to `[30, 60]` and
`[60, 90]` for the two stages, respectively. For more details, please refer
to original paper.
ratios : iterable of list
Aspect ratios of anchors in each output layer. Its length must be equals
to the number of SSD output layers.
steps : list of int
Step size of anchor boxes in each output layer.
classes : iterable of str
Names of categories.
dataset : str
Name of dataset. This is used to identify model name because models trained on
different datasets are going to be very different.
pretrained : bool or str
Boolean value controls whether to load the default pretrained weights for model.
pretrained_base : bool or str, optional, default is True
Load pretrained base network, the extra layers are randomized. Note that
if pretrained is `True`, this has no effect.
root : str
Model weights storing path.
norm_layer : object
Normalization layer used (default: :class:`nn.BatchNorm`)
Can be :class:`nn.BatchNorm` or :class:`other normalization`.
norm_kwargs : dict
Additional `norm_layer` arguments
Returns
-------
nn.Module
A SSD detection network.
"""
pretrained_base = False if pretrained else pretrained_base
base_name = None if callable(features) else name
net = SSD(base_name,
base_size,
features,
filters,
channels,
sizes,
ratios,
steps,
pretrained=pretrained_base,
classes=classes,
**kwargs)
if pretrained:
from model.model_store import get_model_file
full_name = '_'.join(('ssd', str(base_size), name, dataset))
net.load_state_dict(torch.load(get_model_file(full_name, root=root)))
return net
