From 5099a23daf69b7128c5064132198909faf268370 Mon Sep 17 00:00:00 2001 From: Behrooz <3968947+drbeh@users.noreply.github.com> Date: Fri, 13 May 2022 14:11:30 +0000 Subject: [PATCH 1/7] Make Range to be recuresivley applied Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com> --- monai/utils/nvtx.py | 19 ++++++++++++++++++- tests/test_nvtx_decorator.py | 19 ++++++++++++++++++- 2 files changed, 36 insertions(+), 2 deletions(-) diff --git a/monai/utils/nvtx.py b/monai/utils/nvtx.py index 691f900c7d..62e954cb73 100644 --- a/monai/utils/nvtx.py +++ b/monai/utils/nvtx.py @@ -13,6 +13,7 @@ """ from collections import defaultdict +from collections.abc import Iterable from functools import wraps from typing import Any, Optional, Tuple, Union @@ -21,7 +22,9 @@ from torch.optim import Optimizer from torch.utils.data import Dataset -from monai.utils import ensure_tuple, optional_import +# from monai.transforms.compose import Compose +from monai.utils.misc import ensure_tuple +from monai.utils.module import optional_import _nvtx, _ = optional_import("torch._C._nvtx", descriptor="NVTX is not installed. Are you sure you have a CUDA build?") @@ -53,12 +56,26 @@ def __init__( name: Optional[str] = None, methods: Optional[Union[str, Tuple[str, ...]]] = None, append_method_name: Optional[bool] = None, + recursive=False, ) -> None: self.name = name self.methods = methods self.append_method_name = append_method_name + self.recursive = recursive def __call__(self, obj: Any): + if self.recursive is True: + if type(obj).__name__ == "Compose": + annotated_transforms = [Range(recursive=False)(t) for t in obj.transforms] + return Range(self.name, recursive=False)(type(obj)(annotated_transforms)) + + if isinstance(obj, Iterable): + annotated_transforms = [Range(recursive=False)(t) for t in obj] + return Range(self.name, recursive=False)(annotated_transforms) + + self.recursive = False + print("`recursive=True` only works for Iterables or Compose. Falling back to `recursive=False`") + # Define the name to be associated to the range if not provided if self.name is None: name = type(obj).__name__ diff --git a/tests/test_nvtx_decorator.py b/tests/test_nvtx_decorator.py index e81c72efcf..9712a663fb 100644 --- a/tests/test_nvtx_decorator.py +++ b/tests/test_nvtx_decorator.py @@ -26,6 +26,7 @@ Randomizable, Rotate90, ToCupy, + ToNumpy, TorchVision, ToTensor, ToTensorD, @@ -50,7 +51,7 @@ TEST_CASE_WRAPPER = [np.random.randn(3, 10, 10)] -@unittest.skipUnless(has_nvtx, "CUDA is required for NVTX Range!") +@unittest.skipUnless(has_nvtx, "Required torch._C._nvtx for NVTX Range!") class TestNVTXRangeDecorator(unittest.TestCase): @parameterized.expand([TEST_CASE_ARRAY_0, TEST_CASE_ARRAY_1]) def test_tranform_array(self, input): @@ -127,6 +128,22 @@ def test_wrapper_tranforms(self, input): # Check the outputs np.testing.assert_equal(output.get(), output_r.get()) + @parameterized.expand([TEST_CASE_TORCH_0, TEST_CASE_ARRAY_0]) + def test_recursive_tranforms(self, input): + transform_list = [ToNumpy(), Flip(), RandAdjustContrast(prob=0.0), RandFlip(prob=1.0), ToTensor()] + + transforms = Compose(transform_list) + transforms_range = Range(recursive=True)(Compose(transform_list)) + + # Apply transforms + output = transforms(input) + + # Apply transforms with Range + output_r = transforms_range(input) + + # Check the outputs + np.testing.assert_equal(output.numpy(), output_r.numpy()) + @parameterized.expand([TEST_CASE_ARRAY_1]) def test_tranform_randomized(self, input): # Compose deterministic and randomized transforms From 117d7b43ca831658065cd06493de82674a06d957 Mon Sep 17 00:00:00 2001 From: Behrooz <3968947+drbeh@users.noreply.github.com> Date: Fri, 13 May 2022 14:46:53 +0000 Subject: [PATCH 2/7] Deferred import Compose Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com> --- monai/utils/nvtx.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/monai/utils/nvtx.py b/monai/utils/nvtx.py index 62e954cb73..209d6f2d69 100644 --- a/monai/utils/nvtx.py +++ b/monai/utils/nvtx.py @@ -22,7 +22,6 @@ from torch.optim import Optimizer from torch.utils.data import Dataset -# from monai.transforms.compose import Compose from monai.utils.misc import ensure_tuple from monai.utils.module import optional_import @@ -65,7 +64,9 @@ def __init__( def __call__(self, obj: Any): if self.recursive is True: - if type(obj).__name__ == "Compose": + from monai.transforms.compose import Compose + + if isinstance(obj, Compose): annotated_transforms = [Range(recursive=False)(t) for t in obj.transforms] return Range(self.name, recursive=False)(type(obj)(annotated_transforms)) From 1268cd86f62041a53ae52b4532df0e2dbf5fb008 Mon Sep 17 00:00:00 2001 From: Behrooz <3968947+drbeh@users.noreply.github.com> Date: Fri, 13 May 2022 14:58:20 +0000 Subject: [PATCH 3/7] Update docstring Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com> --- monai/utils/nvtx.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/monai/utils/nvtx.py b/monai/utils/nvtx.py index 209d6f2d69..903de24d3d 100644 --- a/monai/utils/nvtx.py +++ b/monai/utils/nvtx.py @@ -45,6 +45,8 @@ class Range: Otherwise, it look up predefined methods: "forward", "__call__", "__next__", "__getitem__" append_method_name: if append the name of the methods to be decorated to the range's name If None (default), it appends the method's name only if we are annotating more than one method. + recursive: if set to True, it will recursively annotate every individual module in a list + or in a chain of modules (chained using Compose). Default to False. """ @@ -55,7 +57,7 @@ def __init__( name: Optional[str] = None, methods: Optional[Union[str, Tuple[str, ...]]] = None, append_method_name: Optional[bool] = None, - recursive=False, + recursive: bool = False, ) -> None: self.name = name self.methods = methods From 1f19703d0fea4838891b51ca4f27a41d2f97fb84 Mon Sep 17 00:00:00 2001 From: Behrooz <3968947+drbeh@users.noreply.github.com> Date: Mon, 16 May 2022 01:09:40 +0000 Subject: [PATCH 4/7] address comment Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com> --- monai/utils/nvtx.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/monai/utils/nvtx.py b/monai/utils/nvtx.py index 903de24d3d..7fc636ae97 100644 --- a/monai/utils/nvtx.py +++ b/monai/utils/nvtx.py @@ -22,8 +22,7 @@ from torch.optim import Optimizer from torch.utils.data import Dataset -from monai.utils.misc import ensure_tuple -from monai.utils.module import optional_import +from monai.utils import ensure_tuple, optional_import _nvtx, _ = optional_import("torch._C._nvtx", descriptor="NVTX is not installed. Are you sure you have a CUDA build?") From 11c1bf696563787c1e2a270ed65d18b2dce7ab8c Mon Sep 17 00:00:00 2001 From: Behrooz <3968947+drbeh@users.noreply.github.com> Date: Mon, 16 May 2022 16:38:46 +0000 Subject: [PATCH 5/7] Support OneOf and Compose/OneOf of Compose/OneOf Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com> --- monai/utils/nvtx.py | 34 +++++++++++++++++++++++++-------- tests/test_nvtx_decorator.py | 37 +++++++++++++++++++++++++++--------- 2 files changed, 54 insertions(+), 17 deletions(-) diff --git a/monai/utils/nvtx.py b/monai/utils/nvtx.py index 7fc636ae97..1a9341ea09 100644 --- a/monai/utils/nvtx.py +++ b/monai/utils/nvtx.py @@ -65,18 +65,36 @@ def __init__( def __call__(self, obj: Any): if self.recursive is True: - from monai.transforms.compose import Compose - - if isinstance(obj, Compose): - annotated_transforms = [Range(recursive=False)(t) for t in obj.transforms] - return Range(self.name, recursive=False)(type(obj)(annotated_transforms)) - if isinstance(obj, Iterable): - annotated_transforms = [Range(recursive=False)(t) for t in obj] + annotated_transforms = [Range(recursive=True)(t) for t in obj] return Range(self.name, recursive=False)(annotated_transforms) + from monai.transforms.compose import Compose, OneOf + + if isinstance(obj, OneOf): + annotated_transforms = [Range(recursive=True)(t) for t in obj.transforms] + return Range(self.name, recursive=False)( + OneOf( + annotated_transforms, + weights=obj.weights, + map_items=obj.map_items, + unpack_items=obj.unpack_items, + log_stats=obj.log_stats, + ) + ) + + if isinstance(obj, Compose): + annotated_transforms = [Range(recursive=True)(t) for t in obj.transforms] + return Range(self.name, recursive=False)( + Compose( + annotated_transforms, + map_items=obj.map_items, + unpack_items=obj.unpack_items, + log_stats=obj.log_stats, + ) + ) + self.recursive = False - print("`recursive=True` only works for Iterables or Compose. Falling back to `recursive=False`") # Define the name to be associated to the range if not provided if self.name is None: diff --git a/tests/test_nvtx_decorator.py b/tests/test_nvtx_decorator.py index 9712a663fb..04b351a00c 100644 --- a/tests/test_nvtx_decorator.py +++ b/tests/test_nvtx_decorator.py @@ -19,7 +19,8 @@ Compose, CuCIM, Flip, - FlipD, + Flipd, + OneOf, RandAdjustContrast, RandCuCIM, RandFlip, @@ -29,7 +30,7 @@ ToNumpy, TorchVision, ToTensor, - ToTensorD, + ToTensord, ) from monai.utils import Range, optional_import from tests.utils import HAS_CUPY @@ -50,6 +51,26 @@ TEST_CASE_WRAPPER = [np.random.randn(3, 10, 10)] +TEST_CASE_RECURSIVE_0 = [ + torch.randn(3, 3), + Compose([ToNumpy(), Flip(), RandAdjustContrast(prob=0.0), RandFlip(prob=1.0), ToTensor()]), +] +TEST_CASE_RECURSIVE_1 = [ + torch.randn(3, 3), + Compose([ToNumpy(), Flip(), Compose([RandAdjustContrast(prob=0.0), RandFlip(prob=1.0)]), ToTensor()]), +] +TEST_CASE_RECURSIVE_2 = [ + torch.randn(3, 3), + Compose( + [ + ToNumpy(), + Flip(), + OneOf([RandAdjustContrast(prob=0.0), RandFlip(prob=1.0)], weights=[0, 1], log_stats=True), + ToTensor(), + ] + ), +] + @unittest.skipUnless(has_nvtx, "Required torch._C._nvtx for NVTX Range!") class TestNVTXRangeDecorator(unittest.TestCase): @@ -80,7 +101,7 @@ def test_tranform_array(self, input): @parameterized.expand([TEST_CASE_DICT_0, TEST_CASE_DICT_1]) def test_tranform_dict(self, input): - transforms = Compose([Range("random flip dict")(FlipD(keys="image")), Range()(ToTensorD("image"))]) + transforms = Compose([Range("random flip dict")(Flipd(keys="image")), Range()(ToTensord("image"))]) # Apply transforms output = transforms(input)["image"] @@ -128,12 +149,9 @@ def test_wrapper_tranforms(self, input): # Check the outputs np.testing.assert_equal(output.get(), output_r.get()) - @parameterized.expand([TEST_CASE_TORCH_0, TEST_CASE_ARRAY_0]) - def test_recursive_tranforms(self, input): - transform_list = [ToNumpy(), Flip(), RandAdjustContrast(prob=0.0), RandFlip(prob=1.0), ToTensor()] - - transforms = Compose(transform_list) - transforms_range = Range(recursive=True)(Compose(transform_list)) + @parameterized.expand([TEST_CASE_RECURSIVE_0, TEST_CASE_RECURSIVE_1, TEST_CASE_RECURSIVE_2]) + def test_recursive_tranforms(self, input, transforms): + transforms_range = Range(recursive=True)(transforms) # Apply transforms output = transforms(input) @@ -142,6 +160,7 @@ def test_recursive_tranforms(self, input): output_r = transforms_range(input) # Check the outputs + self.assertEqual(transforms.log_stats, transforms_range.log_stats) np.testing.assert_equal(output.numpy(), output_r.numpy()) @parameterized.expand([TEST_CASE_ARRAY_1]) From 0bc85ac08b1b6f27b082d69b74dac08adccdbd6e Mon Sep 17 00:00:00 2001 From: Behrooz <3968947+drbeh@users.noreply.github.com> Date: Tue, 17 May 2022 17:09:30 +0000 Subject: [PATCH 6/7] Update recursive and add a new test case Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com> --- monai/utils/nvtx.py | 33 ++++++--------------------------- tests/test_nvtx_decorator.py | 21 ++++++++++++++++++++- 2 files changed, 26 insertions(+), 28 deletions(-) diff --git a/monai/utils/nvtx.py b/monai/utils/nvtx.py index 1a9341ea09..efb16b59bf 100644 --- a/monai/utils/nvtx.py +++ b/monai/utils/nvtx.py @@ -13,7 +13,6 @@ """ from collections import defaultdict -from collections.abc import Iterable from functools import wraps from typing import Any, Optional, Tuple, Union @@ -65,34 +64,14 @@ def __init__( def __call__(self, obj: Any): if self.recursive is True: - if isinstance(obj, Iterable): - annotated_transforms = [Range(recursive=True)(t) for t in obj] - return Range(self.name, recursive=False)(annotated_transforms) - - from monai.transforms.compose import Compose, OneOf - - if isinstance(obj, OneOf): - annotated_transforms = [Range(recursive=True)(t) for t in obj.transforms] - return Range(self.name, recursive=False)( - OneOf( - annotated_transforms, - weights=obj.weights, - map_items=obj.map_items, - unpack_items=obj.unpack_items, - log_stats=obj.log_stats, - ) - ) + if isinstance(obj, (list, tuple)): + return type(obj)(Range(recursive=True)(t) for t in obj) + + from monai.transforms.compose import Compose if isinstance(obj, Compose): - annotated_transforms = [Range(recursive=True)(t) for t in obj.transforms] - return Range(self.name, recursive=False)( - Compose( - annotated_transforms, - map_items=obj.map_items, - unpack_items=obj.unpack_items, - log_stats=obj.log_stats, - ) - ) + obj.transforms = Range(recursive=True)(obj.transforms) + return Range(self.name, recursive=False)(obj) self.recursive = False diff --git a/tests/test_nvtx_decorator.py b/tests/test_nvtx_decorator.py index 04b351a00c..9932b678c9 100644 --- a/tests/test_nvtx_decorator.py +++ b/tests/test_nvtx_decorator.py @@ -70,6 +70,10 @@ ] ), ] +TEST_CASE_RECURSIVE_LIST = [ + torch.randn(3, 3), + [ToNumpy(), Flip(), RandAdjustContrast(prob=0.0), RandFlip(prob=1.0), ToTensor()], +] @unittest.skipUnless(has_nvtx, "Required torch._C._nvtx for NVTX Range!") @@ -151,7 +155,7 @@ def test_wrapper_tranforms(self, input): @parameterized.expand([TEST_CASE_RECURSIVE_0, TEST_CASE_RECURSIVE_1, TEST_CASE_RECURSIVE_2]) def test_recursive_tranforms(self, input, transforms): - transforms_range = Range(recursive=True)(transforms) + transforms_range = Range(name="Recursive Compose", recursive=True)(transforms) # Apply transforms output = transforms(input) @@ -160,9 +164,24 @@ def test_recursive_tranforms(self, input, transforms): output_r = transforms_range(input) # Check the outputs + self.assertEqual(transforms.map_items, transforms_range.map_items) + self.assertEqual(transforms.unpack_items, transforms_range.unpack_items) self.assertEqual(transforms.log_stats, transforms_range.log_stats) np.testing.assert_equal(output.numpy(), output_r.numpy()) + @parameterized.expand([TEST_CASE_RECURSIVE_LIST]) + def test_recursive_list_tranforms(self, input, transform_list): + transforms_list_range = Range(recursive=True)(transform_list) + + # Apply transforms + output = Compose(transform_list)(input) + + # Apply transforms with Range + output_r = Compose(transforms_list_range)(input) + + # Check the outputs + np.testing.assert_equal(output.numpy(), output_r.numpy()) + @parameterized.expand([TEST_CASE_ARRAY_1]) def test_tranform_randomized(self, input): # Compose deterministic and randomized transforms From 09819180a17e9d77a3bcd7855ff4eae59d6f9cff Mon Sep 17 00:00:00 2001 From: Behrooz <3968947+drbeh@users.noreply.github.com> Date: Wed, 18 May 2022 13:18:14 +0000 Subject: [PATCH 7/7] Remove redundant range: Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com> --- monai/utils/nvtx.py | 1 - 1 file changed, 1 deletion(-) diff --git a/monai/utils/nvtx.py b/monai/utils/nvtx.py index efb16b59bf..fefab380f1 100644 --- a/monai/utils/nvtx.py +++ b/monai/utils/nvtx.py @@ -71,7 +71,6 @@ def __call__(self, obj: Any): if isinstance(obj, Compose): obj.transforms = Range(recursive=True)(obj.transforms) - return Range(self.name, recursive=False)(obj) self.recursive = False