diff --git a/mlair/data_handler/default_data_handler.py b/mlair/data_handler/default_data_handler.py index ddf276cf2d88c108d8622c507471f989c4f99e8b..07a866aec1efd43de42f918844abeb7c3bbc9524 100644 --- a/mlair/data_handler/default_data_handler.py +++ b/mlair/data_handler/default_data_handler.py @@ -8,6 +8,7 @@ import gc import logging import os import pickle +import dill import shutil from functools import reduce from typing import Tuple, Union, List @@ -86,7 +87,7 @@ class DefaultDataHandler(AbstractDataHandler): data = {"X": self._X, "Y": self._Y, "X_extreme": self._X_extreme, "Y_extreme": self._Y_extreme} data = self._force_dask_computation(data) with open(self._save_file, "wb") as f: - pickle.dump(data, f) + dill.dump(data, f) logging.debug(f"save pickle data to {self._save_file}") self._reset_data() @@ -101,7 +102,7 @@ class DefaultDataHandler(AbstractDataHandler): def _load(self): try: with open(self._save_file, "rb") as f: - data = pickle.load(f) + data = dill.load(f) logging.debug(f"load pickle data from {self._save_file}") self._X, self._Y = data["X"], data["Y"] self._X_extreme, self._Y_extreme = data["X_extreme"], data["Y_extreme"] diff --git a/mlair/data_handler/iterator.py b/mlair/data_handler/iterator.py index 30c45417a64e949b0c0535a96a20c933641fdcbb..564bf3bfd6e4f5b814c9d090733cfbfbf26a850b 100644 --- a/mlair/data_handler/iterator.py +++ b/mlair/data_handler/iterator.py @@ -9,6 +9,7 @@ import math import os import shutil import pickle +import dill from typing import Tuple, List @@ -109,7 +110,7 @@ class KerasIterator(keras.utils.Sequence): """Load pickle data from disk.""" file = self._path % index with open(file, "rb") as f: - data = pickle.load(f) + data = dill.load(f) return data["X"], data["Y"] @staticmethod @@ -167,7 +168,7 @@ class KerasIterator(keras.utils.Sequence): data = {"X": X, "Y": Y} file = self._path % index with open(file, "wb") as f: - pickle.dump(data, f) + dill.dump(data, f) def _get_number_of_mini_batches(self, number_of_samples: int) -> int: """Return number of mini batches as the floored ration of number of samples to batch size."""