From 527eb044efef2112914f44a23af02485ebe122ba Mon Sep 17 00:00:00 2001 From: Joao Gomes Date: Thu, 17 Mar 2022 14:31:32 +0000 Subject: [PATCH 1/7] set default tracer kwargs always --- torchvision/models/feature_extraction.py | 67 +++++++++++++++++------- 1 file changed, 49 insertions(+), 18 deletions(-) diff --git a/torchvision/models/feature_extraction.py b/torchvision/models/feature_extraction.py index a6c26913093..80d15c7ecdd 100644 --- a/torchvision/models/feature_extraction.py +++ b/torchvision/models/feature_extraction.py @@ -184,6 +184,26 @@ def _get_leaf_modules_for_ops() -> List[type]: return result +def set_default_tracer_kargs(original_tr_kwargs: Optional[Dict[str, Any]]) -> Dict[str, Any]: + default_autowrap_modules = (math, torchvision.ops) + default_leaf_modules = _get_leaf_modules_for_ops() + if original_tr_kwargs is None: + result_tracer_kwargs = {} + else: + result_tracer_kwargs = original_tr_kwargs + if "autowrap_modules" in result_tracer_kwargs: + result_tracer_kwargs["autowrap_modules"] = tuple( + set(result_tracer_kwargs["autowrap_modules"] + default_autowrap_modules) + ) + else: + result_tracer_kwargs["autowrap_modules"] = default_autowrap_modules + if "leaf_modules" in result_tracer_kwargs: + result_tracer_kwargs["leaf_modules"] = list(set(result_tracer_kwargs["leaf_modules"] + default_leaf_modules)) + else: + result_tracer_kwargs["leaf_modules"] = default_leaf_modules + return result_tracer_kwargs + + def get_graph_node_names( model: nn.Module, tracer_kwargs: Optional[Dict[str, Any]] = None, @@ -212,7 +232,20 @@ def get_graph_node_names( tracer_kwargs (dict, optional): a dictionary of keywork arguments for ``NodePathTracer`` (they are eventually passed onto `torch.fx.Tracer `_). - By default it will be set to wrap and make leaf nodes all torchvision ops. + By default it will be set to wrap and make leaf nodes all torchvision ops, i.e. tracer_kwargs + will be set to + + tracer_kwargs = { + "autowrap_modules": ( + math, + torchvision.ops, + ), + "leaf_modules": _get_leaf_modules_for_ops(), + } + + WARNING: In case the user provides tracer_kwargs, above default arguments will be apended to the user + provided dictionary. + 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. @@ -226,14 +259,7 @@ 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(), - } + tracer_kwargs = set_default_tracer_kargs(tracer_kwargs) is_training = model.training train_tracer = NodePathTracer(**tracer_kwargs) train_tracer.trace(model.train()) @@ -378,7 +404,19 @@ 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 `_). - By default it will be set to wrap and make leaf nodes all torchvision ops. + By default it will be set to wrap and make leaf nodes all torchvision ops, i.e. tracer_kwargs + will be set to + + tracer_kwargs = { + "autowrap_modules": ( + math, + torchvision.ops, + ), + "leaf_modules": _get_leaf_modules_for_ops(), + } + + WARNING: In case the user provides tracer_kwargs, above default arguments will be apended to the user + provided dictionary. 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. @@ -423,14 +461,7 @@ 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(), - } + tracer_kwargs = set_default_tracer_kargs(tracer_kwargs) is_training = model.training if all(arg is None for arg in [return_nodes, train_return_nodes, eval_return_nodes]): From a7319882929652e1e3e60624f99ad0cab8539ec3 Mon Sep 17 00:00:00 2001 From: Joao Gomes Date: Thu, 17 Mar 2022 14:41:00 +0000 Subject: [PATCH 2/7] simplify code --- torchvision/models/feature_extraction.py | 25 +++++++++++------------- 1 file changed, 11 insertions(+), 14 deletions(-) diff --git a/torchvision/models/feature_extraction.py b/torchvision/models/feature_extraction.py index 80d15c7ecdd..24f3732427f 100644 --- a/torchvision/models/feature_extraction.py +++ b/torchvision/models/feature_extraction.py @@ -187,20 +187,17 @@ def _get_leaf_modules_for_ops() -> List[type]: def set_default_tracer_kargs(original_tr_kwargs: Optional[Dict[str, Any]]) -> Dict[str, Any]: default_autowrap_modules = (math, torchvision.ops) default_leaf_modules = _get_leaf_modules_for_ops() - if original_tr_kwargs is None: - result_tracer_kwargs = {} - else: - result_tracer_kwargs = original_tr_kwargs - if "autowrap_modules" in result_tracer_kwargs: - result_tracer_kwargs["autowrap_modules"] = tuple( - set(result_tracer_kwargs["autowrap_modules"] + default_autowrap_modules) - ) - else: - result_tracer_kwargs["autowrap_modules"] = default_autowrap_modules - if "leaf_modules" in result_tracer_kwargs: - result_tracer_kwargs["leaf_modules"] = list(set(result_tracer_kwargs["leaf_modules"] + default_leaf_modules)) - else: - result_tracer_kwargs["leaf_modules"] = default_leaf_modules + result_tracer_kwargs = {} if original_tr_kwargs is None else original_tr_kwargs + result_tracer_kwargs["autowrap_modules"] = ( + tuple(set(result_tracer_kwargs["autowrap_modules"] + default_autowrap_modules)) + if "autowrap_modules" in result_tracer_kwargs + else default_autowrap_modules + ) + result_tracer_kwargs["leaf_modules"] = ( + list(set(result_tracer_kwargs["leaf_modules"] + default_leaf_modules)) + if "leaf_modules" in result_tracer_kwargs + else default_leaf_modules + ) return result_tracer_kwargs From 1a891bb97286c3bcbc180179e0ee074ed9b4eb42 Mon Sep 17 00:00:00 2001 From: Joao Gomes Date: Thu, 17 Mar 2022 14:53:06 +0000 Subject: [PATCH 3/7] torchvision/models/feature_extraction.py --- torchvision/models/feature_extraction.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/torchvision/models/feature_extraction.py b/torchvision/models/feature_extraction.py index 24f3732427f..6d6de9306ac 100644 --- a/torchvision/models/feature_extraction.py +++ b/torchvision/models/feature_extraction.py @@ -184,7 +184,7 @@ def _get_leaf_modules_for_ops() -> List[type]: return result -def set_default_tracer_kargs(original_tr_kwargs: Optional[Dict[str, Any]]) -> Dict[str, Any]: +def set_default_tracer_kwargs(original_tr_kwargs: Optional[Dict[str, Any]]) -> Dict[str, Any]: default_autowrap_modules = (math, torchvision.ops) default_leaf_modules = _get_leaf_modules_for_ops() result_tracer_kwargs = {} if original_tr_kwargs is None else original_tr_kwargs @@ -229,7 +229,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 `_). - By default it will be set to wrap and make leaf nodes all torchvision ops, i.e. tracer_kwargs + By default it will be set to wrap and make leaf nodes all torchvision ops, i.e., tracer_kwargs will be set to tracer_kwargs = { @@ -240,7 +240,7 @@ def get_graph_node_names( "leaf_modules": _get_leaf_modules_for_ops(), } - WARNING: In case the user provides tracer_kwargs, above default arguments will be apended to the user + WARNING: In case the user provides tracer_kwargs, above default arguments will be appended to the user provided dictionary. suppress_diff_warning (bool, optional): whether to suppress a warning @@ -256,7 +256,7 @@ def get_graph_node_names( >>> model = torchvision.models.resnet18() >>> train_nodes, eval_nodes = get_graph_node_names(model) """ - tracer_kwargs = set_default_tracer_kargs(tracer_kwargs) + tracer_kwargs = set_default_tracer_kwargs(tracer_kwargs) is_training = model.training train_tracer = NodePathTracer(**tracer_kwargs) train_tracer.trace(model.train()) @@ -401,7 +401,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 `_). - By default it will be set to wrap and make leaf nodes all torchvision ops, i.e. tracer_kwargs + By default it will be set to wrap and make leaf nodes all torchvision ops, i.e., tracer_kwargs will be set to tracer_kwargs = { @@ -412,7 +412,7 @@ def create_feature_extractor( "leaf_modules": _get_leaf_modules_for_ops(), } - WARNING: In case the user provides tracer_kwargs, above default arguments will be apended to the user + WARNING: In case the user provides tracer_kwargs, above default arguments will be appended to the user provided dictionary. suppress_diff_warning (bool, optional): whether to suppress a warning when there are discrepancies between the train and eval version of @@ -458,7 +458,7 @@ def create_feature_extractor( >>> 'autowrap_functions': [leaf_function]}) """ - tracer_kwargs = set_default_tracer_kargs(tracer_kwargs) + tracer_kwargs = set_default_tracer_kwargs(tracer_kwargs) is_training = model.training if all(arg is None for arg in [return_nodes, train_return_nodes, eval_return_nodes]): From 6b412f7ce7bc5ba3f71372963c6ff04c1d364786 Mon Sep 17 00:00:00 2001 From: Joao Gomes Date: Mon, 21 Mar 2022 20:49:25 +0000 Subject: [PATCH 4/7] Adress PR comments --- torchvision/models/feature_extraction.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/torchvision/models/feature_extraction.py b/torchvision/models/feature_extraction.py index 6d6de9306ac..a1491c4e61d 100644 --- a/torchvision/models/feature_extraction.py +++ b/torchvision/models/feature_extraction.py @@ -184,7 +184,7 @@ def _get_leaf_modules_for_ops() -> List[type]: return result -def set_default_tracer_kwargs(original_tr_kwargs: Optional[Dict[str, Any]]) -> Dict[str, Any]: +def _set_default_tracer_kwargs(original_tr_kwargs: Optional[Dict[str, Any]]) -> Dict[str, Any]: default_autowrap_modules = (math, torchvision.ops) default_leaf_modules = _get_leaf_modules_for_ops() result_tracer_kwargs = {} if original_tr_kwargs is None else original_tr_kwargs @@ -256,7 +256,7 @@ def get_graph_node_names( >>> model = torchvision.models.resnet18() >>> train_nodes, eval_nodes = get_graph_node_names(model) """ - tracer_kwargs = set_default_tracer_kwargs(tracer_kwargs) + tracer_kwargs = _set_default_tracer_kwargs(tracer_kwargs) is_training = model.training train_tracer = NodePathTracer(**tracer_kwargs) train_tracer.trace(model.train()) @@ -458,7 +458,7 @@ def create_feature_extractor( >>> 'autowrap_functions': [leaf_function]}) """ - tracer_kwargs = set_default_tracer_kwargs(tracer_kwargs) + tracer_kwargs = _set_default_tracer_kwargs(tracer_kwargs) is_training = model.training if all(arg is None for arg in [return_nodes, train_return_nodes, eval_return_nodes]): From 586c019da746c9f584a71e82ab8c688eb77d6300 Mon Sep 17 00:00:00 2001 From: Joao Gomes Date: Tue, 22 Mar 2022 14:39:42 +0000 Subject: [PATCH 5/7] fix doc format --- torchvision/models/feature_extraction.py | 26 ++++++++++++------------ 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/torchvision/models/feature_extraction.py b/torchvision/models/feature_extraction.py index a1491c4e61d..8d099613768 100644 --- a/torchvision/models/feature_extraction.py +++ b/torchvision/models/feature_extraction.py @@ -401,19 +401,19 @@ 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 `_). - By default it will be set to wrap and make leaf nodes all torchvision ops, i.e., tracer_kwargs - will be set to - - tracer_kwargs = { - "autowrap_modules": ( - math, - torchvision.ops, - ), - "leaf_modules": _get_leaf_modules_for_ops(), - } - - WARNING: In case the user provides tracer_kwargs, above default arguments will be appended to the user - provided dictionary. + By default it will be set to wrap and make leaf nodes all torchvision ops, i.e., + tracer_kwargs will be set to + + tracer_kwargs = { + "autowrap_modules": ( + math, + torchvision.ops, + ), + "leaf_modules": _get_leaf_modules_for_ops(), + } + + WARNING: In case the user provides tracer_kwargs, above default arguments will be appended to the user + provided dictionary. 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. From 741d03520b7667e51f3c034fd232811569df1505 Mon Sep 17 00:00:00 2001 From: Joao Gomes Date: Tue, 22 Mar 2022 17:26:39 +0000 Subject: [PATCH 6/7] fix formatting --- torchvision/models/feature_extraction.py | 26 ++++++++++-------------- 1 file changed, 11 insertions(+), 15 deletions(-) diff --git a/torchvision/models/feature_extraction.py b/torchvision/models/feature_extraction.py index 8d099613768..21b64a8f207 100644 --- a/torchvision/models/feature_extraction.py +++ b/torchvision/models/feature_extraction.py @@ -230,8 +230,7 @@ def get_graph_node_names( ``NodePathTracer`` (they are eventually passed onto `torch.fx.Tracer `_). By default it will be set to wrap and make leaf nodes all torchvision ops, i.e., tracer_kwargs - will be set to - + will be set to: tracer_kwargs = { "autowrap_modules": ( math, @@ -239,7 +238,6 @@ def get_graph_node_names( ), "leaf_modules": _get_leaf_modules_for_ops(), } - WARNING: In case the user provides tracer_kwargs, above default arguments will be appended to the user provided dictionary. @@ -402,18 +400,16 @@ def create_feature_extractor( ``NodePathTracer`` (which passes them onto it's parent class `torch.fx.Tracer `_). By default it will be set to wrap and make leaf nodes all torchvision ops, i.e., - tracer_kwargs will be set to - - tracer_kwargs = { - "autowrap_modules": ( - math, - torchvision.ops, - ), - "leaf_modules": _get_leaf_modules_for_ops(), - } - - WARNING: In case the user provides tracer_kwargs, above default arguments will be appended to the user - provided dictionary. + tracer_kwargs will be set to: + tracer_kwargs = { + "autowrap_modules": ( + math, + torchvision.ops, + ), + "leaf_modules": _get_leaf_modules_for_ops(), + } + WARNING: In case the user provides tracer_kwargs, above default arguments will be appended to the user + provided dictionary. 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. From 81ad361b584fefa1e56fd9137178ea681682b593 Mon Sep 17 00:00:00 2001 From: Joao Gomes Date: Thu, 24 Mar 2022 12:04:31 +0000 Subject: [PATCH 7/7] fix doc error --- torchvision/models/feature_extraction.py | 22 ++++------------------ 1 file changed, 4 insertions(+), 18 deletions(-) diff --git a/torchvision/models/feature_extraction.py b/torchvision/models/feature_extraction.py index 21b64a8f207..8a25e1f7187 100644 --- a/torchvision/models/feature_extraction.py +++ b/torchvision/models/feature_extraction.py @@ -229,15 +229,8 @@ def get_graph_node_names( tracer_kwargs (dict, optional): a dictionary of keywork arguments for ``NodePathTracer`` (they are eventually passed onto `torch.fx.Tracer `_). - By default it will be set to wrap and make leaf nodes all torchvision ops, i.e., tracer_kwargs - will be set to: - tracer_kwargs = { - "autowrap_modules": ( - math, - torchvision.ops, - ), - "leaf_modules": _get_leaf_modules_for_ops(), - } + By default it will be set to wrap and make leaf nodes all torchvision ops: + {"autowrap_modules": (math, torchvision.ops,),"leaf_modules": _get_leaf_modules_for_ops(),} WARNING: In case the user provides tracer_kwargs, above default arguments will be appended to the user provided dictionary. @@ -399,15 +392,8 @@ 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 `_). - By default it will be set to wrap and make leaf nodes all torchvision ops, i.e., - tracer_kwargs will be set to: - tracer_kwargs = { - "autowrap_modules": ( - math, - torchvision.ops, - ), - "leaf_modules": _get_leaf_modules_for_ops(), - } + By default it will be set to wrap and make leaf nodes all torchvision ops: + {"autowrap_modules": (math, torchvision.ops,),"leaf_modules": _get_leaf_modules_for_ops(),} WARNING: In case the user provides tracer_kwargs, above default arguments will be appended to the user provided dictionary. suppress_diff_warning (bool, optional): whether to suppress a warning