torchreid.models

Interface

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

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()

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)

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)

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)

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)

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)

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)

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)

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)

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)

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)

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)

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)

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)

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)

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)

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)

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)

Omni-Scale Network.

Reference:

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

• Zhou et al. Learning Generalisable Omni-Scale Representations for Person Re-Identification. TPAMI, 2021.

class torchreid.models.osnet_ain.OSNet(num_classes, blocks, layers, channels, feature_dim=512, loss='softmax', conv1_IN=False, **kwargs)

Omni-Scale Network.

Reference:

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

• Zhou et al. Learning Generalisable Omni-Scale Representations for Person Re-Identification. TPAMI, 2021.