Refactor Segmentation models #4646

datumbox · 2021-10-18T15:26:18Z

This PR refactors the entire segmentation package to align with the conventions followed by classification, detection and quantization.

Here are the details:

The public model builders located at models.segmentation.segmentation now move to their respective model packages. For example, fcn_resnet50() is now defined in fcn.py.
The private builder methods of each Model are refactored, so that they can be reused on prototype with minimal copy-pasting.
We expose publicly pretrained_backbone instead via kwargs.

Accuracy remains the same for all pre-trained models. No modifications on tests required because the changes are fully-BC. This PR enables the work on #4611.

…tation.

datumbox

Highlighting few important implementation details:

datumbox · 2021-10-18T19:28:06Z

torchvision/models/segmentation/deeplabv3.py

+
+
+def _deeplabv3_resnet(
+    backbone: resnet.ResNet,


We make all private model builders accept the pre-initialized backbones instead of passing the backbone_name. This allows us to reuse the methods on the multi-pretrained weights project.

datumbox · 2021-10-18T19:29:09Z

torchvision/models/segmentation/deeplabv3.py

+
+    aux_classifier = FCNHead(1024, num_classes) if aux else None
+    classifier = DeepLabHead(2048, num_classes)
+    return DeepLabV3(backbone, classifier, aux_classifier)


Overall we simplify the code by splitting the previous massive _segm_model() method.

torchvision/models/segmentation/deeplabv3.py

datumbox · 2021-10-18T19:30:59Z

torchvision/models/segmentation/deeplabv3.py

+    progress: bool = True,
+    num_classes: int = 21,
+    aux_loss: bool = False,
+    pretrained_backbone: bool = True,


Instead of **kwargs we directly expose the pretrained_backbone parameter publicly.

datumbox · 2021-10-18T19:32:29Z

torchvision/models/segmentation/segmentation.py


-from torch import nn
+from . import *  # noqa: F401, F403


Import all methods, classes, etc just in-case someone was using the package:
torchvision.models.segmentation.segmentation.

Do you mind just adding this as a comment so we don't "clean" it later?

datumbox · 2021-10-18T19:33:14Z

torchvision/models/segmentation/segmentation.py

-        _load_weights(model, "lraspp", backbone_name, progress)
-
-    return model
+warnings.warn(


Throw a deprecation warning to use the parent package. Aka torchvision.models.segmentation.

NicolasHug

Thanks for the PR and for writing comments about the important parts.
This looks great, I just have one concern regarding the change of default for aux_loss, LMK what you think

NicolasHug · 2021-10-19T08:15:29Z

torchvision/models/segmentation/segmentation.py


-from torch import nn
+from . import *  # noqa: F401, F403


Do you mind just adding this as a comment so we don't "clean" it later?

torchvision/models/segmentation/deeplabv3.py

NicolasHug

Thanks @datumbox ! LGTM modulo some quick verification of the training perf as you mentioned already

(and the type checks apparently)

fmassa

Looks great, thanks for cleaning this up!

fmassa · 2021-10-19T11:10:36Z

torchvision/models/segmentation/deeplabv3.py

+    if pretrained:
+        arch = "deeplabv3_resnet50_coco"
+        _load_weights(arch, model, model_urls.get(arch, None), progress)


nit: I suppose you've left this here (instead of putting it in _deeplabv3_resnet because it will be more aligned with your changes to the new weights?

Same for the backbone retrieval code

Yes exactly. I was going back and forth about this. If I were to put it in the builder methods, I would have to copy-paste the whole thing...

I think an additional final clean up wold be necessary prior moving the prototype work to main and there we would be able to move things around. This is a great candidate for such clean up.

datumbox · 2021-10-19T11:40:55Z

Thanks for the reviews. I confirmed that all models maintain the right model validation statistics when compared to the main branch:

torchrun --nproc_per_node=2 train.py --test-only --pretrained --model fcn_resnet50
global correct: 91.4
mean IoU: 60.5

torchrun --nproc_per_node=2 train.py --test-only --pretrained --model fcn_resnet101
global correct: 91.9
mean IoU: 63.7

torchrun --nproc_per_node=2 train.py --test-only --pretrained --model deeplabv3_resnet50
global correct: 92.4
mean IoU: 66.4

torchrun --nproc_per_node=2 train.py --test-only --pretrained --model deeplabv3_resnet101
global correct: 92.4
mean IoU: 67.4

torchrun --nproc_per_node=2 train.py --test-only --pretrained --model deeplabv3_mobilenet_v3_large
global correct: 91.2
mean IoU: 60.3

torchrun --nproc_per_node=2 train.py --test-only --pretrained --model lraspp_mobilenet_v3_large
global correct: 91.2
mean IoU: 57.9

I'll merge once the CI passes.

Summary: * Move FCN methods to itsown package. * Fix lint. * Move LRASPP methods to their own package. * Move DeepLabV3 methods to their own package. * Adding deprecation warning for torchvision.models.segmentation.segmentation. * Refactoring deeplab. * Setting aux default to false. * Fixing imports. * Passing backbones instead of backbone names to builders. * Fixing mypy * Addressing review comments. * Correcting typing. * Restoring special handling for references. Reviewed By: datumbox Differential Revision: D31898217 fbshipit-source-id: 187bb6cc50ea14f121a907aaac4770f17696d6af

* Move FCN methods to itsown package. * Fix lint. * Move LRASPP methods to their own package. * Move DeepLabV3 methods to their own package. * Adding deprecation warning for torchvision.models.segmentation.segmentation. * Refactoring deeplab. * Setting aux default to false. * Fixing imports. * Passing backbones instead of backbone names to builders. * Fixing mypy * Addressing review comments. * Correcting typing. * Restoring special handling for references.

Move FCN methods to itsown package.

5e5a68c

pytorch-probot bot added the ciflow/default label Oct 18, 2021

facebook-github-bot added the cla signed label Oct 18, 2021

datumbox marked this pull request as draft October 18, 2021 15:26

datumbox added code quality module: models labels Oct 18, 2021

datumbox added 4 commits October 18, 2021 16:30

Fix lint.

93cdcf6

Move LRASPP methods to their own package.

7213a9d

Move DeepLabV3 methods to their own package.

bee91a1

Adding deprecation warning for torchvision.models.segmentation.segmen…

c4c5ac4

…tation.

datumbox force-pushed the models/segmentation_refactoring branch from 679b871 to c4c5ac4 Compare October 18, 2021 17:50

datumbox added 3 commits October 18, 2021 19:17

Refactoring deeplab.

c341ed3

Setting aux default to false.

43c677e

Fixing imports.

f534046

datumbox changed the title ~~[WIP] Refactor Segmentation models~~ Refactor Segmentation models Oct 18, 2021

datumbox marked this pull request as ready for review October 18, 2021 18:32

Passing backbones instead of backbone names to builders.

de1d2ad

datumbox force-pushed the models/segmentation_refactoring branch from 150ab32 to de1d2ad Compare October 18, 2021 19:02

datumbox commented Oct 18, 2021

View reviewed changes

Fixing mypy

ede7980

datumbox force-pushed the models/segmentation_refactoring branch from 4e19c1d to ede7980 Compare October 18, 2021 19:38

Merge branch 'main' into models/segmentation_refactoring

198683a

NicolasHug reviewed Oct 19, 2021

View reviewed changes

Addressing review comments.

ad34b89

NicolasHug approved these changes Oct 19, 2021

View reviewed changes

datumbox and others added 2 commits October 19, 2021 10:31

Correcting typing.

7a96c40

Merge branch 'main' into models/segmentation_refactoring

39392d7

fmassa approved these changes Oct 19, 2021

View reviewed changes

Restoring special handling for references.

7fef773

Merge branch 'main' into models/segmentation_refactoring

49816fc

datumbox merged commit a1d6b31 into pytorch:main Oct 19, 2021

datumbox deleted the models/segmentation_refactoring branch October 19, 2021 12:40

datumbox mentioned this pull request Oct 19, 2021

Refactor the backbone builders of detection #4656

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Refactor Segmentation models #4646

Refactor Segmentation models #4646

datumbox commented Oct 18, 2021 •

edited

Loading

datumbox left a comment

datumbox Oct 18, 2021 •

edited

Loading

datumbox Oct 18, 2021

datumbox Oct 18, 2021

datumbox Oct 18, 2021

NicolasHug Oct 19, 2021

datumbox Oct 18, 2021

NicolasHug left a comment

NicolasHug Oct 19, 2021

NicolasHug left a comment •

edited

Loading

fmassa left a comment

fmassa Oct 19, 2021

datumbox Oct 19, 2021

datumbox commented Oct 19, 2021

Refactor Segmentation models #4646

Refactor Segmentation models #4646

Conversation

datumbox commented Oct 18, 2021 • edited Loading

datumbox left a comment

Choose a reason for hiding this comment

datumbox Oct 18, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

NicolasHug left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

NicolasHug left a comment • edited Loading

Choose a reason for hiding this comment

fmassa left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

datumbox commented Oct 19, 2021

datumbox commented Oct 18, 2021 •

edited

Loading

datumbox Oct 18, 2021 •

edited

Loading

NicolasHug left a comment •

edited

Loading