models_new.rst 15.1 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
.. _models_new:

Models and pre-trained weights - New
####################################

The ``torchvision.models`` subpackage contains definitions of models for addressing
different tasks, including: image classification, pixelwise semantic
segmentation, object detection, instance segmentation, person
keypoint detection, video classification, and optical flow.

11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
General information on pre-trained weights
==========================================

TorchVision offers pre-trained weights for every provided architecture, using
the PyTorch :mod:`torch.hub`. Instancing a pre-trained model will download its
weights to a cache directory. This directory can be set using the `TORCH_HOME`
environment variable. See :func:`torch.hub.load_state_dict_from_url` for details.

.. note::

    The pre-trained models provided in this library may have their own licenses or
    terms and conditions derived from the dataset used for training. It is your
    responsibility to determine whether you have permission to use the models for
    your use case.

26
.. note ::
27
28
29
30
31
32
33
34
35
36
37
    Backward compatibility is guaranteed for loading a serialized
    ``state_dict`` to the model created using old PyTorch version.
    On the contrary, loading entire saved models or serialized
    ``ScriptModules`` (serialized using older versions of PyTorch)
    may not preserve the historic behaviour. Refer to the following
    `documentation
    <https://pytorch.org/docs/stable/notes/serialization.html#id6>`_


Initializing pre-trained models
-------------------------------
38

39
As of v0.13, TorchVision offers a new `Multi-weight support API
40
41
<https://pytorch.org/blog/introducing-torchvision-new-multi-weight-support-api/>`_
for loading different weights to the existing model builder methods:
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60

.. code:: python

    from torchvision.models import resnet50, ResNet50_Weights

    # Old weights with accuracy 76.130%
    resnet50(weights=ResNet50_Weights.IMAGENET1K_V1)

    # New weights with accuracy 80.858%
    resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)

    # Best available weights (currently alias for IMAGENET1K_V2)
    # Note that these weights may change across versions
    resnet50(weights=ResNet50_Weights.DEFAULT)

    # Strings are also supported
    resnet50(weights="IMAGENET1K_V2")

    # No weights - random initialization
61
    resnet50(weights=None)
62
63
64
65
66
67
68
69
70
71


Migrating to the new API is very straightforward. The following method calls between the 2 APIs are all equivalent:

.. code:: python

    from torchvision.models import resnet50, ResNet50_Weights

    # Using pretrained weights:
    resnet50(weights=ResNet50_Weights.IMAGENET1K_V1)
72
    resnet50(weights="IMAGENET1K_V1")
73
74
75
76
77
    resnet50(pretrained=True)  # deprecated
    resnet50(True)  # deprecated

    # Using no weights:
    resnet50(weights=None)
78
    resnet50()
79
80
81
82
83
    resnet50(pretrained=False)  # deprecated
    resnet50(False)  # deprecated

Note that the ``pretrained`` parameter is now deprecated, using it will emit warnings and will be removed on v0.15.

84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
Using the pre-trained models
----------------------------

Before using the pre-trained models, one must preprocess the image
(resize with right resolution/interpolation, apply inference transforms,
rescale the values etc). There is no standard way to do this as it depends on
how a given model was trained. It can vary across model families, variants or
even weight versions. Using the correct preprocessing method is critical and
failing to do so may lead to decreased accuracy or incorrect outputs.

All the necessary information for the inference transforms of each pre-trained
model is provided on its weights documentation. To simplify inference, TorchVision
bundles the necessary preprocessing transforms into each model weight. These are
accessible via the ``weight.transforms`` attribute:

.. code:: python

    # Initialize the Weight Transforms
    weights = ResNet50_Weights.DEFAULT
    preprocess = weights.transforms()

    # Apply it to the input image
    img_transformed = preprocess(img)


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.

.. code:: python

    # Initialize model
    weights = ResNet50_Weights.DEFAULT
    model = resnet50(weights=weights)

    # Set model to eval mode
    model.eval()

123
124
125
126
127
128
129
130
131
132
133
134

Classification
==============

.. currentmodule:: torchvision.models

The following classification models are available, with or without pre-trained
weights:

.. toctree::
   :maxdepth: 1

135
   models/alexnet
Hu Ye's avatar
Hu Ye committed
136
   models/convnext
137
   models/densenet
138
   models/efficientnet
139
   models/efficientnetv2
140
   models/googlenet
Aditya Oke's avatar
Aditya Oke committed
141
   models/inception
Joao Gomes's avatar
Joao Gomes committed
142
   models/mnasnet
143
   models/mobilenetv2
144
   models/mobilenetv3
145
   models/regnet
146
   models/resnet
147
   models/resnext
148
   models/shufflenetv2
Nicolas Hug's avatar
Nicolas Hug committed
149
   models/squeezenet
150
   models/swin_transformer
151
   models/vgg
152
   models/vision_transformer
153
   models/wide_resnet
154

155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
|

Here is an example of how to use the pre-trained image classification models:

.. code:: python

    from torchvision.io import read_image
    from torchvision.models import resnet50, ResNet50_Weights

    img = read_image("test/assets/encode_jpeg/grace_hopper_517x606.jpg")

    # Step 1: Initialize model with the best available weights
    weights = ResNet50_Weights.DEFAULT
    model = resnet50(weights=weights)
    model.eval()

    # Step 2: Initialize the inference transforms
    preprocess = weights.transforms()

    # Step 3: Apply inference preprocessing transforms
    batch = preprocess(img).unsqueeze(0)

    # Step 4: Use the model and print the predicted category
    prediction = model(batch).squeeze(0).softmax(0)
    class_id = prediction.argmax().item()
    score = prediction[class_id].item()
    category_name = weights.meta["categories"][class_id]
    print(f"{category_name}: {100 * score:.1f}%")
183

184
185
The classes of the pre-trained model outputs can be found at ``weights.meta["categories"]``.

186
187
188
Table of all available classification weights
---------------------------------------------

189
Accuracies are reported on ImageNet-1K using single crops:
190
191
192

.. include:: generated/classification_table.rst

193
194
195
196
197
Quantized models
----------------

.. currentmodule:: torchvision.models.quantization

198
The following architectures provide support for INT8 quantized models, with or without
199
200
201
202
203
204
pre-trained weights:

.. toctree::
   :maxdepth: 1

   models/googlenet_quant
205
   models/inception_quant
206
   models/mobilenetv2_quant
207
   models/mobilenetv3_quant
208
   models/resnet_quant
209
   models/shufflenetv2_quant
210

211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
|

Here is an example of how to use the pre-trained quantized image classification models:

.. code:: python

    from torchvision.io import read_image
    from torchvision.models.quantization import resnet50, ResNet50_QuantizedWeights

    img = read_image("test/assets/encode_jpeg/grace_hopper_517x606.jpg")

    # Step 1: Initialize model with the best available weights
    weights = ResNet50_QuantizedWeights.DEFAULT
    model = resnet50(weights=weights, quantize=True)
    model.eval()

    # Step 2: Initialize the inference transforms
    preprocess = weights.transforms()

    # Step 3: Apply inference preprocessing transforms
    batch = preprocess(img).unsqueeze(0)

    # Step 4: Use the model and print the predicted category
    prediction = model(batch).squeeze(0).softmax(0)
    class_id = prediction.argmax().item()
    score = prediction[class_id].item()
    category_name = weights.meta["categories"][class_id]
    print(f"{category_name}: {100 * score}%")

240
241
The classes of the pre-trained model outputs can be found at ``weights.meta["categories"]``.

242
243
244
245

Table of all available quantized classification weights
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

246
Accuracies are reported on ImageNet-1K using single crops:
247
248
249

.. include:: generated/classification_quant_table.rst

250
251
252
253
254
255
256
257
258
259
260
261
Semantic Segmentation
=====================

.. currentmodule:: torchvision.models.segmentation

The following semantic segmentation models are available, with or without
pre-trained weights:

.. toctree::
   :maxdepth: 1

   models/deeplabv3
262
   models/fcn
Aditya Oke's avatar
Aditya Oke committed
263
   models/lraspp
264

265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
|

Here is an example of how to use the pre-trained semantic segmentation models:

.. code:: python

    from torchvision.io.image import read_image
    from torchvision.models.segmentation import fcn_resnet50, FCN_ResNet50_Weights
    from torchvision.transforms.functional import to_pil_image

    img = read_image("gallery/assets/dog1.jpg")

    # Step 1: Initialize model with the best available weights
    weights = FCN_ResNet50_Weights.DEFAULT
    model = fcn_resnet50(weights=weights)
    model.eval()

    # Step 2: Initialize the inference transforms
    preprocess = weights.transforms()

    # Step 3: Apply inference preprocessing transforms
    batch = preprocess(img).unsqueeze(0)

    # Step 4: Use the model and visualize the prediction
    prediction = model(batch)["out"]
    normalized_masks = prediction.softmax(dim=1)
    class_to_idx = {cls: idx for (idx, cls) in enumerate(weights.meta["categories"])}
    mask = normalized_masks[0, class_to_idx["dog"]]
    to_pil_image(mask).show()

295
296
297
The classes of the pre-trained model outputs can be found at ``weights.meta["categories"]``.
The output format of the models is illustrated in :ref:`semantic_seg_output`.

298

299
300
301
Table of all available semantic segmentation weights
----------------------------------------------------

302
All models are evaluated a subset of COCO val2017, on the 20 categories that are present in the Pascal VOC dataset:
303
304
305
306

.. include:: generated/segmentation_table.rst


307

308
309
310
Object Detection, Instance Segmentation and Person Keypoint Detection
=====================================================================

311
312
313
314
315
The pre-trained models for detection, instance segmentation and
keypoint detection are initialized with the classification models
in torchvision. The models expect a list of ``Tensor[C, H, W]``.
Check the constructor of the models for more information.

316
Object Detection
317
----------------
318

319
320
.. currentmodule:: torchvision.models.detection

321
The following object detection models are available, with or without pre-trained
322
323
324
325
326
weights:

.. toctree::
   :maxdepth: 1

327
   models/faster_rcnn
Hu Ye's avatar
Hu Ye committed
328
329
   models/fcos
   models/retinanet
330
   models/ssd
331
   models/ssdlite
332

333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
|

Here is an example of how to use the pre-trained object detection models:

.. code:: python


    from torchvision.io.image import read_image
    from torchvision.models.detection import fasterrcnn_resnet50_fpn_v2, FasterRCNN_ResNet50_FPN_V2_Weights
    from torchvision.utils import draw_bounding_boxes
    from torchvision.transforms.functional import to_pil_image

    img = read_image("test/assets/encode_jpeg/grace_hopper_517x606.jpg")

    # Step 1: Initialize model with the best available weights
    weights = FasterRCNN_ResNet50_FPN_V2_Weights.DEFAULT
    model = fasterrcnn_resnet50_fpn_v2(weights=weights, box_score_thresh=0.9)
    model.eval()

    # Step 2: Initialize the inference transforms
    preprocess = weights.transforms()

    # Step 3: Apply inference preprocessing transforms
    batch = [preprocess(img)]

    # Step 4: Use the model and visualize the prediction
    prediction = model(batch)[0]
    labels = [weights.meta["categories"][i] for i in prediction["labels"]]
    box = draw_bounding_boxes(img, boxes=prediction["boxes"],
                              labels=labels,
                              colors="red",
                              width=4, font_size=30)
    im = to_pil_image(box.detach())
    im.show()

368
369
370
The classes of the pre-trained model outputs can be found at ``weights.meta["categories"]``.
For details on how to plot the bounding boxes of the models, you may refer to :ref:`instance_seg_output`.

371
372
Table of all available Object detection weights
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
373

374
Box MAPs are reported on COCO val2017:
375
376

.. include:: generated/detection_table.rst
377

378
379
380
381
382
383
384
385
386
387
388
389
390
Instance Segmentation
---------------------

.. currentmodule:: torchvision.models.detection

The following instance segmentation models are available, with or without pre-trained
weights:

.. toctree::
   :maxdepth: 1

   models/mask_rcnn

391
392
393
394
395
|


For details on how to plot the masks of the models, you may refer to :ref:`instance_seg_output`.

396
397
398
Table of all available Instance segmentation weights
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

399
Box and Mask MAPs are reported on COCO val2017:
400
401

.. include:: generated/instance_segmentation_table.rst
402

403
404
Keypoint Detection
------------------
405
406
407

.. currentmodule:: torchvision.models.detection

408
The following person keypoint detection models are available, with or without
409
410
411
412
413
414
415
pre-trained weights:

.. toctree::
   :maxdepth: 1

   models/keypoint_rcnn

416
417
418
419
420
|

The classes of the pre-trained model outputs can be found at ``weights.meta["keypoint_names"]``.
For details on how to plot the bounding boxes of the models, you may refer to :ref:`keypoint_output`.

421
Table of all available Keypoint detection weights
422
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
423

424
Box and Keypoint MAPs are reported on COCO val2017:
425
426
427
428

.. include:: generated/detection_keypoint_table.rst


429
430
431
432
433
434
435
436
437
438
439
440
441
Video Classification
====================

.. currentmodule:: torchvision.models.video

The following video classification models are available, with or without
pre-trained weights:

.. toctree::
   :maxdepth: 1

   models/video_resnet

442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
|

Here is an example of how to use the pre-trained video classification models:

.. code:: python


    from torchvision.io.video import read_video
    from torchvision.models.video import r3d_18, R3D_18_Weights

    vid, _, _ = read_video("test/assets/videos/v_SoccerJuggling_g23_c01.avi")
    vid = vid[:32]  # optionally shorten duration

    # Step 1: Initialize model with the best available weights
    weights = R3D_18_Weights.DEFAULT
    model = r3d_18(weights=weights)
    model.eval()

    # Step 2: Initialize the inference transforms
    preprocess = weights.transforms()

    # Step 3: Apply inference preprocessing transforms
    batch = preprocess(vid).unsqueeze(0)

    # Step 4: Use the model and print the predicted category
    prediction = model(batch).squeeze(0).softmax(0)
    label = prediction.argmax().item()
    score = prediction[label].item()
    category_name = weights.meta["categories"][label]
    print(f"{category_name}: {100 * score}%")

473
474
The classes of the pre-trained model outputs can be found at ``weights.meta["categories"]``.

475

476
477
478
Table of all available video classification weights
---------------------------------------------------

479
Accuracies are reported on Kinetics-400 using single crops for clip length 16:
480
481

.. include:: generated/video_table.rst
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501

Using models from Hub
=====================

Most pre-trained models can be accessed directly via PyTorch Hub without having TorchVision installed:

.. code:: python

    import torch

    # Option 1: passing weights param as string
    model = torch.hub.load("pytorch/vision", "resnet50", weights="IMAGENET1K_V2")

    # Option 2: passing weights param as enum
    weights = torch.hub.load("pytorch/vision", "get_weight", weights="ResNet50_Weights.IMAGENET1K_V2")
    model = torch.hub.load("pytorch/vision", "resnet50", weights=weights)

The only exception to the above are the detection models included on
:mod:`torchvision.models.detection`. These models require TorchVision
to be installed because they depend on custom C++ operators.