提交 8cf01ede authored 作者: Frederic Bastien's avatar Frederic Bastien

Added a Theano flag to disable profiling of memory used by Theano function.

上级 5791145d
......@@ -5,7 +5,7 @@ from theano.gof.cutils import run_cthunk
from theano.compile.mode import Mode, register_mode, predefined_modes, predefined_linkers, predefined_optimizers
from theano.gof.python25 import any
from theano import gof
from theano.configparser import config, AddConfigVar, IntParam
from theano.configparser import config, AddConfigVar, IntParam, BoolParam
from theano.compile.function_module import FunctionMaker
import_time = time.time()
......@@ -23,6 +23,10 @@ AddConfigVar('ProfileMode.min_memory_size',
of their outputs (in bytes) is lower then this threshold""",
IntParam(1024, lambda i: i >= 0))
AddConfigVar('ProfileMode.profile_memory',
"""Enable profiling of memory used by Theano functions""",
BoolParam(True))
class Profile_Maker(FunctionMaker):
def create(self, input_storage=None, trustme=False):
ret = super(Profile_Maker,self).create(input_storage, trustme)
......@@ -142,7 +146,11 @@ class ProfileMode(Mode):
if isinstance(linker, str) or linker is None:
linker = predefined_linkers[linker]
linker = WrapLinker([linker], profile_thunk2)
if config.ProfileMode.profile_memory:
p_thunk = profile_thunk
else:
p_thunk = profile_thunk2
linker = WrapLinker([linker], p_thunk)
self.linker = linker
if isinstance(optimizer, str) or optimizer is None:
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论