models.rst 5.12 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
Sasank Chilamkurthy's avatar
Sasank Chilamkurthy committed
15
16
17
18
19
20
21
22
23
24

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
25
    densenet = models.densenet161()
Sasank Chilamkurthy's avatar
Sasank Chilamkurthy committed
26
    inception = models.inception_v3()
27
    googlenet = models.googlenet()
Bar's avatar
Bar committed
28
    shufflenet = models.shufflenetv2()
Sasank Chilamkurthy's avatar
Sasank Chilamkurthy committed
29
30
31
32
33
34
35
36
37
38
39

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
40
    densenet = models.densenet161(pretrained=True)
Sasank Chilamkurthy's avatar
Sasank Chilamkurthy committed
41
    inception = models.inception_v3(pretrained=True)
42
    googlenet = models.googlenet(pretrained=True)
Bar's avatar
Bar committed
43
    shufflenet = models.shufflenetv2(pretrained=True)
Sasank Chilamkurthy's avatar
Sasank Chilamkurthy committed
44

45
46
47
48
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.

49
50
51
52
53
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
54
55
56
57
58
59
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
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
93
GoogleNet                         30.22           10.47
Bar's avatar
Bar committed
94
ShuffleNet V2                     30.64           11.68
Sasank Chilamkurthy's avatar
Sasank Chilamkurthy committed
95
96
97
98
99
100
101
102
103
================================  =============   =============


.. _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
104
.. _GoogLeNet: https://arxiv.org/abs/1409.4842
Bar's avatar
Bar committed
105
.. _ShuffleNet: https://arxiv.org/abs/1807.11164
Sasank Chilamkurthy's avatar
Sasank Chilamkurthy committed
106
107
108

.. currentmodule:: torchvision.models

Sasank Chilamkurthy's avatar
Sasank Chilamkurthy committed
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
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
142
DenseNet
Sasank Chilamkurthy's avatar
Sasank Chilamkurthy committed
143
144
145
146
147
148
149
150
151
152
153
154
---------

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

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

.. autofunction:: inception_v3

155
156
157
158
159
GoogLeNet
------------

.. autofunction:: googlenet

Bar's avatar
Bar committed
160
161
162
163
164
ShuffleNet v2
-------------

.. autofunction:: shufflenet