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