models.rst 5.88 KB
Newer Older
Sasank Chilamkurthy's avatar
Sasank Chilamkurthy committed
1
2
torchvision.models
==================
Sasank Chilamkurthy's avatar
Sasank Chilamkurthy committed
3

Sasank Chilamkurthy's avatar
Sasank Chilamkurthy committed
4
5
6
7
8
9
10
11
12
The models subpackage contains definitions for the following model
architectures:

-  `AlexNet`_
-  `VGG`_
-  `ResNet`_
-  `SqueezeNet`_
-  `DenseNet`_
-  `Inception`_ v3
13
-  `GoogLeNet`_
Bar's avatar
Bar committed
14
-  `ShuffleNet`_ v2
15
16
-  `MobileNet`_ v2
-  `ResNeXt`_
Sasank Chilamkurthy's avatar
Sasank Chilamkurthy committed
17
18
19
20
21
22
23
24
25
26

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

.. code:: python

    import torchvision.models as models
    resnet18 = models.resnet18()
    alexnet = models.alexnet()
    vgg16 = models.vgg16()
    squeezenet = models.squeezenet1_0()
Ahmed Abdo's avatar
Ahmed Abdo committed
27
    densenet = models.densenet161()
Sasank Chilamkurthy's avatar
Sasank Chilamkurthy committed
28
    inception = models.inception_v3()
29
    googlenet = models.googlenet()
30
31
32
    shufflenet = models.shufflenet_v2_x1_0()
    mobilenet = models.mobilenet_v2()
    resnext50_32x4d = models.resnext50_32x4d()
Sasank Chilamkurthy's avatar
Sasank Chilamkurthy committed
33
34
35
36
37
38
39
40
41
42
43

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

.. code:: python

    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)
Ahmed Abdo's avatar
Ahmed Abdo committed
44
    densenet = models.densenet161(pretrained=True)
Sasank Chilamkurthy's avatar
Sasank Chilamkurthy committed
45
    inception = models.inception_v3(pretrained=True)
46
    googlenet = models.googlenet(pretrained=True)
47
48
49
    shufflenet = models.shufflenet_v2_x1_0(pretrained=True)
    mobilenet = models.mobilenet_v2(pretrained=True)
    resnext50_32x4d = models.resnext50_32x4d(pretrained=True)
Sasank Chilamkurthy's avatar
Sasank Chilamkurthy committed
50

51
52
53
54
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
:func:`torch.utils.model_zoo.load_url` for details.

55
56
57
58
59
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
:meth:`~torch.nn.Module.train` or :meth:`~torch.nn.Module.eval` for details. 

Sasank Chilamkurthy's avatar
Sasank Chilamkurthy committed
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
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 <https://github.com/pytorch/examples/blob/42e5b996718797e45c46a25c55b031e6768f8440/imagenet/main.py#L89-L101>`_

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
99
GoogleNet                         30.22           10.47
Bar's avatar
Bar committed
100
ShuffleNet V2                     30.64           11.68
101
102
103
MobileNet V2                      28.12           9.71
ResNeXt-50-32x4d                  22.38           6.30
ResNeXt-101-32x8d                 20.69           5.47
Sasank Chilamkurthy's avatar
Sasank Chilamkurthy committed
104
105
106
107
108
109
110
111
112
================================  =============   =============


.. _AlexNet: https://arxiv.org/abs/1404.5997
.. _VGG: https://arxiv.org/abs/1409.1556
.. _ResNet: https://arxiv.org/abs/1512.03385
.. _SqueezeNet: https://arxiv.org/abs/1602.07360
.. _DenseNet: https://arxiv.org/abs/1608.06993
.. _Inception: https://arxiv.org/abs/1512.00567
113
.. _GoogLeNet: https://arxiv.org/abs/1409.4842
Bar's avatar
Bar committed
114
.. _ShuffleNet: https://arxiv.org/abs/1807.11164
115
116
.. _MobileNet: https://arxiv.org/abs/1801.04381
.. _ResNeXt: https://arxiv.org/abs/1611.05431
Sasank Chilamkurthy's avatar
Sasank Chilamkurthy committed
117
118
119

.. currentmodule:: torchvision.models

Sasank Chilamkurthy's avatar
Sasank Chilamkurthy committed
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
Alexnet
-------

.. autofunction:: alexnet

VGG
---

.. autofunction:: vgg11
.. autofunction:: vgg11_bn
.. autofunction:: vgg13
.. autofunction:: vgg13_bn
.. autofunction:: vgg16
.. autofunction:: vgg16_bn
.. autofunction:: vgg19
.. autofunction:: vgg19_bn


ResNet
------

.. autofunction:: resnet18
.. autofunction:: resnet34
.. autofunction:: resnet50
.. autofunction:: resnet101
.. autofunction:: resnet152

SqueezeNet
----------

.. autofunction:: squeezenet1_0
.. autofunction:: squeezenet1_1

Sangwhan Moon's avatar
Sangwhan Moon committed
153
DenseNet
Sasank Chilamkurthy's avatar
Sasank Chilamkurthy committed
154
155
156
157
158
159
160
161
162
163
164
165
---------

.. autofunction:: densenet121
.. autofunction:: densenet169
.. autofunction:: densenet161
.. autofunction:: densenet201

Inception v3
------------

.. autofunction:: inception_v3

166
167
168
169
170
GoogLeNet
------------

.. autofunction:: googlenet

Bar's avatar
Bar committed
171
172
173
ShuffleNet v2
-------------

174
175
176
177
.. autofunction:: shufflenet_v2_x0_5
.. autofunction:: shufflenet_v2_x1_0
.. autofunction:: shufflenet_v2_x1_5
.. autofunction:: shufflenet_v2_x2_0
Bar's avatar
Bar committed
178

179
180
181
182
183
184
185
186
187
188
189
MobileNet v2
-------------

.. autofunction:: mobilenet_v2

ResNext
-------------

.. autofunction:: resnext50_32x4d
.. autofunction:: resnext101_32x8d