From 4c5d01fa344916bbfbf323e7ac3f12b5b9cb34ec Mon Sep 17 00:00:00 2001 From: Megvii Engine Team Date: Sun, 26 Apr 2020 17:38:06 +0800 Subject: [PATCH] refactor(mge/examples): add example to use modified sublinear API GitOrigin-RevId: c6f6c65d0bf2011bb8590aec3285bbac99fc9c34 --- python_module/megengine/jit/__init__.py | 59 +++++++++---------- .../megengine/jit/sublinear_memory_config.py | 14 ++++- .../test/integration/test_correctness.py | 21 ++----- python_module/test/unit/jit/test_jit.py | 6 +- 4 files changed, 48 insertions(+), 52 deletions(-) diff --git a/python_module/megengine/jit/__init__.py b/python_module/megengine/jit/__init__.py index cd76bbb0f..e8f063267 100644 --- a/python_module/megengine/jit/__init__.py +++ b/python_module/megengine/jit/__init__.py @@ -18,7 +18,7 @@ import megengine._internal as mgb from megengine._internal.plugin import CompGraphProfiler from ..core import Tensor, graph, tensor -from .sublinear_memory_config import SublinearMemConfig +from .sublinear_memory_config import SublinearMemoryConfig def sideeffect(f): @@ -82,8 +82,8 @@ class trace: :param symbolic: Whether to use symbolic tensor. Default: False :param opt_level: Optimization level for compiling trace. :param log_level: Log level. - :param enable_sublinear: Enable sublinear memory optimization. Default: False - :param sublinear_mem_config: Configuration for sublinear memory optimization. + :param sublinear_memory_config: Configuration for sublinear memory optimization. + If not None, it enables sublinear memory optimization with given setting. :param profiling: Whether to profile compiled trace. Default: False """ @@ -106,16 +106,14 @@ class trace: symbolic: bool = False, opt_level: int = None, log_level: int = None, - enable_sublinear: bool = False, - sublinear_mem_config: SublinearMemConfig = None, + sublinear_memory_config: SublinearMemoryConfig = None, profiling: bool = False ): self.__wrapped__ = func self._symbolic = symbolic self._graph_opt_level = opt_level self._log_level = log_level - self._enable_sublinear = enable_sublinear - self._sublinear_mem_config = sublinear_mem_config + self._sublinear_memory_config = sublinear_memory_config self._status = self._UNSTARTED self._args = None self._kwargs = None @@ -287,35 +285,34 @@ class trace: def _apply_graph_options(self, cg): # graph opt level - if not (self._graph_opt_level is None): + if self._graph_opt_level is not None: cg.set_option("graph_opt_level", self._graph_opt_level) # log level - if not (self._log_level is None): + if self._log_level is not None: cg.set_option("log_level", self._log_level) # sublinear - if self._enable_sublinear: + if self._sublinear_memory_config is not None: cg.set_option("enable_sublinear_memory_opt", True) - if not (self._sublinear_mem_config is None): - cg.set_option( - "sublinear_mem_cofig.lb_memory", - self._sublinear_mem_config.lb_memory, - ) - cg.set_option( - "sublinear_mem_cofig.genetic_nr_iter", - self._sublinear_mem_config.genetic_nr_iter, - ) - cg.set_option( - "sublinear_mem_cofig.genetic_pool_size", - self._sublinear_mem_config.genetic_pool_size, - ) - cg.set_option( - "sublinear_mem_cofig.thresh_nr_try", - self._sublinear_mem_config.thresh_nr_try, - ) - cg.set_option( - "sublinear_mem_cofig.num_worker", - self._sublinear_mem_config.num_worker, - ) + cg.set_option( + "sublinear_mem_cofig.lb_memory", + self._sublinear_memory_config.lb_memory, + ) + cg.set_option( + "sublinear_mem_cofig.genetic_nr_iter", + self._sublinear_memory_config.genetic_nr_iter, + ) + cg.set_option( + "sublinear_mem_cofig.genetic_pool_size", + self._sublinear_memory_config.genetic_pool_size, + ) + cg.set_option( + "sublinear_mem_cofig.thresh_nr_try", + self._sublinear_memory_config.thresh_nr_try, + ) + cg.set_option( + "sublinear_mem_cofig.num_worker", + self._sublinear_memory_config.num_worker, + ) # profile if self._profiling: self._profiler = CompGraphProfiler(cg) diff --git a/python_module/megengine/jit/sublinear_memory_config.py b/python_module/megengine/jit/sublinear_memory_config.py index 7603a843e..fc3e46ddd 100644 --- a/python_module/megengine/jit/sublinear_memory_config.py +++ b/python_module/megengine/jit/sublinear_memory_config.py @@ -10,7 +10,7 @@ from ..core.device import get_device_count -class SublinearMemConfig: +class SublinearMemoryConfig: r""" Configuration for sublinear memory optimization. @@ -28,7 +28,11 @@ class SublinearMemConfig: It can also be set through the environmental variable 'MGB_SUBLINEAR_MEMORY_LOWER_BOUND_MB'. :param num_worker: number of thread workers to search the optimum checkpoints in sublinear memory optimization. Default: half of cpu number in the system. + Note: the value must be greater or equal to one. It can also be set through the environmental variable 'MGB_SUBLINEAR_MEMORY_WORKERS'. + + Note that the environmental variable MGB_COMP_GRAPH_OPT must be set to 'enable_sublinear_memory_opt=1' + in order for the above environmental variable to be effective. """ def __init__( @@ -37,10 +41,16 @@ class SublinearMemConfig: genetic_nr_iter: int = 0, genetic_pool_size: int = 20, lb_memory: int = 0, - num_worker: int = get_device_count("cpu") / 2, + num_worker: int = max(1, get_device_count("cpu") // 2), ): + assert thresh_nr_try >= 0, "thresh_nr_try must be greater or equal to zero" self.thresh_nr_try = thresh_nr_try + assert genetic_nr_iter >= 0, "genetic_nr_iter must be greater or equal to zero" self.genetic_nr_iter = genetic_nr_iter + assert ( + genetic_pool_size >= 0 + ), "genetic_pool_size must be greater or equal to zero" self.genetic_pool_size = genetic_pool_size self.lb_memory = lb_memory + assert num_worker > 0, "num_worker must be greater or equal to one" self.num_worker = num_worker diff --git a/python_module/test/integration/test_correctness.py b/python_module/test/integration/test_correctness.py index 282818317..70a9861f6 100644 --- a/python_module/test/integration/test_correctness.py +++ b/python_module/test/integration/test_correctness.py @@ -17,7 +17,7 @@ import megengine as mge import megengine.functional as F from megengine import jit, tensor from megengine.functional.debug_param import set_conv_execution_strategy -from megengine.jit import SublinearMemConfig +from megengine.jit import SublinearMemoryConfig from megengine.module import AvgPool2d, BatchNorm2d, Conv2d, Linear, Module from megengine.optimizer import SGD from megengine.test import assertTensorClose @@ -132,12 +132,7 @@ def update_model(model_path): def run_test( - model_path, - use_jit, - use_symbolic, - enable_sublinear=False, - sublinear_mem_config=None, - max_err=None, + model_path, use_jit, use_symbolic, sublinear_memory_config=None, max_err=None, ): """ @@ -168,8 +163,7 @@ def run_test( train_func = jit.trace( train_func, symbolic=use_symbolic, - enable_sublinear=enable_sublinear, - sublinear_mem_config=sublinear_mem_config, + sublinear_memory_config=sublinear_memory_config, ) opt.zero_grad() @@ -199,12 +193,7 @@ def test_correctness(): run_test(model_path, True, True) # sublinear - config = SublinearMemConfig(genetic_nr_iter=10) + config = SublinearMemoryConfig(genetic_nr_iter=10) run_test( - model_path, - True, - True, - enable_sublinear=True, - sublinear_mem_config=config, - max_err=1e-5, + model_path, True, True, sublinear_memory_config=config, max_err=1e-5, ) diff --git a/python_module/test/unit/jit/test_jit.py b/python_module/test/unit/jit/test_jit.py index d5104037f..b90655aaa 100644 --- a/python_module/test/unit/jit/test_jit.py +++ b/python_module/test/unit/jit/test_jit.py @@ -18,7 +18,7 @@ import megengine._internal as mgb import megengine.module as M from megengine import jit, tensor from megengine.core.tensor import Tensor -from megengine.jit import SublinearMemConfig +from megengine.jit import SublinearMemoryConfig from megengine.test import assertTensorClose @@ -190,9 +190,9 @@ def test_dump_bn_fused(): # Simply verify the options passed down def test_sublinear(): - config = SublinearMemConfig(genetic_nr_iter=10) + config = SublinearMemoryConfig(genetic_nr_iter=10) - @jit.trace(symbolic=True, enable_sublinear=True, sublinear_mem_config=config) + @jit.trace(symbolic=True, sublinear_memory_config=config) def f(x): return x + x -- GitLab