Shortcuts

torchvision.models

The models subpackage contains definitions for the following model architectures:

You can construct a model with random weights by calling its constructor:

import torchvision.models as models
resnet18 = models.resnet18()
alexnet = models.alexnet()
vgg16 = models.vgg16()
squeezenet = models.squeezenet1_0()
densenet = models.densenet161()
inception = models.inception_v3()
googlenet = models.googlenet()
shufflenet = models.shufflenetv2()

We provide pre-trained models, using the PyTorch torch.utils.model_zoo. These can be constructed by passing pretrained=True:

import torchvision.models as models
resnet18 = models.resnet18(pretrained=True)
alexnet = models.alexnet(pretrained=True)
squeezenet = models.squeezenet1_0(pretrained=True)
vgg16 = models.vgg16(pretrained=True)
densenet = models.densenet161(pretrained=True)
inception = models.inception_v3(pretrained=True)
googlenet = models.googlenet(pretrained=True)
shufflenet = models.shufflenetv2(pretrained=True)

Instancing a pre-trained model will download its weights to a cache directory. This directory can be set using the TORCH_MODEL_ZOO environment variable. See torch.utils.model_zoo.load_url() for details.

Some models use modules which have different training and evaluation behavior, such as batch normalization. To switch between these modes, use model.train() or model.eval() as appropriate. See train() or eval() for details.

All pre-trained models expect input images normalized in the same way, i.e. mini-batches of 3-channel RGB images of shape (3 x H x W), where H and W are expected to be at least 224. The images have to be loaded in to a range of [0, 1] and then normalized using mean = [0.485, 0.456, 0.406] and std = [0.229, 0.224, 0.225]. You can use the following transform to normalize:

normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225])

An example of such normalization can be found in the imagenet example here

ImageNet 1-crop error rates (224x224)

Network

Top-1 error

Top-5 error

AlexNet

43.45

20.91

VGG-11

30.98

11.37

VGG-13

30.07

10.75

VGG-16

28.41

9.62

VGG-19

27.62

9.12

VGG-11 with batch normalization

29.62

10.19

VGG-13 with batch normalization

28.45

9.63

VGG-16 with batch normalization

26.63

8.50

VGG-19 with batch normalization

25.76

8.15

ResNet-18

30.24

10.92

ResNet-34

26.70

8.58

ResNet-50

23.85

7.13

ResNet-101

22.63

6.44

ResNet-152

21.69

5.94

SqueezeNet 1.0

41.90

19.58

SqueezeNet 1.1

41.81

19.38

Densenet-121

25.35

7.83

Densenet-169

24.00

7.00

Densenet-201

22.80

6.43

Densenet-161

22.35

6.20

Inception v3

22.55

6.44

GoogleNet

30.22

10.47

ShuffleNet V2

30.64

11.68

Alexnet

torchvision.models.alexnet(pretrained=False, progress=True, **kwargs)[source]

AlexNet model architecture from the “One weird trick…” paper.

Parameters
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

VGG

torchvision.models.vgg11(pretrained=False, progress=True, **kwargs)[source]

VGG 11-layer model (configuration “A”)

Parameters
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

torchvision.models.vgg11_bn(pretrained=False, progress=True, **kwargs)[source]

VGG 11-layer model (configuration “A”) with batch normalization

Parameters
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

torchvision.models.vgg13(pretrained=False, progress=True, **kwargs)[source]

VGG 13-layer model (configuration “B”)

Parameters
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

torchvision.models.vgg13_bn(pretrained=False, progress=True, **kwargs)[source]

VGG 13-layer model (configuration “B”) with batch normalization

Parameters
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

torchvision.models.vgg16(pretrained=False, progress=True, **kwargs)[source]

VGG 16-layer model (configuration “D”)

Parameters
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

torchvision.models.vgg16_bn(pretrained=False, progress=True, **kwargs)[source]

VGG 16-layer model (configuration “D”) with batch normalization

Parameters
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

torchvision.models.vgg19(pretrained=False, progress=True, **kwargs)[source]

VGG 19-layer model (configuration “E”)

Parameters
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

torchvision.models.vgg19_bn(pretrained=False, progress=True, **kwargs)[source]

VGG 19-layer model (configuration ‘E’) with batch normalization

Parameters
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

ResNet

torchvision.models.resnet18(pretrained=False, progress=True, **kwargs)[source]

Constructs a ResNet-18 model.

Parameters
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

torchvision.models.resnet34(pretrained=False, progress=True, **kwargs)[source]

Constructs a ResNet-34 model.

Parameters
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

torchvision.models.resnet50(pretrained=False, progress=True, **kwargs)[source]

Constructs a ResNet-50 model.

Parameters
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

torchvision.models.resnet101(pretrained=False, progress=True, **kwargs)[source]

Constructs a ResNet-101 model.

Parameters
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

torchvision.models.resnet152(pretrained=False, progress=True, **kwargs)[source]

Constructs a ResNet-152 model.

Parameters
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

SqueezeNet

torchvision.models.squeezenet1_0(pretrained=False, progress=True, **kwargs)[source]

SqueezeNet model architecture from the “SqueezeNet: AlexNet-level accuracy with 50x fewer parameters and <0.5MB model size” paper.

Parameters
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

torchvision.models.squeezenet1_1(pretrained=False, progress=True, **kwargs)[source]

SqueezeNet 1.1 model from the official SqueezeNet repo. SqueezeNet 1.1 has 2.4x less computation and slightly fewer parameters than SqueezeNet 1.0, without sacrificing accuracy.

Parameters
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

DenseNet

torchvision.models.densenet121(pretrained=False, progress=True, **kwargs)[source]

Densenet-121 model from “Densely Connected Convolutional Networks”

Parameters
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

torchvision.models.densenet169(pretrained=False, progress=True, **kwargs)[source]

Densenet-169 model from “Densely Connected Convolutional Networks”

Parameters
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

torchvision.models.densenet161(pretrained=False, progress=True, **kwargs)[source]

Densenet-161 model from “Densely Connected Convolutional Networks”

Parameters
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

torchvision.models.densenet201(pretrained=False, progress=True, **kwargs)[source]

Densenet-201 model from “Densely Connected Convolutional Networks”

Parameters
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

Inception v3

torchvision.models.inception_v3(pretrained=False, progress=True, **kwargs)[source]

Inception v3 model architecture from “Rethinking the Inception Architecture for Computer Vision”.

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.

Parameters
  • 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

GoogLeNet

torchvision.models.googlenet(pretrained=False, progress=True, **kwargs)[source]

GoogLeNet (Inception v1) model architecture from “Going Deeper with Convolutions”.

Parameters
  • 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, adds two auxiliary branches that can improve training. Default: False when pretrained is True otherwise True

  • transform_input (bool) – If True, preprocesses the input according to the method with which it was trained on ImageNet. Default: False

ShuffleNet v2

Docs

Access comprehensive developer documentation for PyTorch

View Docs

Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

Resources

Find development resources and get your questions answered

View Resources