diff --git a/imperative/python/megengine/__init__.py b/imperative/python/megengine/__init__.py index 95f32e9ad3835f4de61a47f779e9ce328f5c5fc0..fc352f11aa84e51adb9cfdd2b6a4a4614b15c17c 100644 --- a/imperative/python/megengine/__init__.py +++ b/imperative/python/megengine/__init__.py @@ -76,7 +76,6 @@ from .core._imperative_rt.core2 import full_sync as _full_sync from .core._imperative_rt.core2 import sync as _sync from .core._imperative_rt.utils import _set_fork_exec_path_for_timed_func from .device import * -from .dtr import * from .logger import enable_debug_log, get_logger, set_log_file, set_log_level from .serialization import load, save from .tensor import Parameter, Tensor, tensor @@ -100,6 +99,7 @@ del _set_fork_exec_path_for_timed_func import megengine.autodiff import megengine.data import megengine.distributed +import megengine.dtr import megengine.functional import megengine.hub import megengine.jit diff --git a/imperative/python/megengine/dtr.py b/imperative/python/megengine/dtr.py index d586b69e0f73a3faab30a15e1eb163f5fb1d8473..928a799b2ef4a3f43750339b57f4f7c39e1d77b9 100644 --- a/imperative/python/megengine/dtr.py +++ b/imperative/python/megengine/dtr.py @@ -11,14 +11,14 @@ from typing import Union from mprop import mproperty -from .core._imperative_rt.core2 import set_option +from .core._imperative_rt.core2 import set_option as _set_option from .core._imperative_rt.utils import _set_defrag _eviction_threshold = 0 _evictee_minimum_size = 1024 ** 2 -def str2bytes(text: str) -> int: +def _str2bytes(text: str) -> int: regex = re.compile(r"(\d+(?:\.\d+)?)\s*([kmg]?b)", re.IGNORECASE) order = ["b", "kb", "mb", "gb"] result = regex.findall(text) @@ -32,7 +32,9 @@ def str2bytes(text: str) -> int: @mproperty def eviction_threshold(mod): r""" - Returns the eviction threshold in bytes. + Get or set the eviction threshold in bytes. It can also be set to a string, + whose formatting supports byte(B), kilobyte(KB), megabyte(MB) and + gigabyte(GB) units. .. note:: @@ -40,40 +42,34 @@ def eviction_threshold(mod): and evict resident tensors until the amount of used memory falls below this threshold. - """ - return mod._eviction_threshold - - -@eviction_threshold.setter -def eviction_threshold(mod, value: Union[int, str]): - r""" - Change the eviction threshold. If `value` is an int, it represents the - number of bytes. If `value` is a string, its formatting supports bytes(B), - kilobyte(KB), megabyte(MB) and gigabyte(GB) units. - Examples: .. code-block:: import megengine as mge - mge.dtr.eviction_threshold = 2 * 1024 ** 3 mge.dtr.eviction_threshold = "2GB" - mge.dtr.eviction_threshold = "2048MB" """ + return mod._eviction_threshold + + +@eviction_threshold.setter +def eviction_threshold(mod, value: Union[int, str]): if isinstance(value, str): - mod._eviction_threshold = mod.str2bytes(value) + mod._eviction_threshold = mod._str2bytes(value) elif isinstance(value, int): mod._eviction_threshold = value else: raise TypeError("`value` should be a str or an int") - set_option("dtr_eviction_threshold", mod._eviction_threshold) + _set_option("dtr_eviction_threshold", mod._eviction_threshold) @mproperty def evictee_minimum_size(mod): r""" - Returns the memory threshold of tensors in bytes. + Get or set the memory threshold of tensors in bytes. It can also be set to a + string, whose formatting supports byte(B), kilobyte(KB), megabyte(MB) and + gigabyte(GB) units. .. note:: @@ -81,34 +77,26 @@ def evictee_minimum_size(mod): candidate set. A tensor that is not added to the candidate set will never be evicted during its lifetime. - """ - return mod._evictee_minimum_size - - -@evictee_minimum_size.setter -def evictee_minimum_size(mod, value: Union[int, str]): - r""" - Change the memory threshold of tensors. If `value` is an int, it represents - the number of bytes. If `value` is a string, its formatting supports bytes(B), - kilobyte(KB), megabyte(MB) and gigabyte(GB) units. - Examples: .. code-block:: import megengine as mge - mge.dtr.evictee_minimum_size = 2 * 1024 ** 2 mge.dtr.evictee_minimum_size = "2MB" - mge.dtr.evictee_minimum_size = "2048KB" """ + return mod._evictee_minimum_size + + +@evictee_minimum_size.setter +def evictee_minimum_size(mod, value: Union[int, str]): if isinstance(value, str): - mod._evictee_minimum_size = mod.str2bytes(value) + mod._evictee_minimum_size = mod._str2bytes(value) elif isinstance(value, int): mod._evictee_minimum_size = value else: raise TypeError("`value` should be a str or an int") - set_option("dtr_evictee_minimum_size", mod._evictee_minimum_size) + _set_option("dtr_evictee_minimum_size", mod._evictee_minimum_size) def enable(): @@ -116,16 +104,16 @@ def enable(): Enable to record computing path of tensors and to perform DTR policy. """ _set_defrag(True) - set_option("enable_dtr_auto_drop", 1) - set_option("enable_drop", 1) - set_option("buffer_length", 0) - set_option("record_computing_path", 1) + _set_option("enable_dtr_auto_drop", 1) + _set_option("enable_drop", 1) + _set_option("buffer_length", 0) + _set_option("record_computing_path", 1) def disable(): r""" Stop recording computing path of tensors and performing DTR policy. """ - set_option("enable_dtr_auto_drop", 0) - set_option("enable_drop", 0) - set_option("record_computing_path", 0) + _set_option("enable_dtr_auto_drop", 0) + _set_option("enable_drop", 0) + _set_option("record_computing_path", 0)