MegEngine Lite python 接口介绍

LiteTensor 相关 API

LiteLayout

def __init__(self, shape=None, dtype=None)

指定 shape 和 dtype,将构造出一个 LiteLayout。

参数:

  • shape:LiteLayout 中的 shape 信息。

  • dtype:LiteLayout 中的数据类型,这个 dtype 可以为如下类型:

    • 字符:”int32”,”float32”,”uint8”,”int8”,”int16”,”uint16”,”float16”。

    • LiteDataType:LITE_FLOAT,LITE_HALF,LITE_INT,LITE_INT16,LITE_INT8,LITE_UINT8,LITE_UINT16。

    • numpy dtype 实例:np.dtype(“int32”),np.dtype(“float32”),np.dtype(“uint8”),np.dtype(“int8”),np.dtype(“int16”),np.dtype(“uint16”),np.dtype(“float16”)。

    • numpy dtype:numpy.int32,numpy.float32,numpy.uint8,numpy.int8,numpy.int16,numpy.uint16,numpy.float16。

LiteTensor

def __init__(
    self,
    layout=None,
    device_type=LiteDeviceType.LITE_CPU,
    device_id=0,
    is_pinned_host=False,
    shapes=None,
    dtype=None,
):

构造 LiteTensor 时候可以指定 layout, device_type,device_id,以及是否内存是 锁页内存,另外 如果也可以通过传递关键词参数 shapes 和 dtype 在 MegEngineLite 内部构造 layout,让后再构造 Tensor。

参数

  • layout:LiteTensor 对应的 layout 信息。

  • device_type:LiteDeviceType 对应的数据类型,包含:

    class LiteDeviceType(IntEnum):
    LITE_CPU = 0
    LITE_CUDA = 1
    LITE_ATLAS = 3
    LITE_NPU = 4
    LITE_DEVICE_DEFAULT = 5
    
  • device_id:LiteTensor 所在设备的 id。

  • is_pinned_host:LiteTensor 是否为 锁页内存

警告

LiteTensor 构造之后,内存没有立即申请,只有当 LiteTensor 需要用到内存时候才会申请。

示例:

# 直接从 shapes 创建 LiteTensor
tensor_cuda2 = LiteTensor(shapes=[4,16], dtype="float32", device_type=LiteDeviceType.LITE_CUDA)
# 从 layout 创建 LiteTensor
layout = LiteLayout([4, 16], "float32")
tensor = LiteTensor(layout, LiteDeviceType.LITE_CPU)
tensor_cuda = LiteTensor(layout=layout, device_type=LiteDeviceType.LITE_CUDA)

LiteTensor 信息获取

# 获取或者设置 LiteTensor 的 layout 信息
@property
def layout(self):
@layout.setter
def layout(self, layout):

# 获取 LiteTensor 是否是锁页内存
@property
def is_pinned_host(self):

# 获取 LiteTensor 所在的设备类型
@property
def device_type(self):

# 获取 LiteTensor 所在的设备 id
@property
def device_id(self):

# 获取 LiteTensor 的内存是否是连续的
@property
def is_continue(self):

# 获取 LiteTensor 的内存的大小,单位是字节
@property
def nbytes(self):

注解

上面 LiteTensor 的 layout 信息具有装饰器 @property 和 @layout.setter,可以直接作为成员一样访问和赋值, 其他信息都具有 @property 的装饰器,因此都可以通过成员一样的访问。

get_ctypes_memory

def get_ctypes_memory(self)
  • get_ctypes_memory:将返回 ctypes.c_void_p 类型,其指向 Tensor 的内存地址,如果 Tensor 没有申请内存,将会申请内存。

reshape

def reshape(self, shape):

改变这个 LiteTensor 的 LiteLayout 中的 shape 为新的 shape,其中 新的 shape 中元素个数需要和老的 shape 里面的元素个数相等

slice

def slice(self, start, end, step=None):

对 LiteTensor 进行切片,返回一个新的 LiteTensor,新的 LiteTensor 和原来 LiteTensor 共享内存, 新的 LiteTensor 可能不连续

参数: start,end 的长度必须相等,长度可以小于 Tensor 的 Layout 的维度,如果传递了 step,则 step 也需要和 start,end 的长度相等

  • start:Tensor 每一维度的起始 index 组成的数组,从高维到低维。

  • end:Tensor 每一维度的结束 index 组成的数组,从高维到低维。

  • step:Tensor 每一维度切片的间距,从高维到低维,默认为1。

返回值:返回一个新的 LiteTensor。

示例:

layout = LiteLayout([4, 8], "int32")
tensor1 = LiteTensor(layout)

tensor1.set_data_by_copy([i for i in range(32)])
real_data_org = tensor1.to_numpy()

tensor2 = tensor1.slice([1, 4], [3, 8])
assert tensor2.layout.shapes[0] == 2
assert tensor2.layout.shapes[1] == 4
assert tensor2.is_continue == False

real_data = tensor2.to_numpy()
for i in range(8):
    row = i // 4
    col = i % 4
    assert real_data[row][col] == real_data_org[row + 1][col + 4]

fill_zero

def fill_zero(self):

将 LiteTensor 内存里面的数据全部设置为 0。

copy_from

def copy_from(self, src_tensor):

从 src_Tensor 中拷贝数据到自己内存中, 如果 src_tensor 和自己的 layout 不相同时,会更改自身 Layout 信息为 src layout

share_memory_with

def share_memory_with(self, src_tensor):

将会和 src_tensor 共享内存数据, 如果 src_tensor 和自己的 LiteTensor 信息(layout,device_type,device_id等)不相同时,会更改自身信息为 src 的信息

示例:

layout = LiteLayout([4, 8], "int16")
tensor1 = LiteTensor(layout)
tensor2 = LiteTensor(layout)

tensor1.set_data_by_copy([i for i in range(32)])
tensor2.share_memory_with(tensor1)
real_data = tensor2.to_numpy()
for i in range(32):
    assert real_data[i // 8][i % 8] == i

update

def update(self):

将 LiteTensor 底层的信息更新到 python 中的 LiteTensor 中,包括 LiteTensor 的设备,设备 id,layout等信息。

set_data_by_copy

def set_data_by_copy(self, data, data_length=0, layout=None):

将用户指定的 data 以 复制的方式 到该 LiteTensor 中。

参数:

  • data: data 可以是 list 或者 numpy ndarray 或者 ctypes 的 c_void_p。

    • 当 data 类型为 list 时候,LiteTensor 的 Layout 不会被修改,用户需要保证 tensor 的内存大小大于 list 的长度。

    • 当 data 为 numpy ndarray 时候,如果 data 的长度和 LiteTensor 的内存大小不等时,将修改 LiteTensor 的 layout 为 data 的 layout。

    • 当 data 为 ctypes 的 c_void_p 时候,用户要么设置 data_length 并且必须 data_length LiteTensor 的长度相等,要么设置新的 Layout。

  • data_length:当用户输入的 data 为 ctypes 的 c_void_p 时候,指明数据长度。

  • layout 当需要改变 LiteTensor 的 layout 时,可以通过这个接口传递新的 layout。

警告

  • LiteTensor 必须是 锁页内存 或者是 CPU 上的内存

示例:

layout = LiteLayout([2, 16], "int8")
tensor = LiteTensor(layout)

data = [i for i in range(32)]
tensor.set_data_by_copy(data)
real_data = tensor.to_numpy()
for i in range(32):
    assert real_data[i // 16][i % 16] == i

set_data_by_share

def set_data_by_share(self, data, length=0, layout=None):

将用户传递进来的 data 通过 共享的方式 保存在 LiteTensor 中,避免 copy 带来的性能影响。

参数:

  • data: data 可以是 numpy ndarray 或者 ctypes 的 c_void_p。

    • 当 data 为 numpy ndarray 时候,如果 data 的长度和 LiteTensor 的内存大小不等时,将修改 LiteTensor 的 layout 为 data 的 layout。

    • 当 data 为 ctypes 的 c_void_p 时候,用户要么设置 data_length 并且必须 data_length LiteTensor 的长度相等,要么设置新的 Layout。

  • data_length:当用户输入的 data 为 ctypes 的 c_void_p 时候,指明数据长度。

  • layout 当需要改变 LiteTensor 的 layout 时,可以通过这个接口传递新的 layout。

警告

  • 当 data 为 numpy ndarray 时候,LiteTensor 要么是 锁页内存 要么是 CPU 上的内存

  • 当 data 为 ctypes 的 c_void_p 时候,对 LiteTensor 没有要求,这时候需要用户自己保证内存的设备属性。

示例:

layout = LiteLayout([2, 16], "int8")
tensor = LiteTensor(layout)
arr = np.ones([2, 16], "int8")
for i in range(32):
    arr[i // 16][i % 16] = i
tensor.set_data_by_share(arr)
real_data = tensor.to_numpy()
for i in range(32):
    assert real_data[i // 16][i % 16] == i

get_data_by_share

def get_data_by_share(self):

将 LiteTensor 中的数据以 共享 的方式构建一个 numpy 的数组,并返回给用户, 当这个 LiteTensor 中内存数据被修改时候,返回的这个 numpy 数组中的 数据也将会被修改

返回值:

  • 返回一个和 LiteTensor 共享内存的 numpy ndarray。

警告

  • 当这个 LiteTensor 中内存数据被修改时候,返回的这个 numpy 数组中的数据也将会被修改,如:

  • 当第一次 Network Forward 之后通过输出 LiteTensor 获得的这个 numpy 数组,在下一次 Network Forward 的时候会被修改

示例:

layout = LiteLayout([4, 32], "int16")
tensor = LiteTensor(layout)
assert tensor.nbytes == 4 * 32 * 2

arr = np.ones([4, 32], "int16")
for i in range(128):
    arr[i // 32][i % 32] = i
tensor.set_data_by_copy(arr)
test_data = tensor.get_data_by_share()

for i in range(128):
    assert test_data[i // 32][i % 32] == i

arr[1][18] = 5
arr[3][7] = 345
tensor.set_data_by_copy(arr)
assert test_data[1][18] == 5
assert test_data[3][7] == 345

to_numpy

def to_numpy(self):

将 LiteTensor 中数据 copy 到一个 numpy 的 ndarray 中,可以方便查看 LiteTensor 中的数据。

注解

  • 当 LiteTensor 是 锁页内存 或者是 CPU 上的 LiteTensor,则会直接 copy 到 numpy ndarray 中

  • 当 LiteTensor 在其他设备上,这时会先 copy 到 CPU LiteTensor 中,再从新的 LiteTensor copy 到 numpy ndarray 中,所以可能有 性能问题

LiteOptions

_fields_ = [
    ("weight_preprocess", c_int),
    ("fuse_preprocess", c_int),
    ("fake_next_exec", c_int),
    ("var_sanity_check_first_run", c_int),
    ("const_shape", c_int),
    ("force_dynamic_alloc", c_int),
    ("force_output_dynamic_alloc", c_int),
    ("force_output_use_user_specified_memory", c_int),
    ("no_profiling_on_shape_change", c_int),
    ("jit_level", c_int),
    ("comp_node_seq_record_level", c_int),
    ("graph_opt_level", c_int),
    ("async_exec_level", c_int),
    # layout transform options
    ("enable_nchw44", c_int),
    ("enable_nchw44_dot", c_int),
    ("enable_nchw88", c_int),
    ("enable_nhwcd4", c_int),
    ("enable_nchw4", c_int),
    ("enable_nchw32", c_int),
    ("enable_nchw64", c_int),
]

LiteOptions 是一个包含 MegEngine Network 优化选项集合的结构体,每个选项的解释如下:

  • weight_preprocess:在推理时候,部分 Kernel 执行前需要对权重进行转换,或者 Relayout,开启这个选项之后,将权重处理放到 Kernel 执行之前, 优化 Kernel 执行时间,但是 Network 初始化时间变长。

  • fuse_preprocess:开启该选项之后,模型中的部分前后处理 Operator 将会被融合在一起,优化模型执行的性能。

  • fake_next_exec:下一次执行 Inference 时候,是否为假的执行:仅仅完成内存分配等和计算无关的操作。这次假的执行完成之后将被设置为 false。

  • var_sanity_check_first_run:第一次执行 Inference 时候是否需要对每一个 Operator 的输入输出 Tensor 的正确性进行检查,默认为 true。

  • const_shape:指定 Network 的输入 shape 不会变化,这样不用在后面的执行时检查是否需要重新分配内存等操作。

  • force_dynamic_alloc:强制要求所有的 Tensor 都是运行时动态分配,且不进行内存优化,MegEngine 默认所有的 Tensor 都是执行前进行内存优化并静态申请。

  • force_output_dynamic_alloc:强制最后输出的 Tensor 的内存为动态申请,这样输出 Tensor 不用 copy 到用户的内存中,可以直接代理到返回内存给用户。

  • force_output_use_user_specified_memory:强制让输出 Tensor 的内存由用户指定,这样输出 Tensor 将不需要 copy 到用户内存,在最后一个 Kernel 计算时就写到了用户的内存地址中。

  • no_profiling_on_shape_change:当 Network 的输入 Tensor 的 shape 改变的时候,这时候 fast-run 将不会进行重新搜索最优的 kernel 算法实现。

  • jit_level:JIT 的级别,设置为 0 时:将关闭 JIT,设置为 1 时:仅仅只开启基本的 elemwise 的 JIT,当是指为 2 时:将开启 elemwise 和 reduce Operator 的 JIT。

  • comp_node_seq_record_level:设置 MegEngine 的录制模式,当设置为 0 时:将不开启录制模式,设置为 1 时:将开启录制模式,不会析构这个计算图结构,当设置为 2 时:将开启录制模式,并释放掉整个计算图。

  • graph_opt_level:设置图优化等级,当设置为 0 时:关闭图优化,当设置为 1 时:算术计算 inplace 优化,当设置为 2 时:在 1 的基础上在加上全局优化,当设置为 3 时:在 2 的基础上再使能 JIT。

  • enable_xxxx:开启对应的 layout 转换优化,不同的平台上不同的 layout 性能差异较大,见下表:

参数

作用

适用平台

enable-nchw88

将输入nchw layout的模型转为nchw88 layout的模型

X86 avx256

enable-nchw44

将输入nchw layout的模型转为nchw44 layout的模型

Arm float32

enable-nchw44-dot

将输入nchw layout的模型转为nchw44-dot layout的模型

Arm V8.2

enable-nchw4

将输入nchw layout的模型转为nchw4 layout的模型

CUDA

enable-chwn4

将输入nchw layout的模型转为chwn4 layout的模型

CUDA

enable-nchw32

将输入nchw layout的模型转为nchw32 layout的模型

CUDA

enable-nhwcd4

将输入nchw layout的模型转为nhcw4 layout的模型

移动平台GPU

LiteConfig

_fields_ = [
    ("has_compression", c_int),
    ("device_id", c_int),
    ("device_type", c_int),
    ("backend", c_int),
    ("bare_model_cryption_name", c_char_p),
    ("options", LiteOptions),
]
  • has_compression: 模型是否压缩过。

  • device_id: LiteNetwork 创建所在的设备 id。

  • device_type:LiteNetwork 创建所在的设备类型。

  • backend:指运行 LiteNetwork 的后端推理框架,目前默认是:MegEngine。

  • bare_model_cryption_name:如果模型有加密,则指明加密算法的名字,如果没有加密,则不用配置。

  • options 模型的优化参数,如上面所示。

LiteIO

_fields_ = [
    ("name", c_char_p),
    ("is_host", c_int),
    ("io_type", c_int),
    ("config_layout", LiteLayout),
]

LiteIO 为指定模型中输入输出 LiteTensor 所在的位置,可以在 device 端,也可以配置在 CPU 端,如果不配置,默认为 CPU 端。

  • name:LiteTensor 的名字,字符串。

  • is_host:LiteNetwork 创建所在的设备 id。

  • io_type:指定该 LiteTensor 对应的IO类型,目前支持两种类型,分别是:LITE_IO_VALUE 和 LITE_IO_SHAPE,默认为 LITE_IO_VALUE 。

  • config_layout:提前配置好的 layout,不配置默认为模型中的 layout。

LiteNetworkIO

def __init__(self, inputs=None, outputs=None):

LiteNetworkIO 是 LiteNetwork 构造时候的 IO 信息的集合,包含 inputs 和 outputs,为用户指定的上述 LiteIO,另外用户可以通过 add_input,add_output 接口添加 LiteIO 到 LiteNetworkIO 中。

示例:

input_io1 = LiteIO("data1", is_host=False, io_type=LiteIOType.LITE_IO_VALUE)
input_io2 = LiteIO(
    "data2",
    is_host=True,
    io_type=LiteIOType.LITE_IO_SHAPE,
    layout=LiteLayout([2, 4, 4]),
)
io = LiteNetworkIO([input_io1, input_io2])

io.add_output("out1", is_host=False)
io.add_output("out2", is_host=True, layout=LiteLayout([1, 1000]))

assert len(io.inputs) == 2
assert len(io.outputs) == 2

LiteNetwork 相关 API

LiteNetwork

def __init__(self, config=None, io=None):

构造一个 LiteNetwork,可以传递两个参数分别是 config 和 io。

参数:

  • config:模型优化需要的 LiteConfig 类型配置,默认为 None。

  • io: LiteNetworkIO 类型,指定用户输入输出 LiteTensor 的信息。

示例:

config = LiteConfig()
config.options.var_sanity_check_first_run = 0
config.device_type = LiteDeviceType.LITE_CUDA

ios = LiteNetworkIO(inputs=[LiteIO("data", False)])
network = LiteNetwork(config=config, io=ios)

load

def load(self, path):

指定创建 LiteNetwork 的模型路径,并解析这个模型,加载到内存中。

forward

def forward(self):

对指定创建 LiteNetwork 进行 forward。

wait

def wait(self):

等待指定创建 LiteNetwork 进行 forward 完成。

获取 LiteNetwork 相关信息

# 获取 LiteNetwork 的运行所在的设备 id
@property
def device_id(self):

# 获取 LiteNetwork 的运行所在的执行流 id
@property
def stream_id(self):

# 获取 LiteNetwork 的运行在 CPU 多线程时候的线程个数
@property
def threads_number(self):

# 获取 LiteNetwork 中输入 LiteTensor 中第 index 个的名字
def get_input_name(self, index):

# 获取 LiteNetwork 中输出 LiteTensor 中第 index 个的名字
def get_output_name(self, index):

# 获取 LiteNetwork 中所有输入 LiteTensor 的名字,返回一个 list
def get_all_input_name(self):

# 获取 LiteNetwork 中所有输出 LiteTensor 的名字,返回一个 list
def get_all_output_name(self):

# 获取 LiteNetwork 运行时候需要的内存信息,并将内存信息 dump 到 log_dir 指定的目录下
def get_static_memory_alloc_info(self, log_dir="logs/test"):

# 获取该 LiteNetwork 在 CPU 上运行时,是否为 inplace 模式
def is_cpu_inplace_mode(self):

注解

inplace 模式为:运行模型时候只有一个线程,这个线程发送 Kernel 任务的同时,inplace 地将 kernel 执行计算任务。非 inplace 模式:将有2个线程,一个线程发送 Kernel 任务,一个线程执行 Kernel 任务。在一些单核处理器 或者低端 cpu 上,设置 inplace 模式性能会好一些

设置 LiteNetwork 相关信息

# 设置模型运行使用的设备 id
@device_id.setter
def device_id(self, device_id):

# 设置模型运行使用的执行流 id
@stream_id.setter
def stream_id(self, stream_id):

# 如果模型执行在 CPU 多线程的情况下,设置模型运行时候需要的线程数量
@threads_number.setter
def threads_number(self, nr_threads):

# 如果模型在 CPU 上执行,设置模型运行模式为:inplace 模式
def enable_cpu_inplace_mode(self):

# 设置模型运行使用 TensorRT 进行推理
def use_tensorrt(self):

警告

上面这些 LiteNetwork 的运行时的信息设置需要在 LiteNetwork 创建之后,模型 load 之前进行设置,否则将报错。

get_io_tensor

def get_io_tensor(self, name, phase=LiteTensorPhase.LITE_IO):

获取 LiteNetwork 中名字为 name 的输入或者输出 LiteTensor。

参数:

  • name:字符串,指定输入或者输出 LiteTensor 的名字。

  • phase:当有输入和输出 LiteTensor 名字重复时候,指明获取的 LiteTensor 来自输入或者输出,可以设置为:

    • LiteTensorPhase.LITE_IO:在输入和输出的所有 LiteTensor 中寻找指定 name 的 LiteTensor,名字不会重复的情况下。

    • LiteTensorPhase.LITE_INPUT:在输入的所有 LiteTensor 中寻找指定 name 的 LiteTensor。

    • LiteTensorPhase.LITE_OUTPUT:在输出的所有 LiteTensor 中寻找指定 name 的 LiteTensor。

share_weights_with

def share_weights_with(self, src_network):

设置 LiteNetwork 运行和 src_network 共享同一份权重,两个 LiteNetwork 可以对不同的输入数据进行推理,也可以同时运行。

share_runtime_memroy

def share_runtime_memroy(self, src_network):

设置 LiteNetwork 运行和 src_network 共享运行时候的内存, 这时 self 和 src_network 不能同时执行, 运行时内存指:除了保存模型 weights 和图结构以外的所有需要的运行时内存。

async_with_callback

def async_with_callback(self, async_callback):

设置模型 forward 运行在异步模式,异步模式中,主线程将不会被阻塞,当 LiteNetwork 执行完成之后将执行 async_callback,告诉主线程执行完成。

示例:

count = 0
finished = False

def async_callback():
    nonlocal finished
    finished = True
    return 0

config = LiteConfig()
config.options.var_sanity_check_first_run = 0

network = LiteNetwork(config=config)
network.load(model_path)

network.async_with_callback(async_callback)

network.forward()

while not finished:
    count += 1

assert count > 0
output_data = output_tensor.to_numpy()

set_start_callback

def set_start_callback(self, start_callback):

设置模型运行之前的回调函数,用户可以通过这个回调函数检查输入数据是否满足要求。

set_finish_callback

def set_finish_callback(self, finish_callback):

设置模型运行之后的回调函数,用户可以通过这个回调函数检查输出数据是否满足要求。

示例:

network = LiteNetwork()
network.load(model_path)
finish_checked = False

def finish_callback(ios):
    nonlocal finish_checked
    finish_checked = True
    assert len(ios) == 1
    for key in ios:
        io = key
        data = ios[key].to_numpy().flatten()
        output_data = self.correct_data.flatten()
        ...
    return 0

network.set_finish_callback(finish_callback)
network.forward()
network.wait()
assert finish_checked == True

enable_profile_performance

def enable_profile_performance(self, profile_file):

模型运行时候,对模型中的各个 Operator 进行速度测试,并将测试结果写到指定的 profile_file 中,得到的这个 profile 文件为 json 文件, 可以使用 MegEngine 中指定的 tool 进行解析。

set_network_algo_workspace_limit

def set_network_algo_workspace_limit(self, size_limit):

模型运行时候,模型中每一个 Operator 运行时候选择的算法最大能够用到的 workspace 大小,超过 size_limit 大小的算法将不会被选择,其中 size_limit 的单位为字节。

set_network_algo_policy

def set_network_algo_policy(
    self, policy, shared_batch_size=0, binary_equal_between_batch=False
):

设置模型运行时候选择每个 Operator 算法的策略。

参数:

  • policy 选择算法的策略,MegEngine Lite 中支持以下策略:

    class LiteAlgoSelectStrategy(IntEnum):
        """
        operation algorithm seletion strategy type, some operations have
        multi algorithms, different algorithm has different attribute, according to
        the strategy, the best algorithm will be selected.
    
        Note: These strategies can be combined
    
        LITE_ALGO_HEURISTIC | LITE_ALGO_PROFILE means: if profile cache not valid,
        use heuristic instead
    
        LITE_ALGO_HEURISTIC | LITE_ALGO_REPRODUCIBLE means: heuristic choice the
        reproducible algo
    
        LITE_ALGO_PROFILE | LITE_ALGO_REPRODUCIBLE means: profile the best
        algorithm from the reproducible algorithms set
    
        LITE_ALGO_PROFILE | LITE_ALGO_OPTIMIZED means: profile the best
        algorithm form the optimzed algorithms, thus profile will process fast
    
        LITE_ALGO_PROFILE | LITE_ALGO_OPTIMIZED | LITE_ALGO_REPRODUCIBLE means:
        profile the best algorithm form the optimzed and reproducible algorithms
        """
    
        LITE_ALGO_HEURISTIC = 1
        LITE_ALGO_PROFILE = 2
        LITE_ALGO_REPRODUCIBLE = 4
        LITE_ALGO_OPTIMIZED = 8
    

    其中上面的策略在不冲突的情况下,可以进行与操作,然后组合在一起。

  • shared_batch_size:binary_equal_between_batch 的时候,选择最优算法所依据的 batch 大小,设置 0 将使用模型默认的 batch size。

  • binary_equal_between_batch: 多个 batch 同时进行计算时,如果输入完全一样,保证所有 batch 的计算结果完全一样。

io_txt_dump

def io_txt_dump(self, txt_file):

将 LiteNetwork 运行时候的所有 IO tensor 输出到文本文件 io_txt_out_file 中。

io_bin_dump

def io_bin_dump(self, bin_dir):

将 LiteNetwork 运行时候的所有 IO tensor 以二进制的形式保存在 bin_dir 文件夹中。

全局设置相关 API

全局接口在 MegEngine Lite 中都封装在 LiteGlobal 中,都作为它的静态函数存在。

register_decryption_and_key

@staticmethod
 def register_decryption_and_key(decryption_name, decryption_func, key):

注册用户自定义的模型解密算法到 MegEngine Lite 中,包括解密方法和解密需要的秘钥。

参数:

  • decryption_name:解密算法的名字,字符串。

  • decryption_func:解密算法的方法,以及闭包函数。

  • key:解密算法的秘钥。

示例:

@decryption_func
def function(in_arr, key_arr, out_arr):
    if not out_arr:
        return in_arr.size
    else:
        for i in range(in_arr.size):
            out_arr[i] = in_arr[i] ^ key_arr[0] ^ key_arr[0]
        return out_arr.size

LiteGlobal.register_decryption_and_key("just_for_test", function, [15])
config = LiteConfig()
config.bare_model_cryption_name = "just_for_test".encode("utf-8")

network = LiteNetwork(config)
model_path = os.path.join(self.source_dir, "shufflenet.mge")
network.load(model_path)

update_decryption_key

@staticmethod
def update_decryption_key(decryption_name, key):

更新 MegEngine Lite 中 build-in 的解密算法的秘钥。

  • decryption_name:解密算法的名字,目前 MegEngine Lite 中写了三种加密算法,分别是:”AES_default”,”RC4_default” 和 “SIMPLE_FAST_RC4_default”。

  • 对应的秘钥:”AES_default” 为 32 字节数组,”RC4_default” 和 “SIMPLE_FAST_RC4_default” 为 16 自己数组。

示例:

wrong_key = [0] * 32
LiteGlobal.update_decryption_key("AES_default", wrong_key)

with self.assertRaises(RuntimeError):
    config = LiteConfig()
    config.bare_model_cryption_name = "AES_default".encode("utf-8")
    network = LiteNetwork(config)
    model_path = os.path.join(self.source_dir, "shufflenet_crypt_aes.mge")
    network.load(model_path)

right_key = [i for i in range(32)]
LiteGlobal.update_decryption_key("AES_default", right_key)

config = LiteConfig()
config.bare_model_cryption_name = "AES_default".encode("utf-8")
network = LiteNetwork(config)
model_path = os.path.join(self.source_dir, "shufflenet_crypt_aes.mge")
network.load(model_path)

set_loader_lib_path

@staticmethod
def set_loader_lib_path(path):

当第三方硬件以 loader 的形式接入到 MegEngine 中,该接口用于用户设置对应 loader 的执行动态库,path 为执行的动态库的路径。

set_persistent_cache

@staticmethod
def set_persistent_cache(path, always_sync=False):

设置当前 MegEngine Lite 中模型运行的算法 cache,模型运行时将从这个 cache 中取出对应 Operator 的算法信息,并解析找到执行算法,并运行, 这将节省模型运行时候搜索最优的算法时候,用户可以提前搜索好对应的 cache。

dump_persistent_cache

@staticmethod
def dump_persistent_cache(path):

将当前 MegEngine Lite 中模型运行的算法的 cache 从内存中 dump 到指定文件中,该方法可以用户用户提前所有最优算法的 cache。

dump_persistent_cache

@staticmethod
def get_device_count(device_type):

获取指定 device_type 类型的设备数量。

try_coalesce_all_free_memory

def try_coalesce_all_free_memory():

释放当前 MegEngine Lite 中所有不在需要的内存,这样将减少当前系统内存使用峰值。

tensorrt_cache

def set_tensorrt_cache(path):
def dump_tensorrt_cache():

设置以及下载 tensorRT 的 cache。

set_log_level

def set_log_level(LiteLogLevel):

class LiteLogLevel(IntEnum):
    """
    DEBUG: The most verbose level, printing debugging info
    INFO: The default level
    WARN: Printing warnings
    ERROR: The least verbose level, printing errors only
    """

    DEBUG = 0
    INFO = 1
    WARN = 2
    ERROR = 3

设置 MegEngine Lite 的 log 级别,改函数不在 LiteGlobal 类中,是一个独立的全局函数。

物理地址和虚拟地址操作

def register_memory_pair(
        vir_ptr, phy_ptr, length, device, backend=LiteBackend.LITE_DEFAULT
    ):
def clear_memory_pair(vir_ptr, phy_ptr, device, backend=LiteBackend.LITE_DEFAULT):
def lookup_physic_ptr(vir_ptr, device, backend=LiteBackend.LITE_DEFAULT):

部分设备上有虚拟地址和物理地址的概念,这里提供用户操作虚拟地址和物理地址的接口,主要有: * 设置全局的物理地址和虚拟地址对 * 清除这些地址对 * 通过虚拟地址查询物理地址

Utils API

MegEngine Lite 现在有一个 utils ,TensorBatchCollector,主要为方便用户在进行推理之前收集多个 batch 数据,然后将攒出来的一个多个 batch 的数据 同时放到 LiteNetwork 中进行推理,避免不必要的内存拷贝。

TensorBatchCollector

def __init__(
    self,
    shape,
    dtype=LiteDataType.LITE_INT8,
    device_type=LiteDeviceType.LITE_CUDA,
    device_id=0,
    is_pinned_host=False,
    tensor=None,
):

创建一个 TensorBatchCollector,这个 TensorBatchCollector 默认数据类型是 INT8,设备为 CUDA。

参数:

  • shape:用户指定 TensorBatchCollector 的 shape。

  • dtype:具体的数据类型,可以是 LITE_FLOAT,LITE_HALF,LITE_INT,LITE_INT16,LITE_INT8,LITE_UINT8,LITE_UINT16。

  • device_type:具体的设备类型。

  • device_id:TensorBatchCollector 所在的设备 id。

  • is_pinned_host:该 TensorBatchCollector 申请的内存是否为: 锁页内存

  • tensor:可选的用户设置已经创建好的 LiteTensor 到 TensorBatchCollector 中。

collect_id

def collect_id(self, array, batch_id):

设置该 TensorBatchCollector 中指定 batch_id 的数据为用户输入的 array。

参数:

  • array:可以是 numpy 的 ndarry,也可以是 LiteTensor 类型。

    • 如果是 numpy 的 ndarry,MegEngine Lite 将调用 LiteTensor 的 set_data_by_copy 将数据 copy 到指定的 batch_id 的内存中。

    • 如果是 LiteTensor 类型,MegEngine Lite 将调用 LiteTensor 的 copy_from 完成数据 copy。

  • batch_id:用户指定将要拷贝 array 数据的目标 batch。

collect_by_ctypes

def collect_by_ctypes(self, data, length):

当用户的数据为 ctypes 的 c_void_p,可以调用该接口将数据设置到第一个空着的 batch 中。

collect

def collect(self, array):

当用户需要顺序的搜集batch,如从 0 一直到最大 batch,可以直接调用该接口。

free

def free(self, indexes):

释放指定的 indexes,indexes 是一个 list。

get

def get(self):

获得该 TensorBatchCollector 中内部存储数据的完整 LiteTensor。

to_numpy

def to_numpy(self):

获得该 TensorBatchCollector 中的数据保存在 numpy 的 array 中,并返回。

  • 示例1:顺序的进行攒 batch

    batch_tensor = TensorBatchCollector(
        [4, 8, 8], dtype=LiteDataType.LITE_INT, device_type=LiteDeviceType.LITE_CUDA
    )
    arr = np.ones([8, 8], "int32")
    for j in range(2):
        for i in range(4):
            batch_tensor.collect(arr)
            arr += 1
        batch_tensor.free(range(4))
    data = batch_tensor.to_numpy()
    assert data.shape[0] == 4
    assert data.shape[1] == 8
    assert data.shape[2] == 8
    for i in range(4):
        for j in range(64):
            assert data[i][j // 8][j % 8] == i + 4 + 1
    
  • 示例2:通过指定 batch_id 进行攒 batch

    batch_tensor = TensorBatchCollector(
        [4, 8, 8], dtype=LiteDataType.LITE_INT, device_type=LiteDeviceType.LITE_CUDA
    )
    arr = np.ones([8, 8], "int32")
    arr += 1  # ==2
    batch_tensor.collect_id(arr, 1)
    arr -= 1  # ==1
    batch_tensor.collect_id(arr, 0)
    arr += 2  # ==3
    batch_tensor.collect_id(arr, 2)
    arr += 1  # ==4
    batch_tensor.collect_id(arr, 3)
    
    data = batch_tensor.to_numpy()
    batch_tensor.free(range(4))
    assert data.shape[0] == 4
    assert data.shape[1] == 8
    assert data.shape[2] == 8
    for i in range(4):
        for j in range(64):
            assert data[i][j // 8][j % 8] == i + 1
    
  • 示例3:通过 ctpes 进行攒 batch

    all_tensor = LiteTensor(
        LiteLayout([4, 6, 8], dtype=LiteDataType.LITE_INT),
        device_type=LiteDeviceType.LITE_CUDA,
    )
    batch_tensor = TensorBatchCollector([4, 6, 8], tensor=all_tensor)
    nparr = np.ones([6, 8], "int32")
    for j in range(2):
        for i in range(4):
            batch_tensor.collect(nparr)
            nparr += 1
        batch_tensor.free(range(4))
    data = batch_tensor.to_numpy()
    assert data.shape[0] == 4
    assert data.shape[1] == 6
    assert data.shape[2] == 8
    for i in range(4):
        for j in range(48):
            assert data[i][j // 8][j % 8] == i + 4 + 1
    
  • 示例4:通过 LiteTensor 进行攒 batch

    batch_tensor = TensorBatchCollector(
        [4, 6, 8], dtype=LiteDataType.LITE_INT, device_type=LiteDeviceType.LITE_CPU
    )
    nparr = np.ones([6, 8], "int32")
    tensor = LiteTensor(LiteLayout([6, 8], LiteDataType.LITE_INT))
    for j in range(2):
        for i in range(4):
            tensor.set_data_by_share(nparr)
            batch_tensor.collect(tensor)
            nparr += 1
        batch_tensor.free(range(4))
    data = batch_tensor.to_numpy()
    assert data.shape[0] == 4
    assert data.shape[1] == 6
    assert data.shape[2] == 8
    for i in range(4):
        for j in range(48):
            assert data[i][j // 8][j % 8] == i + 4 + 1