From 8726dd5dbfd44c01f3d5ae10bf2b1067fb8deb54 Mon Sep 17 00:00:00 2001 From: Behrooz <3968947+drbeh@users.noreply.github.com> Date: Thu, 12 Aug 2021 11:40:27 -0400 Subject: [PATCH] Nvtx transform (#2713) --- docs/source/transforms.rst | 29 +++++++ monai/transforms/__init__.py | 26 ++++++ monai/transforms/nvtx.py | 125 ++++++++++++++++++++++++++++ tests/test_nvtx_transform.py | 154 +++++++++++++++++++++++++++++++++++ 4 files changed, 334 insertions(+) create mode 100644 monai/transforms/nvtx.py create mode 100644 tests/test_nvtx_transform.py diff --git a/docs/source/transforms.rst b/docs/source/transforms.rst index f97be395d1..7f970dfb15 100644 --- a/docs/source/transforms.rst +++ b/docs/source/transforms.rst @@ -341,6 +341,35 @@ IO :members: :special-members: __call__ + +NVIDIA Tool Extension (NVTX) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +`RangePush` +""""""""""" +.. autoclass:: RangePush + +`RandRangePush` +""""""""""""""" +.. autoclass:: RandRangePush + +`RangePop` +"""""""""" +.. autoclass:: RangePop + +`RandRangePop` +"""""""""""""" +.. autoclass:: RandRangePop + +`Mark` +"""""" +.. autoclass:: Mark + +`RandMark` +"""""""""" +.. autoclass:: RandMark + + Post-processing ^^^^^^^^^^^^^^^ diff --git a/monai/transforms/__init__.py b/monai/transforms/__init__.py index 390b85a1b8..33d7fba26e 100644 --- a/monai/transforms/__init__.py +++ b/monai/transforms/__init__.py @@ -195,6 +195,32 @@ from .inverse_batch_transform import BatchInverseTransform, Decollated from .io.array import LoadImage, SaveImage from .io.dictionary import LoadImaged, LoadImageD, LoadImageDict, SaveImaged, SaveImageD, SaveImageDict +from .nvtx import ( + Mark, + Markd, + MarkD, + MarkDict, + RandMark, + RandMarkd, + RandMarkD, + RandMarkDict, + RandRangePop, + RandRangePopd, + RandRangePopD, + RandRangePopDict, + RandRangePush, + RandRangePushd, + RandRangePushD, + RandRangePushDict, + RangePop, + RangePopd, + RangePopD, + RangePopDict, + RangePush, + RangePushd, + RangePushD, + RangePushDict, +) from .post.array import ( Activations, AsDiscrete, diff --git a/monai/transforms/nvtx.py b/monai/transforms/nvtx.py new file mode 100644 index 0000000000..12c03dc028 --- /dev/null +++ b/monai/transforms/nvtx.py @@ -0,0 +1,125 @@ +# Copyright 2020 - 2021 MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Wrapper around NVIDIA Tools Extension for profiling MONAI transformations +""" + +from monai.transforms.transform import RandomizableTransform, Transform +from monai.utils import optional_import + +_nvtx, _ = optional_import("torch._C._nvtx", descriptor="NVTX is not installed. Are you sure you have a CUDA build?") + +__all__ = [ + "Mark", + "Markd", + "MarkD", + "MarkDict", + "RandMark", + "RandMarkd", + "RandMarkD", + "RandMarkDict", + "RandRangePop", + "RandRangePopd", + "RandRangePopD", + "RandRangePopDict", + "RandRangePush", + "RandRangePushd", + "RandRangePushD", + "RandRangePushDict", + "RangePop", + "RangePopd", + "RangePopD", + "RangePopDict", + "RangePush", + "RangePushd", + "RangePushD", + "RangePushDict", +] + + +class RangePush(Transform): + """ + Pushes a range onto a stack of nested range span. + Stores zero-based depth of the range that is started. + + Args: + msg: ASCII message to associate with range + """ + + def __init__(self, msg: str) -> None: + self.msg = msg + self.depth = None + + def __call__(self, data): + self.depth = _nvtx.rangePushA(self.msg) + return data + + +class RandRangePush(RangePush, RandomizableTransform): + """ + Pushes a range onto a stack of nested range span (RandomizableTransform). + Stores zero-based depth of the range that is started. + + Args: + msg: ASCII message to associate with range + """ + + +class RangePop(Transform): + """ + Pops a range off of a stack of nested range spans. + Stores zero-based depth of the range that is ended. + """ + + def __call__(self, data): + _nvtx.rangePop() + return data + + +class RandRangePop(RangePop, RandomizableTransform): + """ + Pops a range off of a stack of nested range spans (RandomizableTransform). + Stores zero-based depth of the range that is ended. + """ + + +class Mark(Transform): + """ + Mark an instantaneous event that occurred at some point. + + Args: + msg: ASCII message to associate with the event. + """ + + def __init__(self, msg: str) -> None: + self.msg = msg + + def __call__(self, data): + _nvtx.markA(self.msg) + return data + + +class RandMark(Mark, RandomizableTransform): + """ + Mark an instantaneous event that occurred at some point. + (RandomizableTransform) + + Args: + msg: ASCII message to associate with the event. + """ + + +MarkDict = MarkD = Markd = Mark +RandMarkDict = RandMarkD = RandMarkd = RandMark +RandRangePopDict = RandRangePopD = RandRangePopd = RandRangePop +RandRangePushDict = RandRangePushD = RandRangePushd = RandRangePush +RangePopDict = RangePopD = RangePopd = RangePop +RangePushDict = RangePushD = RangePushd = RangePush diff --git a/tests/test_nvtx_transform.py b/tests/test_nvtx_transform.py new file mode 100644 index 0000000000..d1887377ba --- /dev/null +++ b/tests/test_nvtx_transform.py @@ -0,0 +1,154 @@ +# Copyright 2020 - 2021 MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import numpy as np +import torch +from parameterized import parameterized + +from monai.transforms import Compose, Flip, RandFlip, RandFlipD, Randomizable, ToTensor, ToTensorD +from monai.transforms.nvtx import ( + Mark, + MarkD, + RandMark, + RandMarkD, + RandRangePop, + RandRangePopD, + RandRangePush, + RandRangePushD, + RangePop, + RangePopD, + RangePush, + RangePushD, +) +from monai.utils import optional_import + +_, has_nvtx = optional_import("torch._C._nvtx", descriptor="NVTX is not installed. Are you sure you have a CUDA build?") + + +TEST_CASE_ARRAY_0 = [ + np.random.randn(3, 3), +] +TEST_CASE_ARRAY_1 = [ + np.random.randn(3, 10, 10), +] +TEST_CASE_DICT_0 = [ + {"image": np.random.randn(3, 3)}, +] +TEST_CASE_DICT_1 = [ + {"image": np.random.randn(3, 10, 10)}, +] + + +class TestNVTXTransforms(unittest.TestCase): + @parameterized.expand( + [ + TEST_CASE_ARRAY_0, + TEST_CASE_ARRAY_1, + TEST_CASE_DICT_0, + TEST_CASE_DICT_1, + ] + ) + @unittest.skipUnless(has_nvtx, "CUDA is required for NVTX!") + def test_nvtx_transfroms_alone(self, input): + transforms = Compose( + [ + Mark("Mark: Transform Starts!"), + RangePush("Range: RandFlipD"), + RangePop(), + RandRangePush("Range: ToTensorD"), + RandRangePop(), + RandMark("Mark: Transform Ends!"), + ] + ) + output = transforms(input) + self.assertEqual(id(input), id(output)) + + # Check if chain of randomizable/non-randomizable transforms is not broken + for tran in transforms.transforms: + if isinstance(tran, Randomizable): + self.assertIsInstance(tran, RangePush) + break + + @parameterized.expand([TEST_CASE_ARRAY_0, TEST_CASE_ARRAY_1]) + @unittest.skipUnless(has_nvtx, "CUDA is required for NVTX!") + def test_nvtx_transfroms_array(self, input): + transforms = Compose( + [ + RandMark("Mark: Transform Starts!"), + RandRangePush("Range: RandFlip"), + RandFlip(prob=0.0), + RandRangePop(), + RangePush("Range: ToTensor"), + ToTensor(), + RangePop(), + Mark("Mark: Transform Ends!"), + ] + ) + output = transforms(input) + self.assertIsInstance(output, torch.Tensor) + np.testing.assert_array_equal(input, output) + + transforms = Compose( + [ + RandMark("Mark: Transform Starts!"), + RandRangePush("Range: RandFlip"), + RandFlip(prob=1.0), + RandRangePop(), + RangePush("Range: ToTensor"), + ToTensor(), + RangePop(), + Mark("Mark: Transform Ends!"), + ] + ) + output = transforms(input) + self.assertIsInstance(output, torch.Tensor) + np.testing.assert_array_equal(input, Flip()(output.numpy())) + + @parameterized.expand([TEST_CASE_DICT_0, TEST_CASE_DICT_1]) + @unittest.skipUnless(has_nvtx, "CUDA is required for NVTX!") + def test_nvtx_transfromsd(self, input): + transforms = Compose( + [ + RandMarkD("Mark: Transform Starts!"), + RandRangePushD("Range: RandFlipD"), + RandFlipD(keys="image", prob=0.0), + RandRangePopD(), + RangePushD("Range: ToTensorD"), + ToTensorD(keys=("image")), + RangePopD(), + MarkD("Mark: Transform Ends!"), + ] + ) + output = transforms(input) + self.assertIsInstance(output["image"], torch.Tensor) + np.testing.assert_array_equal(input["image"], output["image"]) + + transforms = Compose( + [ + RandMarkD("Mark: Transform Starts!"), + RandRangePushD("Range: RandFlipD"), + RandFlipD(keys="image", prob=1.0), + RandRangePopD(), + RangePushD("Range: ToTensorD"), + ToTensorD(keys=("image")), + RangePopD(), + MarkD("Mark: Transform Ends!"), + ] + ) + output = transforms(input) + self.assertIsInstance(output["image"], torch.Tensor) + np.testing.assert_array_equal(input["image"], Flip()(output["image"].numpy())) + + +if __name__ == "__main__": + unittest.main()