megengine.jit package

megengine.jit.sublinear_memory_config

class megengine.jit.sublinear_memory_config.SublinearMemoryConfig(thresh_nr_try=10, genetic_nr_iter=0, genetic_pool_size=20, lb_memory=0, num_worker=2)[源代码]

基类:object

亚线性内存优化的配置。

参数
  • thresh_nr_try (int) – 线性空间以及亚线性内存优化的当前范围搜索的样本数目。默认:10。也可以通过环境变量 ‘MGB_SUBLINEAR_MEMORY_THRESH_NR_TRY’ 进行设置。

  • genetic_nr_iter (int) – 使用遗传算法寻找最优切分策略时的迭代轮数。默认:0。也可以通过环境变量 ‘MGB_SUBLINEAR_MEMORY_GENETIC_NR_ITER’ 进行设置。

  • genetic_pool_size (int) – 遗传优化算法进行交叉随机选择(crossover)时所使用的样本数。默认:20。也可以通过环境变量 ‘MGB_SUBLINEAR_MEMORY_GENETIC_POOL_SIZE’ 进行设置。

  • lb_memory (int) – 次线性内存优化瓶颈大小的下界(以MB为单位)。它可用于在内存和速度之间进行手动权衡。默认:0。也可以通过设置环境变量 ‘MGB_SUBLINEAR_MEMORY_LOWER_BOUND_MB’ 来实现。

  • num_worker (int) – 搜索次线性内存优化最优切分策略时使用的线程数。默认:当前系统中CPU数目的一半。注意:该参数值需要大于等于1。也可以通过设置环境变量 ‘MGB_SUBLINEAR_MEMORY_WORKERS’来实现。

注意,为了使上述环境变量生效,需要将环境变量 MGB_COMP_GRAPH_OPT 设置为 ‘enable_sublinear_memory_opt=1’ 。

megengine.jit tracing

class megengine.jit.tracing.BrokenRawTensor(handle=None)[源代码]

基类:megengine.core.tensor.raw_tensor.RawTensor

class megengine.jit.tracing.CompiledTensorProxy(handle)[源代码]

基类:megengine.core.tensor.raw_tensor.RawTensor

鸭子类型的RawTensor

property device
property dtype
numpy()[源代码]
property shape
class megengine.jit.tracing.LazyEvalTensor(varnode)[源代码]

基类:megengine.core.tensor.raw_tensor.RawTensor

property device
property dtype
numpy()[源代码]
property shape
class megengine.jit.tracing.TensorInfo[源代码]

基类:object

__slots__ = ('external', 'exported', 'data_read', 'shape_read', 'value_read', 'device', 'dtype', 'shape', 'bound_data', 'varnode', 'data_setter', 'shape_reader', 'value_reader', 'data_reader')
bound_data
data_read
data_reader
data_setter
device
dtype
exported
external
shape
shape_read
shape_reader
value_read
value_reader
varnode
exception megengine.jit.tracing.TraceMismatchError[源代码]

基类:RuntimeError

class megengine.jit.tracing.TraceMixin[源代码]

基类:object

numpy()[源代码]
property shape
class megengine.jit.tracing.TracedLazyTensor(varnode)[源代码]

基类:megengine.jit.tracing.TraceMixin, megengine.jit.tracing.LazyEvalTensor

class megengine.jit.tracing.TracedRawTensor(handle=None)[源代码]

基类:megengine.jit.tracing.TraceMixin, megengine.core.tensor.raw_tensor.RawTensor

megengine.jit.tracing.apply_compiled_mode(op, *args)[源代码]
megengine.jit.tracing.apply_const_symbolic_mode(op, *args)[源代码]
megengine.jit.tracing.apply_const_with_tracing(op, *args)[源代码]
megengine.jit.tracing.apply_symbolic_mode(op, *args)[源代码]
megengine.jit.tracing.apply_with_tracing(op, *args)[源代码]
megengine.jit.tracing.assign_raw_tensor(lhs, rhs)[源代码]
megengine.jit.tracing.exclude_from_trace()[源代码]
megengine.jit.tracing.find_raw_tensor(x)[源代码]
megengine.jit.tracing.find_raw_tensor(x: megengine.core.tensor.raw_tensor.RawTensor)
megengine.jit.tracing.find_raw_tensor(x: megengine.core.tensor.core.TensorWrapperBase)
megengine.jit.tracing.find_raw_tensor(x: megengine.core.tensor.tensor.Tensor)
megengine.jit.tracing.is_tracing()[源代码]
class megengine.jit.tracing.trace(function, symbolic=False, capture_as_const=False, sublinear_memory_config=None, profiling=False, opt_level=None, tensor_shape=True)[源代码]

基类:object

包装一个可调用对象,并提供以下功能:

  • 通过 trace()dump() 实现追溯(tracing)

  • 通过调用 __call__() 加速执行

参数
  • function – 将被追溯的函数。

  • symbolic – 是否为追溯使用符号执行。默认:False

  • capture_as_const – 将全局变量或闭包捕获为常值。默认:False

  • sublinear_memory_config (Optional[SublinearMemoryConfig]) – 配置亚线性内存优化。如果不是None,则使用给定设置进行亚线性内存优化。

  • profiling (bool) – 是否对编译好的函数追溯(trace)进行性能评估(profile)。默认:False

  • opt_level (Optional[int]) – 编译追踪函数时的优化级别。

  • symbolic_shape – 是否为追溯使用符号形状。默认:True

dump(file, *, arg_names=None, output_names=None, append=False, optimize_for_inference=True, **kwargs)[源代码]

序列化被追溯 (trace) 的模型并保存到文件。

参数
  • file – 输出文件,可以是文件对象或文件名

  • arg_names – 被追溯(traced)函数的输入张量的名字。

  • output_names – 被追溯(traced)函数的输出张量的名字,如果未指明则使用默认名字。

  • append – 是否在 file 后追加输出。仅当 file 是文件名时可用。

  • optimize_for_inference – 打开推理优化,如果是False则关闭所有优化选项。默认:True

关键字参数
  • enable_io16xc32 –

    是否使用float16作为算子间I/O的数据精度,同时float32作为内部计算的数据精度。注意输出变量的类型也随之更改为float16。

  • enable_ioc16 –

    是否使用float16同时作为算子间I/O和内部计算的数据精度。

  • enable_hwcd4 –

    是否使用NHWCD4数据格式。在某些OpenCL设备上,会提高计算速度。

  • enable_nchw88 –

    是否使用NCHW88数据格式。当前用于X86 AVX后端。

  • enable_nchw44 –

    是否使用NCHW44数据格式。当前用于arm后端。

  • enable_nchw44_dot –

    是否使用NCHW4_dot数据格式。当前用于armv8.2+dotprod后端。

  • enable_nchw4 –

    是否使用NCHW4数据格式。当前用于nvidia后端(基于cudnn)。

  • enable_nchw32 –

    是否使用NCHW32数据格式。当前与tensorcore用于nvidia后端(基于cudnn)。

  • enable_chwn4 –

    是否使用CHWN4数据格式。当前与tensorcore用于nvidia后端。

  • enable_fuse_conv_bias_nonlinearity:是否融合 conv bias nonlinearty 算子

    成一个算子。

  • enable_fuse_conv_bias_with_z:

    推理阶段是否在nvidia后端对输入z融合 CONV + bias 成一个算子(这个优化会导致训练和推理的输出精度不一致)

get_profile()[源代码]

获取被追溯(trace)函数在编译后运行的性能结果。

返回

一个兼容json的对象。

trace(*args, **kwargs)[源代码]