From f6ea34d63b56daf70bd285457e220ba8b1dc3c5d Mon Sep 17 00:00:00 2001 From: Richard Brown <33289025+rijobro@users.noreply.github.com> Date: Fri, 22 Jan 2021 13:35:14 +0000 Subject: [PATCH 1/7] CopyToDevice transform Signed-off-by: Richard Brown <33289025+rijobro@users.noreply.github.com> --- monai/transforms/__init__.py | 4 ++ monai/transforms/utility/array.py | 28 +++++++++- monai/transforms/utility/dictionary.py | 26 +++++++++- tests/test_array_copytodevice.py | 49 ++++++++++++++++++ tests/test_copy_to_device.py | 71 ++++++++++++++++++++++++++ tests/test_dict_copytodevice.py | 54 ++++++++++++++++++++ 6 files changed, 230 insertions(+), 2 deletions(-) create mode 100644 tests/test_array_copytodevice.py create mode 100644 tests/test_copy_to_device.py create mode 100644 tests/test_dict_copytodevice.py diff --git a/monai/transforms/__init__.py b/monai/transforms/__init__.py index 9eaedd6b15..5adbb5ec3b 100644 --- a/monai/transforms/__init__.py +++ b/monai/transforms/__init__.py @@ -254,6 +254,7 @@ TorchVision, ToTensor, Transpose, + CopyToDevice, ) from .utility.dictionary import ( AddChanneld, @@ -316,6 +317,9 @@ ToTensord, ToTensorD, ToTensorDict, + CopyToDeviced, + CopyToDeviceD, + CopyToDeviceDict, ) from .utils import ( apply_transform, diff --git a/monai/transforms/utility/array.py b/monai/transforms/utility/array.py index 5476e800f4..dfdfb7d4a8 100644 --- a/monai/transforms/utility/array.py +++ b/monai/transforms/utility/array.py @@ -22,7 +22,7 @@ from monai.transforms.compose import Randomizable, Transform from monai.transforms.utils import extreme_points_to_image, get_extreme_points, map_binary_to_indices -from monai.utils import ensure_tuple, min_version, optional_import +from monai.utils import ensure_tuple, min_version, optional_import, copy_to_device __all__ = [ "Identity", @@ -44,6 +44,7 @@ "ConvertToMultiChannelBasedOnBratsClasses", "AddExtremePointsChannel", "TorchVision", + "CopyToDevice", ] # Generic type which can represent either a numpy.ndarray or a torch.Tensor @@ -671,3 +672,28 @@ def __call__(self, img: torch.Tensor): """ return self.trans(img) + + +class CopyToDevice(Transform): + """ + Copy to ``device`` where possible. + """ + def __init__( + self, + device: Optional[Union[str, torch.device]], + non_blocking: bool = True, + verbose: bool = False, + ) -> None: + self.device = device + self.non_blocking = non_blocking + self.verbose = verbose + + def __call__( + self, + img: Union[torch.Tensor, np.ndarray], + ) -> np.ndarray: + """ + Args: + img: the image to be moved to ``device``. + """ + return copy_to_device(img, self.device, self.non_blocking, self.verbose) diff --git a/monai/transforms/utility/dictionary.py b/monai/transforms/utility/dictionary.py index 1427f24356..7ac8c2e457 100644 --- a/monai/transforms/utility/dictionary.py +++ b/monai/transforms/utility/dictionary.py @@ -29,7 +29,7 @@ AsChannelFirst, AsChannelLast, CastToType, - ConvertToMultiChannelBasedOnBratsClasses, + ConvertToMultiChannelBasedOnBratsClasses, CopyToDevice, DataStats, FgBgToIndices, Identity, @@ -109,6 +109,9 @@ "TorchVisiond", "TorchVisionD", "TorchVisionDict", + "CopyToDeviceD", + "CopyToDeviceDict", + "CopyToDeviced", ] @@ -801,6 +804,26 @@ def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torc return d +class CopyToDeviced(MapTransform): + """ + Dictionary-based wrapper of :py:class:`monai.transforms.CopyToDevice`. + """ + + def __init__(self, + keys: KeysCollection, + device: Optional[Union[str, torch.device]], + non_blocking: bool = True, + verbose: bool = False, + ) -> None: + super().__init__(keys) + self.converter = CopyToDevice(device, non_blocking, verbose) + + def __call__(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndarray]: + d = dict(data) + for key in self.keys: + d[key] = self.converter(d[key]) + return d + IdentityD = IdentityDict = Identityd AsChannelFirstD = AsChannelFirstDict = AsChannelFirstd AsChannelLastD = AsChannelLastDict = AsChannelLastd @@ -823,3 +846,4 @@ def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torc ) = ConvertToMultiChannelBasedOnBratsClassesd AddExtremePointsChannelD = AddExtremePointsChannelDict = AddExtremePointsChanneld TorchVisionD = TorchVisionDict = TorchVisiond +CopyToDeviceD = CopyToDeviceDict = CopyToDeviced diff --git a/tests/test_array_copytodevice.py b/tests/test_array_copytodevice.py new file mode 100644 index 0000000000..1bc6cbc911 --- /dev/null +++ b/tests/test_array_copytodevice.py @@ -0,0 +1,49 @@ +# 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 +from parameterized import parameterized +from tests.utils import skip_if_no_cuda + +from monai.data import ArrayDataset +from monai.transforms import Compose, CopyToDevice, ToTensor + +DEVICE="cuda:0" + +TEST_CASE_0 = [ + Compose([ToTensor(), CopyToDevice(device=DEVICE)]), + Compose([ToTensor()]), + DEVICE, + "cpu", +] + +@skip_if_no_cuda +class TestArrayCopyToDevice(unittest.TestCase): + @parameterized.expand([TEST_CASE_0]) + def test_array_copy_to_device(self, img_transform, label_transform, img_device, label_device): + numel = 2 + test_imgs = [np.zeros((3,3,3)) for _ in range(numel)] + test_segs = [np.zeros((3,3,3)) for _ in range(numel)] + + test_labels = [1, 1] + dataset = ArrayDataset(test_imgs, img_transform, test_segs, label_transform, test_labels, None) + self.assertEqual(len(dataset), 2) + for data in dataset: + im, seg = data[0], data[1] + self.assertTrue(str(im.device) == img_device) + self.assertTrue(str(seg.device) == label_device) + + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_copy_to_device.py b/tests/test_copy_to_device.py new file mode 100644 index 0000000000..8d8d0c8784 --- /dev/null +++ b/tests/test_copy_to_device.py @@ -0,0 +1,71 @@ +# 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.utils import copy_to_device +from tests.utils import skip_if_no_cuda + +TEST_CASE_TENSOR = [ + torch.Tensor([1.0]).to("cuda:0"), + "cuda:0", + "cpu", +] +TEST_CASE_LIST = [ + 2 * [torch.Tensor([1.0])], + "cpu", + "cuda:0", +] +TEST_CASE_TUPLE = [ + 2 * (torch.Tensor([1.0]),), + "cpu", + "cuda:0", +] +TEST_CASE_MIXED_LIST = [ + [torch.Tensor([1.0]), np.array([1])], + "cpu", + "cuda:0", +] +TEST_CASE_DICT = [ + { + "x": torch.Tensor([1.0]), + "y": 2 * [torch.Tensor([1.0])], + "z": np.array([1]), + }, + "cpu", + "cuda:0", +] +TEST_CASES = [TEST_CASE_TENSOR, TEST_CASE_LIST, TEST_CASE_TUPLE, TEST_CASE_MIXED_LIST, TEST_CASE_DICT] + + +@skip_if_no_cuda +class TestCopyToDevice(unittest.TestCase): + def _check_on_device(self, obj, device): + if hasattr(obj, "device"): + self.assertTrue(str(obj.device) == device) + elif any(isinstance(obj, x) for x in [list, tuple]): + _ = [self._check_on_device(o, device) for o in obj] + elif isinstance(obj, dict): + _ = [self._check_on_device(o, device) for o in obj.values()] + + @parameterized.expand(TEST_CASES) + def test_copy(self, input, in_device, out_device): + out = copy_to_device(input, out_device) + self._check_on_device(input, in_device) + self._check_on_device(out, out_device) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_dict_copytodevice.py b/tests/test_dict_copytodevice.py new file mode 100644 index 0000000000..d0af5e32db --- /dev/null +++ b/tests/test_dict_copytodevice.py @@ -0,0 +1,54 @@ +# 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 +from parameterized import parameterized +from tests.utils import skip_if_no_cuda + +from monai.data import Dataset +from monai.transforms import Compose, CopyToDeviced, ToTensord + +DEVICE="cuda:0" + +TEST_CASE_0 = [ + Compose([ + ToTensord(keys=["image", "label", "other"]), + CopyToDeviced(keys=["image", "label"], device=DEVICE) + ]), + DEVICE, + "cpu", +] + +@skip_if_no_cuda +class TestDictCopyToDevice(unittest.TestCase): + @parameterized.expand([TEST_CASE_0]) + def test_dict_copy_to_device(self, transform, modified_device, unmodified_device): + + numel = 2 + test_data = [{ + "image": np.zeros((3,3,3)), + "label": np.zeros((3,3,3)), + "other": np.zeros((3,3,3)), + } for _ in range(numel)] + + dataset = Dataset(data=test_data, transform=transform) + self.assertEqual(len(dataset), 2) + for data in dataset: + self.assertTrue(str(data["image"].device) == modified_device) + self.assertTrue(str(data["label"].device) == modified_device) + self.assertTrue(str(data["other"].device) == unmodified_device) + + + +if __name__ == "__main__": + unittest.main() From bd2f2c7509410718963ff0ea8dd6e3417ac8f429 Mon Sep 17 00:00:00 2001 From: Rich <33289025+rijobro@users.noreply.github.com> Date: Mon, 25 Jan 2021 13:57:37 +0000 Subject: [PATCH 2/7] add copy_to_device Signed-off-by: Richard Brown <33289025+rijobro@users.noreply.github.com> --- monai/utils/__init__.py | 1 + monai/utils/misc.py | 43 ++++++++++++++++++++++++++++++++++++++++- 2 files changed, 43 insertions(+), 1 deletion(-) diff --git a/monai/utils/__init__.py b/monai/utils/__init__.py index 9bb25d723a..6430fae75a 100644 --- a/monai/utils/__init__.py +++ b/monai/utils/__init__.py @@ -32,6 +32,7 @@ ) from .misc import ( MAX_SEED, + copy_to_device, dtype_numpy_to_torch, dtype_torch_to_numpy, ensure_tuple, diff --git a/monai/utils/misc.py b/monai/utils/misc.py index bf1ff60cbc..2b31392a46 100644 --- a/monai/utils/misc.py +++ b/monai/utils/misc.py @@ -10,11 +10,14 @@ # limitations under the License. import collections.abc +import inspect import itertools import random +import types +import warnings from ast import literal_eval from distutils.util import strtobool -from typing import Any, Callable, Optional, Sequence, Tuple, Union +from typing import Any, Callable, Optional, Sequence, Tuple, Union, cast import numpy as np import torch @@ -37,6 +40,7 @@ "dtype_torch_to_numpy", "dtype_numpy_to_torch", "MAX_SEED", + "copy_to_device", ] _seed = None @@ -306,3 +310,40 @@ def dtype_torch_to_numpy(dtype): def dtype_numpy_to_torch(dtype): """Convert a numpy dtype to its torch equivalent.""" return _np_to_torch_dtype[dtype] + + +def copy_to_device( + obj: Any, + device: Optional[Union[str, torch.device]], + non_blocking: bool = True, + verbose: bool = False, +) -> Any: + """ + Copy object or tuple/list/dictionary of objects to ``device``. + + Args: + obj: object or tuple/list/dictionary of objects to move to ``device``. + device: move ``obj`` to this device. Can be a string (e.g., ``cpu``, ``cuda``, + ``cuda:0``, etc.) or of type ``torch.device``. + non_blocking_transfer: when `True`, moves data to device asynchronously if + possible, e.g., moving CPU Tensors with pinned memory to CUDA devices. + verbose: when `True`, will print a warning for any elements of incompatible type + not copied to ``device``. + Returns: + Same as input, copied to ``device`` where possible. Original input will be + unchanged. + """ + + if hasattr(obj, "to"): + return obj.to(device, non_blocking=non_blocking) + elif isinstance(obj, tuple): + return tuple(copy_to_device(o, device, non_blocking) for o in obj) + elif isinstance(obj, list): + return [copy_to_device(o, device, non_blocking) for o in obj] + elif isinstance(obj, dict): + return {k: copy_to_device(o, device, non_blocking) for k, o in obj.items()} + elif verbose: + fn_name = cast(types.FrameType, inspect.currentframe()).f_code.co_name + warnings.warn(f"{fn_name} called with incompatible type: " + f"{type(obj)}. Data will be returned unchanged.") + + return obj From fd08522b822180cb899c4221ac9079782913997f Mon Sep 17 00:00:00 2001 From: Richard Brown <33289025+rijobro@users.noreply.github.com> Date: Mon, 25 Jan 2021 14:01:56 +0000 Subject: [PATCH 3/7] add copy_to_device fn Signed-off-by: Richard Brown <33289025+rijobro@users.noreply.github.com> --- monai/transforms/__init__.py | 8 ++++---- monai/transforms/utility/array.py | 3 ++- monai/transforms/utility/dictionary.py | 7 +++++-- tests/test_array_copytodevice.py | 10 +++++----- tests/test_dict_copytodevice.py | 24 ++++++++++++------------ 5 files changed, 28 insertions(+), 24 deletions(-) diff --git a/monai/transforms/__init__.py b/monai/transforms/__init__.py index 5adbb5ec3b..874edbee42 100644 --- a/monai/transforms/__init__.py +++ b/monai/transforms/__init__.py @@ -241,6 +241,7 @@ AsChannelLast, CastToType, ConvertToMultiChannelBasedOnBratsClasses, + CopyToDevice, DataStats, FgBgToIndices, Identity, @@ -254,7 +255,6 @@ TorchVision, ToTensor, Transpose, - CopyToDevice, ) from .utility.dictionary import ( AddChanneld, @@ -281,6 +281,9 @@ CopyItemsd, CopyItemsD, CopyItemsDict, + CopyToDeviced, + CopyToDeviceD, + CopyToDeviceDict, DataStatsd, DataStatsD, DataStatsDict, @@ -317,9 +320,6 @@ ToTensord, ToTensorD, ToTensorDict, - CopyToDeviced, - CopyToDeviceD, - CopyToDeviceDict, ) from .utils import ( apply_transform, diff --git a/monai/transforms/utility/array.py b/monai/transforms/utility/array.py index dfdfb7d4a8..97fa78ca6f 100644 --- a/monai/transforms/utility/array.py +++ b/monai/transforms/utility/array.py @@ -22,7 +22,7 @@ from monai.transforms.compose import Randomizable, Transform from monai.transforms.utils import extreme_points_to_image, get_extreme_points, map_binary_to_indices -from monai.utils import ensure_tuple, min_version, optional_import, copy_to_device +from monai.utils import copy_to_device, ensure_tuple, min_version, optional_import __all__ = [ "Identity", @@ -678,6 +678,7 @@ class CopyToDevice(Transform): """ Copy to ``device`` where possible. """ + def __init__( self, device: Optional[Union[str, torch.device]], diff --git a/monai/transforms/utility/dictionary.py b/monai/transforms/utility/dictionary.py index 7ac8c2e457..091a9de8a5 100644 --- a/monai/transforms/utility/dictionary.py +++ b/monai/transforms/utility/dictionary.py @@ -29,7 +29,8 @@ AsChannelFirst, AsChannelLast, CastToType, - ConvertToMultiChannelBasedOnBratsClasses, CopyToDevice, + ConvertToMultiChannelBasedOnBratsClasses, + CopyToDevice, DataStats, FgBgToIndices, Identity, @@ -809,7 +810,8 @@ class CopyToDeviced(MapTransform): Dictionary-based wrapper of :py:class:`monai.transforms.CopyToDevice`. """ - def __init__(self, + def __init__( + self, keys: KeysCollection, device: Optional[Union[str, torch.device]], non_blocking: bool = True, @@ -824,6 +826,7 @@ def __call__(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.nda d[key] = self.converter(d[key]) return d + IdentityD = IdentityDict = Identityd AsChannelFirstD = AsChannelFirstDict = AsChannelFirstd AsChannelLastD = AsChannelLastDict = AsChannelLastd diff --git a/tests/test_array_copytodevice.py b/tests/test_array_copytodevice.py index 1bc6cbc911..ee73cad753 100644 --- a/tests/test_array_copytodevice.py +++ b/tests/test_array_copytodevice.py @@ -13,12 +13,12 @@ import numpy as np from parameterized import parameterized -from tests.utils import skip_if_no_cuda from monai.data import ArrayDataset from monai.transforms import Compose, CopyToDevice, ToTensor +from tests.utils import skip_if_no_cuda -DEVICE="cuda:0" +DEVICE = "cuda:0" TEST_CASE_0 = [ Compose([ToTensor(), CopyToDevice(device=DEVICE)]), @@ -27,13 +27,14 @@ "cpu", ] + @skip_if_no_cuda class TestArrayCopyToDevice(unittest.TestCase): @parameterized.expand([TEST_CASE_0]) def test_array_copy_to_device(self, img_transform, label_transform, img_device, label_device): numel = 2 - test_imgs = [np.zeros((3,3,3)) for _ in range(numel)] - test_segs = [np.zeros((3,3,3)) for _ in range(numel)] + test_imgs = [np.zeros((3, 3, 3)) for _ in range(numel)] + test_segs = [np.zeros((3, 3, 3)) for _ in range(numel)] test_labels = [1, 1] dataset = ArrayDataset(test_imgs, img_transform, test_segs, label_transform, test_labels, None) @@ -44,6 +45,5 @@ def test_array_copy_to_device(self, img_transform, label_transform, img_device, self.assertTrue(str(seg.device) == label_device) - if __name__ == "__main__": unittest.main() diff --git a/tests/test_dict_copytodevice.py b/tests/test_dict_copytodevice.py index d0af5e32db..cdcb645c04 100644 --- a/tests/test_dict_copytodevice.py +++ b/tests/test_dict_copytodevice.py @@ -13,33 +13,34 @@ import numpy as np from parameterized import parameterized -from tests.utils import skip_if_no_cuda from monai.data import Dataset from monai.transforms import Compose, CopyToDeviced, ToTensord +from tests.utils import skip_if_no_cuda -DEVICE="cuda:0" +DEVICE = "cuda:0" TEST_CASE_0 = [ - Compose([ - ToTensord(keys=["image", "label", "other"]), - CopyToDeviced(keys=["image", "label"], device=DEVICE) - ]), + Compose([ToTensord(keys=["image", "label", "other"]), CopyToDeviced(keys=["image", "label"], device=DEVICE)]), DEVICE, "cpu", ] + @skip_if_no_cuda class TestDictCopyToDevice(unittest.TestCase): @parameterized.expand([TEST_CASE_0]) def test_dict_copy_to_device(self, transform, modified_device, unmodified_device): numel = 2 - test_data = [{ - "image": np.zeros((3,3,3)), - "label": np.zeros((3,3,3)), - "other": np.zeros((3,3,3)), - } for _ in range(numel)] + test_data = [ + { + "image": np.zeros((3, 3, 3)), + "label": np.zeros((3, 3, 3)), + "other": np.zeros((3, 3, 3)), + } + for _ in range(numel) + ] dataset = Dataset(data=test_data, transform=transform) self.assertEqual(len(dataset), 2) @@ -49,6 +50,5 @@ def test_dict_copy_to_device(self, transform, modified_device, unmodified_device self.assertTrue(str(data["other"].device) == unmodified_device) - if __name__ == "__main__": unittest.main() From e56e38611828c7fc0830f9cde124dcd36de9be44 Mon Sep 17 00:00:00 2001 From: Richard Brown <33289025+rijobro@users.noreply.github.com> Date: Mon, 25 Jan 2021 14:16:45 +0000 Subject: [PATCH 4/7] update skip with no cuda Signed-off-by: Richard Brown <33289025+rijobro@users.noreply.github.com> --- tests/test_copy_to_device.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_copy_to_device.py b/tests/test_copy_to_device.py index 8d8d0c8784..b9cd2c219d 100644 --- a/tests/test_copy_to_device.py +++ b/tests/test_copy_to_device.py @@ -10,13 +10,13 @@ # limitations under the License. import unittest +from unittest.case import skipUnless import numpy as np import torch from parameterized import parameterized from monai.utils import copy_to_device -from tests.utils import skip_if_no_cuda TEST_CASE_TENSOR = [ torch.Tensor([1.0]).to("cuda:0"), @@ -50,7 +50,7 @@ TEST_CASES = [TEST_CASE_TENSOR, TEST_CASE_LIST, TEST_CASE_TUPLE, TEST_CASE_MIXED_LIST, TEST_CASE_DICT] -@skip_if_no_cuda +@skipUnless(torch.cuda.is_available(), "torch required to be built with CUDA.") class TestCopyToDevice(unittest.TestCase): def _check_on_device(self, obj, device): if hasattr(obj, "device"): From b11138ff11b4b9c6481c3e7f53a603c2df16aa34 Mon Sep 17 00:00:00 2001 From: Richard Brown <33289025+rijobro@users.noreply.github.com> Date: Mon, 25 Jan 2021 14:27:30 +0000 Subject: [PATCH 5/7] Revert "update skip with no cuda" This reverts commit e56e38611828c7fc0830f9cde124dcd36de9be44. Signed-off-by: Richard Brown <33289025+rijobro@users.noreply.github.com> --- tests/test_copy_to_device.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_copy_to_device.py b/tests/test_copy_to_device.py index b9cd2c219d..8d8d0c8784 100644 --- a/tests/test_copy_to_device.py +++ b/tests/test_copy_to_device.py @@ -10,13 +10,13 @@ # limitations under the License. import unittest -from unittest.case import skipUnless import numpy as np import torch from parameterized import parameterized from monai.utils import copy_to_device +from tests.utils import skip_if_no_cuda TEST_CASE_TENSOR = [ torch.Tensor([1.0]).to("cuda:0"), @@ -50,7 +50,7 @@ TEST_CASES = [TEST_CASE_TENSOR, TEST_CASE_LIST, TEST_CASE_TUPLE, TEST_CASE_MIXED_LIST, TEST_CASE_DICT] -@skipUnless(torch.cuda.is_available(), "torch required to be built with CUDA.") +@skip_if_no_cuda class TestCopyToDevice(unittest.TestCase): def _check_on_device(self, obj, device): if hasattr(obj, "device"): From a1784d1a082f9c4a1ff4f6879ed065703f2443f2 Mon Sep 17 00:00:00 2001 From: Richard Brown <33289025+rijobro@users.noreply.github.com> Date: Mon, 25 Jan 2021 14:34:48 +0000 Subject: [PATCH 6/7] fix copy to Signed-off-by: Richard Brown <33289025+rijobro@users.noreply.github.com> --- tests/test_copy_to_device.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_copy_to_device.py b/tests/test_copy_to_device.py index 8d8d0c8784..d9e0afc8f6 100644 --- a/tests/test_copy_to_device.py +++ b/tests/test_copy_to_device.py @@ -18,8 +18,10 @@ from monai.utils import copy_to_device from tests.utils import skip_if_no_cuda +DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu:0" + TEST_CASE_TENSOR = [ - torch.Tensor([1.0]).to("cuda:0"), + torch.Tensor([1.0]).to(DEVICE), "cuda:0", "cpu", ] From 704904baf42ba445418a04a521b492eea8853041 Mon Sep 17 00:00:00 2001 From: Richard Brown <33289025+rijobro@users.noreply.github.com> Date: Mon, 25 Jan 2021 16:53:22 +0000 Subject: [PATCH 7/7] add comparison Signed-off-by: Richard Brown <33289025+rijobro@users.noreply.github.com> --- .../test_dict_copytodevice_timecomparison.py | 95 +++++++++++++++++++ 1 file changed, 95 insertions(+) create mode 100644 tests/test_dict_copytodevice_timecomparison.py diff --git a/tests/test_dict_copytodevice_timecomparison.py b/tests/test_dict_copytodevice_timecomparison.py new file mode 100644 index 0000000000..aba907d8c4 --- /dev/null +++ b/tests/test_dict_copytodevice_timecomparison.py @@ -0,0 +1,95 @@ +# 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 os +import time +import unittest + +import torch +from torch.utils.data import DataLoader + +from monai.apps import MedNISTDataset +from monai.networks.nets import densenet121 +from monai.transforms import Compose, CopyToDeviced, ToTensord, LoadImaged, AddChanneld +from tests.utils import skip_if_no_cuda + +# This test is only run with cuda +DEVICE = "cuda:0" + +@skip_if_no_cuda +class TestDictCopyToDeviceTimeComparison(unittest.TestCase): + + @staticmethod + def get_data(use_copy_to_device_transform): + + root_dir = os.environ.get("MONAI_DATA_DIRECTORY") + if not root_dir: + root_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "testing_data") + + transforms = Compose( + [ + LoadImaged(keys="image"), + AddChanneld(keys="image"), + ToTensord(keys="image"), + ] + ) + # If necessary, append the transform + if use_copy_to_device_transform: + transforms.transforms = transforms.transforms + (CopyToDeviced(keys="image", device=DEVICE),) + + train_ds = MedNISTDataset( + root_dir=root_dir, + transform=transforms, + section="validation", + val_frac=0.001, + download=True, + num_workers=10, + ) + train_loader = DataLoader(train_ds, batch_size=300, shuffle=True, num_workers=10) + num_classes = train_ds.get_num_classes() + + model = densenet121(spatial_dims=2, in_channels=1, out_channels=num_classes).to(DEVICE) + + return train_loader, model + + def test_dict_copy_to_device_time_comparison(self): + + + for use_copy_transform in [True, False]: + start_time = time.time() + + train_loader, model = self.get_data(use_copy_transform) + + model.train() + for batch_data in train_loader: + inputs, labels = batch_data["image"], batch_data["label"] + # If using the copy transform, check they're on the GPU + if use_copy_transform: + self.assertEqual(str(inputs.device), DEVICE) + # Assert not already on device, and then copy them there + else: + self.assertNotEqual(str(inputs.device), DEVICE) + inputs = inputs.to(DEVICE) + labels = labels.to(DEVICE) + + loss_function = torch.nn.CrossEntropyLoss() + optimizer = torch.optim.Adam(model.parameters(), 1e-5) + optimizer.zero_grad() + outputs = model(inputs) + loss = loss_function(outputs, labels) + loss.backward() + optimizer.step() + + print(f"--- {time.time() - start_time} seconds ---") + + +if __name__ == "__main__": + unittest.main()