Add registration mechanism for models (#6333)

* Model registration mechanism.

* Add overwrite options to the dataset prototype registration mechanism.

* Adding example models.

* Fix module filtering

* Fix linter

* Fix docs

* Make name optional if same as model builder

* Apply updates from code-review.

* fix minor bug

* Adding getter for model weight enum

* Support both strings and callables on get_model_weight.

* linter fixes

* Fixing mypy.

* Renaming `get_model_weight` to `get_model_weights`

* Registering all classification models.

* Registering all video models.

* Registering all detection models.

* Registering all optical flow models.

* Fixing mypy.

* Registering all segmentation models.

* Registering all quantization models.

* Fixing linter

* Registering all prototype depth perception models.

* Adding tests and updating existing tests.

* Fix linters

* Fix tests.

* Add beta annotation on docs.

* Fix tests.

* Apply changes from code-review.

* Adding documentation.

* Fix docs.

Author: datumbox

docs/source/models.rst

@@ -120,6 +120,46 @@ behavior, such as batch normalization. To switch between these modes, use
     # Set model to eval mode
     model.eval()
 
+Model Registration Mechanism
+----------------------------
+
+.. betastatus:: registration mechanism
+
+As of v0.14, TorchVision offers a new model registration mechanism which allows retreaving models
+and weights by their names. Here are a few examples on how to use them:
+
+.. code:: python
+
+    # List available models
+    all_models = list_models()
+    classification_models = list_models(module=torchvision.models)
+
+    # Initialize models
+    m1 = get_model("mobilenet_v3_large", weights=None)
+    m2 = get_model("quantized_mobilenet_v3_large", weights="DEFAULT")
+
+    # Fetch weights
+    weights = get_weight("MobileNet_V3_Large_QuantizedWeights.DEFAULT")
+    assert weights == MobileNet_V3_Large_QuantizedWeights.DEFAULT
+
+    weights_enum = get_model_weights("quantized_mobilenet_v3_large")
+    assert weights_enum == MobileNet_V3_Large_QuantizedWeights
+
+    weights_enum2 = get_model_weights(torchvision.models.quantization.mobilenet_v3_large)
+    assert weights_enum == weights_enum2
+
+Here are the available public methods of the model registration mechanism:
+
+.. currentmodule:: torchvision.models
+.. autosummary::
+    :toctree: generated/
+    :template: function.rst
+
+    get_model
+    get_model_weights
+    get_weight
+    list_models
+
 Using models from Hub
 ---------------------
 

test/test_backbone_utils.py

@@ -11,15 +11,6 @@
 from torchvision.models.feature_extraction import create_feature_extractor, get_graph_node_names
 
 
-def get_available_models():
-    # TODO add a registration mechanism to torchvision.models
-    return [
-        k
-        for k, v in models.__dict__.items()
-        if callable(v) and k[0].lower() == k[0] and k[0] != "_" and k != "get_weight"
-    ]
-
-
 @pytest.mark.parametrize("backbone_name", ("resnet18", "resnet50"))
 def test_resnet_fpn_backbone(backbone_name):
     x = torch.rand(1, 3, 300, 300, dtype=torch.float32, device="cpu")
@@ -135,10 +126,10 @@ def _get_return_nodes(self, model):
         eval_nodes = [n for n in eval_nodes if not any(x in n for x in exclude_nodes_filter)]
         return random.sample(train_nodes, 10), random.sample(eval_nodes, 10)
 
-    @pytest.mark.parametrize("model_name", get_available_models())
+    @pytest.mark.parametrize("model_name", models.list_models(models))
     def test_build_fx_feature_extractor(self, model_name):
         set_rng_seed(0)
-        model = models.__dict__[model_name](**self.model_defaults).eval()
+        model = models.get_model(model_name, **self.model_defaults).eval()
         train_return_nodes, eval_return_nodes = self._get_return_nodes(model)
         # Check that it works with both a list and dict for return nodes
         self._create_feature_extractor(
@@ -172,9 +163,9 @@ def test_node_name_conventions(self):
         train_nodes, _ = get_graph_node_names(model)
         assert all(a == b for a, b in zip(train_nodes, test_module_nodes))
 
-    @pytest.mark.parametrize("model_name", get_available_models())
+    @pytest.mark.parametrize("model_name", models.list_models(models))
     def test_forward_backward(self, model_name):
-        model = models.__dict__[model_name](**self.model_defaults).train()
+        model = models.get_model(model_name, **self.model_defaults).train()
         train_return_nodes, eval_return_nodes = self._get_return_nodes(model)
         model = self._create_feature_extractor(
             model, train_return_nodes=train_return_nodes, eval_return_nodes=eval_return_nodes
@@ -211,10 +202,10 @@ def test_feature_extraction_methods_equivalence(self):
         for k in ilg_out.keys():
             assert ilg_out[k].equal(fgn_out[k])
 
-    @pytest.mark.parametrize("model_name", get_available_models())
+    @pytest.mark.parametrize("model_name", models.list_models(models))
     def test_jit_forward_backward(self, model_name):
         set_rng_seed(0)
-        model = models.__dict__[model_name](**self.model_defaults).train()
+        model = models.get_model(model_name, **self.model_defaults).train()
         train_return_nodes, eval_return_nodes = self._get_return_nodes(model)
         model = self._create_feature_extractor(
             model, train_return_nodes=train_return_nodes, eval_return_nodes=eval_return_nodes