提交 4c5d01fa 编写于 作者: M Megvii Engine Team 提交者: Xinran Xu

refactor(mge/examples): add example to use modified sublinear API

GitOrigin-RevId: c6f6c65d0bf2011bb8590aec3285bbac99fc9c34
上级 396653df
......@@ -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)
......
......@@ -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
......@@ -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,
)
......@@ -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
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册