From 35afe1a8a0f7305b7d4a1c52864a9406f0719068 Mon Sep 17 00:00:00 2001 From: Wenqi Li Date: Mon, 29 Nov 2021 14:56:34 +0000 Subject: [PATCH 1/3] fixes pickle mod Signed-off-by: Wenqi Li --- monai/data/dataset.py | 14 +++++++------- monai/data/utils.py | 3 +++ tests/test_persistentdataset.py | 2 +- 3 files changed, 11 insertions(+), 8 deletions(-) diff --git a/monai/data/dataset.py b/monai/data/dataset.py index 8b8845f20a..675c852a06 100644 --- a/monai/data/dataset.py +++ b/monai/data/dataset.py @@ -29,9 +29,9 @@ from torch.utils.data import Dataset as _TorchDataset from torch.utils.data import Subset -from monai.data.utils import convert_tables_to_dicts, pickle_hashing +from monai.data.utils import SUPPORTED_PICKLE_MOD, convert_tables_to_dicts, pickle_hashing from monai.transforms import Compose, Randomizable, ThreadUnsafe, Transform, apply_transform -from monai.utils import MAX_SEED, ensure_tuple, get_seed, min_version, optional_import +from monai.utils import MAX_SEED, ensure_tuple, get_seed, look_up_option, min_version, optional_import from monai.utils.misc import first if TYPE_CHECKING: @@ -152,7 +152,7 @@ def __init__( transform: Union[Sequence[Callable], Callable], cache_dir: Optional[Union[Path, str]], hash_func: Callable[..., bytes] = pickle_hashing, - pickle_module=pickle, + pickle_module: str = "pickle", pickle_protocol=pickle.DEFAULT_PROTOCOL, ) -> None: """ @@ -169,7 +169,7 @@ def __init__( If `cache_dir` is `None`, there is effectively no caching. hash_func: a callable to compute hash from data items to be cached. defaults to `monai.data.utils.pickle_hashing`. - pickle_module: module used for pickling metadata and objects, default to `pickle`. + pickle_module: module used for pickling metadata and objects, default to `"pickle"`. this arg is used by `torch.save`, for more details, please check: https://pytorch.org/docs/stable/generated/torch.save.html#torch.save. pickle_protocol: can be specified to override the default protocol, default to `2`. @@ -182,7 +182,7 @@ def __init__( super().__init__(data=data, transform=transform) self.cache_dir = Path(cache_dir) if cache_dir is not None else None self.hash_func = hash_func - self.pickle_module = pickle_module + self.pickle_module = look_up_option(pickle_module, SUPPORTED_PICKLE_MOD) self.pickle_protocol = pickle_protocol if self.cache_dir is not None: if not self.cache_dir.exists(): @@ -317,7 +317,7 @@ def __init__( cache_n_trans: int, cache_dir: Optional[Union[Path, str]], hash_func: Callable[..., bytes] = pickle_hashing, - pickle_module=pickle, + pickle_module: str = "pickle", pickle_protocol=pickle.DEFAULT_PROTOCOL, ) -> None: """ @@ -335,7 +335,7 @@ def __init__( If `cache_dir` is `None`, there is effectively no caching. hash_func: a callable to compute hash from data items to be cached. defaults to `monai.data.utils.pickle_hashing`. - pickle_module: module used for pickling metadata and objects, default to `pickle`. + pickle_module: module used for pickling metadata and objects, default to `"pickle"`. this arg is used by `torch.save`, for more details, please check: https://pytorch.org/docs/stable/generated/torch.save.html#torch.save. pickle_protocol: can be specified to override the default protocol, default to `2`. diff --git a/monai/data/utils.py b/monai/data/utils.py index e68dd387ff..cf153eac11 100644 --- a/monai/data/utils.py +++ b/monai/data/utils.py @@ -76,8 +76,11 @@ "pad_list_data_collate", "no_collation", "convert_tables_to_dicts", + "SUPPORTED_PICKLE_MOD", ] +SUPPORTED_PICKLE_MOD = {"pickle": pickle} + def get_random_patch( dims: Sequence[int], patch_size: Sequence[int], rand_state: Optional[np.random.RandomState] = None diff --git a/tests/test_persistentdataset.py b/tests/test_persistentdataset.py index e7475df1f0..9c66d3d10c 100644 --- a/tests/test_persistentdataset.py +++ b/tests/test_persistentdataset.py @@ -61,7 +61,7 @@ def test_cache(self): data=items, transform=_InplaceXform(), cache_dir=tempdir, - pickle_module=pickle, + pickle_module="pickle", pickle_protocol=pickle.HIGHEST_PROTOCOL, ) self.assertEqual(items, [[[]], [[0]], [[0, 1]], [[0, 1, 2]], [[0, 1, 2, 3]]]) From b830c03b14e0f3cc1bc281e301bd2c401ef28ab2 Mon Sep 17 00:00:00 2001 From: Wenqi Li Date: Mon, 29 Nov 2021 19:10:16 +0000 Subject: [PATCH 2/3] update pickle Signed-off-by: Wenqi Li --- monai/data/dataset.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/monai/data/dataset.py b/monai/data/dataset.py index 675c852a06..480b1455c8 100644 --- a/monai/data/dataset.py +++ b/monai/data/dataset.py @@ -182,7 +182,7 @@ def __init__( super().__init__(data=data, transform=transform) self.cache_dir = Path(cache_dir) if cache_dir is not None else None self.hash_func = hash_func - self.pickle_module = look_up_option(pickle_module, SUPPORTED_PICKLE_MOD) + self.pickle_module = pickle_module self.pickle_protocol = pickle_protocol if self.cache_dir is not None: if not self.cache_dir.exists(): @@ -287,7 +287,7 @@ def _cachecheck(self, item_transformed): torch.save( obj=_item_transformed, f=temp_hash_file, - pickle_module=self.pickle_module, + pickle_module=look_up_option(self.pickle_module, SUPPORTED_PICKLE_MOD), pickle_protocol=self.pickle_protocol, ) if temp_hash_file.is_file() and not hashfile.is_file(): From ed3f6b3814e16ffd33cbe4da1c49615a29bcfed3 Mon Sep 17 00:00:00 2001 From: Wenqi Li Date: Mon, 29 Nov 2021 19:24:44 +0000 Subject: [PATCH 3/3] docstring Signed-off-by: Wenqi Li --- monai/data/dataset.py | 10 ++++++---- monai/data/utils.py | 1 + 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/monai/data/dataset.py b/monai/data/dataset.py index 480b1455c8..989de09c18 100644 --- a/monai/data/dataset.py +++ b/monai/data/dataset.py @@ -169,9 +169,10 @@ def __init__( If `cache_dir` is `None`, there is effectively no caching. hash_func: a callable to compute hash from data items to be cached. defaults to `monai.data.utils.pickle_hashing`. - pickle_module: module used for pickling metadata and objects, default to `"pickle"`. + pickle_module: string representing the module used for pickling metadata and objects, default to `"pickle"`. this arg is used by `torch.save`, for more details, please check: - https://pytorch.org/docs/stable/generated/torch.save.html#torch.save. + https://pytorch.org/docs/stable/generated/torch.save.html#torch.save, + and ``monai.data.utils.SUPPORTED_PICKLE_MOD``. pickle_protocol: can be specified to override the default protocol, default to `2`. this arg is used by `torch.save`, for more details, please check: https://pytorch.org/docs/stable/generated/torch.save.html#torch.save. @@ -335,9 +336,10 @@ def __init__( If `cache_dir` is `None`, there is effectively no caching. hash_func: a callable to compute hash from data items to be cached. defaults to `monai.data.utils.pickle_hashing`. - pickle_module: module used for pickling metadata and objects, default to `"pickle"`. + pickle_module: string representing the module used for pickling metadata and objects, default to `"pickle"`. this arg is used by `torch.save`, for more details, please check: - https://pytorch.org/docs/stable/generated/torch.save.html#torch.save. + https://pytorch.org/docs/stable/generated/torch.save.html#torch.save, + and ``monai.data.utils.SUPPORTED_PICKLE_MOD``. pickle_protocol: can be specified to override the default protocol, default to `2`. this arg is used by `torch.save`, for more details, please check: https://pytorch.org/docs/stable/generated/torch.save.html#torch.save. diff --git a/monai/data/utils.py b/monai/data/utils.py index cf153eac11..61c773c204 100644 --- a/monai/data/utils.py +++ b/monai/data/utils.py @@ -79,6 +79,7 @@ "SUPPORTED_PICKLE_MOD", ] +# module to be used by `torch.save` SUPPORTED_PICKLE_MOD = {"pickle": pickle}