Merge branch 'main' into completely-remove-requests

Author: prabhat00155

test/test_ops.py

@@ -7,12 +7,54 @@
 import numpy as np
 import pytest
 import torch
+import torch.fx
 from common_utils import needs_cuda, cpu_and_gpu, assert_equal
 from PIL import Image
 from torch import nn, Tensor
 from torch.autograd import gradcheck
 from torch.nn.modules.utils import _pair
 from torchvision import models, ops
+from torchvision.models.feature_extraction import get_graph_node_names
+
+
+class RoIOpTesterModuleWrapper(nn.Module):
+    def __init__(self, obj):
+        super().__init__()
+        self.layer = obj
+        self.n_inputs = 2
+
+    def forward(self, a, b):
+        self.layer(a, b)
+
+
+class MultiScaleRoIAlignModuleWrapper(nn.Module):
+    def __init__(self, obj):
+        super().__init__()
+        self.layer = obj
+        self.n_inputs = 3
+
+    def forward(self, a, b, c):
+        self.layer(a, b, c)
+
+
+class DeformConvModuleWrapper(nn.Module):
+    def __init__(self, obj):
+        super().__init__()
+        self.layer = obj
+        self.n_inputs = 3
+
+    def forward(self, a, b, c):
+        self.layer(a, b, c)
+
+
+class StochasticDepthWrapper(nn.Module):
+    def __init__(self, obj):
+        super().__init__()
+        self.layer = obj
+        self.n_inputs = 1
+
+    def forward(self, a):
+        self.layer(a)
 
 
 class RoIOpTester(ABC):
@@ -46,6 +88,15 @@ def test_forward(self, device, contiguous, x_dtype=None, rois_dtype=None, **kwar
         tol = 1e-3 if (x_dtype is torch.half or rois_dtype is torch.half) else 1e-5
         torch.testing.assert_close(gt_y.to(y), y, rtol=tol, atol=tol)
 
+    @pytest.mark.parametrize("device", cpu_and_gpu())
+    def test_is_leaf_node(self, device):
+        op_obj = self.make_obj(wrap=True).to(device=device)
+        graph_node_names = get_graph_node_names(op_obj)
+
+        assert len(graph_node_names) == 2
+        assert len(graph_node_names[0]) == len(graph_node_names[1])
+        assert len(graph_node_names[0]) == 1 + op_obj.n_inputs
+
     @pytest.mark.parametrize("seed", range(10))
     @pytest.mark.parametrize("device", cpu_and_gpu())
     @pytest.mark.parametrize("contiguous", (True, False))
@@ -91,6 +142,10 @@ def _helper_boxes_shape(self, func):
     def fn(*args, **kwargs):
         pass
 
+    @abstractmethod
+    def make_obj(*args, **kwargs):
+        pass
+
     @abstractmethod
     def get_script_fn(*args, **kwargs):
         pass
@@ -104,6 +159,10 @@ class TestRoiPool(RoIOpTester):
     def fn(self, x, rois, pool_h, pool_w, spatial_scale=1, sampling_ratio=-1, **kwargs):
         return ops.RoIPool((pool_h, pool_w), spatial_scale)(x, rois)
 
+    def make_obj(self, pool_h=5, pool_w=5, spatial_scale=1, wrap=False):
+        obj = ops.RoIPool((pool_h, pool_w), spatial_scale)
+        return RoIOpTesterModuleWrapper(obj) if wrap else obj
+
     def get_script_fn(self, rois, pool_size):
         scriped = torch.jit.script(ops.roi_pool)
         return lambda x: scriped(x, rois, pool_size)
@@ -144,6 +203,10 @@ class TestPSRoIPool(RoIOpTester):
     def fn(self, x, rois, pool_h, pool_w, spatial_scale=1, sampling_ratio=-1, **kwargs):
         return ops.PSRoIPool((pool_h, pool_w), 1)(x, rois)
 
+    def make_obj(self, pool_h=5, pool_w=5, spatial_scale=1, wrap=False):
+        obj = ops.PSRoIPool((pool_h, pool_w), spatial_scale)
+        return RoIOpTesterModuleWrapper(obj) if wrap else obj
+
     def get_script_fn(self, rois, pool_size):
         scriped = torch.jit.script(ops.ps_roi_pool)
         return lambda x: scriped(x, rois, pool_size)
@@ -223,6 +286,12 @@ def fn(self, x, rois, pool_h, pool_w, spatial_scale=1, sampling_ratio=-1, aligne
             (pool_h, pool_w), spatial_scale=spatial_scale, sampling_ratio=sampling_ratio, aligned=aligned
         )(x, rois)
 
+    def make_obj(self, pool_h=5, pool_w=5, spatial_scale=1, sampling_ratio=-1, aligned=False, wrap=False):
+        obj = ops.RoIAlign(
+            (pool_h, pool_w), spatial_scale=spatial_scale, sampling_ratio=sampling_ratio, aligned=aligned
+        )
+        return RoIOpTesterModuleWrapper(obj) if wrap else obj
+
     def get_script_fn(self, rois, pool_size):
         scriped = torch.jit.script(ops.roi_align)
         return lambda x: scriped(x, rois, pool_size)
@@ -374,6 +443,10 @@ class TestPSRoIAlign(RoIOpTester):
     def fn(self, x, rois, pool_h, pool_w, spatial_scale=1, sampling_ratio=-1, **kwargs):
         return ops.PSRoIAlign((pool_h, pool_w), spatial_scale=spatial_scale, sampling_ratio=sampling_ratio)(x, rois)
 
+    def make_obj(self, pool_h=5, pool_w=5, spatial_scale=1, sampling_ratio=-1, wrap=False):
+        obj = ops.PSRoIAlign((pool_h, pool_w), spatial_scale=spatial_scale, sampling_ratio=sampling_ratio)
+        return RoIOpTesterModuleWrapper(obj) if wrap else obj
+
     def get_script_fn(self, rois, pool_size):
         scriped = torch.jit.script(ops.ps_roi_align)
         return lambda x: scriped(x, rois, pool_size)
@@ -422,12 +495,18 @@ def test_boxes_shape(self):
 
 
 class TestMultiScaleRoIAlign:
+    def make_obj(self, fmap_names=None, output_size=(7, 7), sampling_ratio=2, wrap=False):
+        if fmap_names is None:
+            fmap_names = ["0"]
+        obj = ops.poolers.MultiScaleRoIAlign(fmap_names, output_size, sampling_ratio)
+        return MultiScaleRoIAlignModuleWrapper(obj) if wrap else obj
+
     def test_msroialign_repr(self):
         fmap_names = ["0"]
         output_size = (7, 7)
         sampling_ratio = 2
         # Pass mock feature map names
-        t = ops.poolers.MultiScaleRoIAlign(fmap_names, output_size, sampling_ratio)
+        t = self.make_obj(fmap_names, output_size, sampling_ratio, wrap=False)
 
         # Check integrity of object __repr__ attribute
         expected_string = (
@@ -436,6 +515,15 @@ def test_msroialign_repr(self):
         )
         assert repr(t) == expected_string
 
+    @pytest.mark.parametrize("device", cpu_and_gpu())
+    def test_is_leaf_node(self, device):
+        op_obj = self.make_obj(wrap=True).to(device=device)
+        graph_node_names = get_graph_node_names(op_obj)
+
+        assert len(graph_node_names) == 2
+        assert len(graph_node_names[0]) == len(graph_node_names[1])
+        assert len(graph_node_names[0]) == 1 + op_obj.n_inputs
+
 
 class TestNMS:
     def _reference_nms(self, boxes, scores, iou_threshold):
@@ -693,6 +781,21 @@ def get_fn_args(self, device, contiguous, batch_sz, dtype):
 
         return x, weight, offset, mask, bias, stride, pad, dilation
 
+    def make_obj(self, in_channels=6, out_channels=2, kernel_size=(3, 2), groups=2, wrap=False):
+        obj = ops.DeformConv2d(
+            in_channels, out_channels, kernel_size, stride=(2, 1), padding=(1, 0), dilation=(2, 1), groups=groups
+        )
+        return DeformConvModuleWrapper(obj) if wrap else obj
+
+    @pytest.mark.parametrize("device", cpu_and_gpu())
+    def test_is_leaf_node(self, device):
+        op_obj = self.make_obj(wrap=True).to(device=device)
+        graph_node_names = get_graph_node_names(op_obj)
+
+        assert len(graph_node_names) == 2
+        assert len(graph_node_names[0]) == len(graph_node_names[1])
+        assert len(graph_node_names[0]) == 1 + op_obj.n_inputs
+
     @pytest.mark.parametrize("device", cpu_and_gpu())
     @pytest.mark.parametrize("contiguous", (True, False))
     @pytest.mark.parametrize("batch_sz", (0, 33))
@@ -705,9 +808,9 @@ def test_forward(self, device, contiguous, batch_sz, dtype=None):
         groups = 2
         tol = 2e-3 if dtype is torch.half else 1e-5
 
-        layer = ops.DeformConv2d(
-            in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups
-        ).to(device=x.device, dtype=dtype)
+        layer = self.make_obj(in_channels, out_channels, kernel_size, groups, wrap=False).to(
+            device=x.device, dtype=dtype
+        )
         res = layer(x, offset, mask)
 
         weight = layer.weight.data
@@ -1200,6 +1303,20 @@ def test_stochastic_depth(self, seed, mode, p):
         elif p == 1:
             assert out.equal(torch.zeros_like(x))
 
+    def make_obj(self, p, mode, wrap=False):
+        obj = ops.StochasticDepth(p, mode)
+        return StochasticDepthWrapper(obj) if wrap else obj
+
+    @pytest.mark.parametrize("p", (0, 1))
+    @pytest.mark.parametrize("mode", ["batch", "row"])
+    def test_is_leaf_node(self, p, mode):
+        op_obj = self.make_obj(p, mode, wrap=True)
+        graph_node_names = get_graph_node_names(op_obj)
+
+        assert len(graph_node_names) == 2
+        assert len(graph_node_names[0]) == len(graph_node_names[1])
+        assert len(graph_node_names[0]) == 1 + op_obj.n_inputs
+
 
 class TestUtils:
     @pytest.mark.parametrize("norm_layer", [None, nn.BatchNorm2d, nn.LayerNorm])

torchvision/models/feature_extraction.py

@@ -1,11 +1,14 @@
+import inspect
+import math
 import re
 import warnings
 from collections import OrderedDict
 from copy import deepcopy
 from itertools import chain
-from typing import Dict, Callable, List, Union, Optional, Tuple
+from typing import Dict, Callable, List, Union, Optional, Tuple, Any
 
 import torch
+import torchvision
 from torch import fx
 from torch import nn
 from torch.fx.graph_module import _copy_attr
@@ -172,8 +175,19 @@ def _warn_graph_differences(train_tracer: NodePathTracer, eval_tracer: NodePathT
     warnings.warn(msg + suggestion_msg)
 
 
+def _get_leaf_modules_for_ops() -> List[type]:
+    members = inspect.getmembers(torchvision.ops)
+    result = []
+    for _, obj in members:
+        if inspect.isclass(obj) and issubclass(obj, torch.nn.Module):
+            result.append(obj)
+    return result
+
+
 def get_graph_node_names(
-    model: nn.Module, tracer_kwargs: Dict = {}, suppress_diff_warning: bool = False
+    model: nn.Module,
+    tracer_kwargs: Optional[Dict[str, Any]] = None,
+    suppress_diff_warning: bool = False,
 ) -> Tuple[List[str], List[str]]:
     """
     Dev utility to return node names in order of execution. See note on node
@@ -198,6 +212,7 @@ def get_graph_node_names(
         tracer_kwargs (dict, optional): a dictionary of keywork arguments for
             ``NodePathTracer`` (they are eventually passed onto
             `torch.fx.Tracer <https://pytorch.org/docs/stable/fx.html#torch.fx.Tracer>`_).
+            By default it will be set to wrap and make leaf nodes all torchvision ops.
         suppress_diff_warning (bool, optional): whether to suppress a warning
             when there are discrepancies between the train and eval version of
             the graph. Defaults to False.
@@ -211,6 +226,14 @@ def get_graph_node_names(
         >>> model = torchvision.models.resnet18()
         >>> train_nodes, eval_nodes = get_graph_node_names(model)
     """
+    if tracer_kwargs is None:
+        tracer_kwargs = {
+            "autowrap_modules": (
+                math,
+                torchvision.ops,
+            ),
+            "leaf_modules": _get_leaf_modules_for_ops(),
+        }
     is_training = model.training
     train_tracer = NodePathTracer(**tracer_kwargs)
     train_tracer.trace(model.train())
@@ -294,7 +317,7 @@ def create_feature_extractor(
     return_nodes: Optional[Union[List[str], Dict[str, str]]] = None,
     train_return_nodes: Optional[Union[List[str], Dict[str, str]]] = None,
     eval_return_nodes: Optional[Union[List[str], Dict[str, str]]] = None,
-    tracer_kwargs: Dict = {},
+    tracer_kwargs: Optional[Dict[str, Any]] = None,
     suppress_diff_warning: bool = False,
 ) -> fx.GraphModule:
     """
@@ -353,6 +376,7 @@ def create_feature_extractor(
         tracer_kwargs (dict, optional): a dictionary of keywork arguments for
             ``NodePathTracer`` (which passes them onto it's parent class
             `torch.fx.Tracer <https://pytorch.org/docs/stable/fx.html#torch.fx.Tracer>`_).
+            By default it will be set to wrap and make leaf nodes all torchvision ops.
         suppress_diff_warning (bool, optional): whether to suppress a warning
             when there are discrepancies between the train and eval version of
             the graph. Defaults to False.
@@ -397,6 +421,14 @@ def create_feature_extractor(
         >>>                    'autowrap_functions': [leaf_function]})
 
     """
+    if tracer_kwargs is None:
+        tracer_kwargs = {
+            "autowrap_modules": (
+                math,
+                torchvision.ops,
+            ),
+            "leaf_modules": _get_leaf_modules_for_ops(),
+        }
     is_training = model.training
 
     assert any(