torchreid.models

Interface

torchreid.models.__init__.build_model(name, num_classes, loss='softmax', pretrained=True, use_gpu=True)[source]

A function wrapper for building a model.

Parameters:
  • name (str) – model name.
  • num_classes (int) – number of training identities.
  • loss (str, optional) – loss function to optimize the model. Currently supports “softmax” and “triplet”. Default is “softmax”.
  • pretrained (bool, optional) – whether to load ImageNet-pretrained weights. Default is True.
  • use_gpu (bool, optional) – whether to use gpu. Default is True.
Returns:

nn.Module

Examples::
>>> from torchreid import models
>>> model = models.build_model('resnet50', 751, loss='softmax')
torchreid.models.__init__.show_avai_models()[source]

Displays available models.

Examples::
>>> from torchreid import models
>>> models.show_avai_models()

ImageNet Classification Models

class torchreid.models.resnet.ResNet(num_classes, loss, block, layers, zero_init_residual=False, groups=1, width_per_group=64, replace_stride_with_dilation=None, norm_layer=None, last_stride=2, fc_dims=None, dropout_p=None, **kwargs)[source]

Residual network.

Reference:
  • He et al. Deep Residual Learning for Image Recognition. CVPR 2016.
  • Xie et al. Aggregated Residual Transformations for Deep Neural Networks. CVPR 2017.
Public keys:
  • resnet18: ResNet18.
  • resnet34: ResNet34.
  • resnet50: ResNet50.
  • resnet101: ResNet101.
  • resnet152: ResNet152.
  • resnext50_32x4d: ResNeXt50.
  • resnext101_32x8d: ResNeXt101.
  • resnet50_fc512: ResNet50 + FC.
class torchreid.models.senet.SENet(num_classes, loss, block, layers, groups, reduction, dropout_p=0.2, inplanes=128, input_3x3=True, downsample_kernel_size=3, downsample_padding=1, last_stride=2, fc_dims=None, **kwargs)[source]

Squeeze-and-excitation network.

Reference:
Hu et al. Squeeze-and-Excitation Networks. CVPR 2018.
Public keys:
  • senet154: SENet154.
  • se_resnet50: ResNet50 + SE.
  • se_resnet101: ResNet101 + SE.
  • se_resnet152: ResNet152 + SE.
  • se_resnext50_32x4d: ResNeXt50 (groups=32, width=4) + SE.
  • se_resnext101_32x4d: ResNeXt101 (groups=32, width=4) + SE.
  • se_resnet50_fc512: (ResNet50 + SE) + FC.
class torchreid.models.densenet.DenseNet(num_classes, loss, growth_rate=32, block_config=(6, 12, 24, 16), num_init_features=64, bn_size=4, drop_rate=0, fc_dims=None, dropout_p=None, **kwargs)[source]

Densely connected network.

Reference:
Huang et al. Densely Connected Convolutional Networks. CVPR 2017.
Public keys:
  • densenet121: DenseNet121.
  • densenet169: DenseNet169.
  • densenet201: DenseNet201.
  • densenet161: DenseNet161.
  • densenet121_fc512: DenseNet121 + FC.
class torchreid.models.inceptionresnetv2.InceptionResNetV2(num_classes, loss='softmax', **kwargs)[source]

Inception-ResNet-V2.

Reference:
Szegedy et al. Inception-v4, Inception-ResNet and the Impact of Residual Connections on Learning. AAAI 2017.
Public keys:
  • inceptionresnetv2: Inception-ResNet-V2.
class torchreid.models.inceptionv4.InceptionV4(num_classes, loss, **kwargs)[source]

Inception-v4.

Reference:
Szegedy et al. Inception-v4, Inception-ResNet and the Impact of Residual Connections on Learning. AAAI 2017.
Public keys:
  • inceptionv4: InceptionV4.
class torchreid.models.xception.Xception(num_classes, loss, fc_dims=None, dropout_p=None, **kwargs)[source]

Xception.

Reference:
Chollet. Xception: Deep Learning with Depthwise Separable Convolutions. CVPR 2017.
Public keys:
  • xception: Xception.

Lightweight Models

class torchreid.models.nasnet.NASNetAMobile(num_classes, loss, stem_filters=32, penultimate_filters=1056, filters_multiplier=2, **kwargs)[source]

Neural Architecture Search (NAS).

Reference:
Zoph et al. Learning Transferable Architectures for Scalable Image Recognition. CVPR 2018.
Public keys:
  • nasnetamobile: NASNet-A Mobile.
class torchreid.models.mobilenetv2.MobileNetV2(num_classes, width_mult=1, loss='softmax', fc_dims=None, dropout_p=None, **kwargs)[source]

MobileNetV2.

Reference:
Sandler et al. MobileNetV2: Inverted Residuals and Linear Bottlenecks. CVPR 2018.
Public keys:
  • mobilenetv2_x1_0: MobileNetV2 x1.0.
  • mobilenetv2_x1_4: MobileNetV2 x1.4.
class torchreid.models.shufflenet.ShuffleNet(num_classes, loss='softmax', num_groups=3, **kwargs)[source]

ShuffleNet.

Reference:
Zhang et al. ShuffleNet: An Extremely Efficient Convolutional Neural Network for Mobile Devices. CVPR 2018.
Public keys:
  • shufflenet: ShuffleNet (groups=3).
class torchreid.models.squeezenet.SqueezeNet(num_classes, loss, version=1.0, fc_dims=None, dropout_p=None, **kwargs)[source]

SqueezeNet.

Reference:
Iandola et al. SqueezeNet: AlexNet-level accuracy with 50x fewer parameters and< 0.5 MB model size. arXiv:1602.07360.
Public keys:
  • squeezenet1_0: SqueezeNet (version=1.0).
  • squeezenet1_1: SqueezeNet (version=1.1).
  • squeezenet1_0_fc512: SqueezeNet (version=1.0) + FC.
class torchreid.models.shufflenetv2.ShuffleNetV2(num_classes, loss, stages_repeats, stages_out_channels, **kwargs)[source]

ShuffleNetV2.

Reference:
Ma et al. ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design. ECCV 2018.
Public keys:
  • shufflenet_v2_x0_5: ShuffleNetV2 x0.5.
  • shufflenet_v2_x1_0: ShuffleNetV2 x1.0.
  • shufflenet_v2_x1_5: ShuffleNetV2 x1.5.
  • shufflenet_v2_x2_0: ShuffleNetV2 x2.0.

ReID-specific Models

class torchreid.models.mudeep.MuDeep(num_classes, loss='softmax', **kwargs)[source]

Multiscale deep neural network.

Reference:
Qian et al. Multi-scale Deep Learning Architectures for Person Re-identification. ICCV 2017.
Public keys:
  • mudeep: Multiscale deep neural network.
class torchreid.models.resnetmid.ResNetMid(num_classes, loss, block, layers, last_stride=2, fc_dims=None, **kwargs)[source]

Residual network + mid-level features.

Reference:
Yu et al. The Devil is in the Middle: Exploiting Mid-level Representations for Cross-Domain Instance Matching. arXiv:1711.08106.
Public keys:
  • resnet50mid: ResNet50 + mid-level feature fusion.
class torchreid.models.hacnn.HACNN(num_classes, loss='softmax', nchannels=[128, 256, 384], feat_dim=512, learn_region=True, use_gpu=True, **kwargs)[source]

Harmonious Attention Convolutional Neural Network.

Reference:
Li et al. Harmonious Attention Network for Person Re-identification. CVPR 2018.
Public keys:
  • hacnn: HACNN.
class torchreid.models.pcb.PCB(num_classes, loss, block, layers, parts=6, reduced_dim=256, nonlinear='relu', **kwargs)[source]

Part-based Convolutional Baseline.

Reference:
Sun et al. Beyond Part Models: Person Retrieval with Refined Part Pooling (and A Strong Convolutional Baseline). ECCV 2018.
Public keys:
  • pcb_p4: PCB with 4-part strips.
  • pcb_p6: PCB with 6-part strips.
class torchreid.models.mlfn.MLFN(num_classes, loss='softmax', groups=32, channels=[64, 256, 512, 1024, 2048], embed_dim=1024, **kwargs)[source]

Multi-Level Factorisation Net.

Reference:
Chang et al. Multi-Level Factorisation Net for Person Re-Identification. CVPR 2018.
Public keys:
  • mlfn: MLFN (Multi-Level Factorisation Net).
class torchreid.models.osnet.OSNet(num_classes, blocks, layers, channels, feature_dim=512, loss='softmax', IN=False, **kwargs)[source]

Omni-Scale Network.

Reference:
  • Zhou et al. Omni-Scale Feature Learning for Person Re-Identification. ICCV, 2019.