Skip to content
体验新版
项目
组织
正在加载...
登录
切换导航
打开侧边栏
MegEngine 天元
MegEngine
提交
92ded721
MegEngine
项目概览
MegEngine 天元
/
MegEngine
1 年多 前同步成功
通知
403
Star
4705
Fork
582
代码
文件
提交
分支
Tags
贡献者
分支图
Diff
Issue
0
列表
看板
标记
里程碑
合并请求
0
DevOps
流水线
流水线任务
计划
Wiki
0
Wiki
分析
仓库
DevOps
项目成员
Pages
MegEngine
项目概览
项目概览
详情
发布
仓库
仓库
文件
提交
分支
标签
贡献者
分支图
比较
Issue
0
Issue
0
列表
看板
标记
里程碑
合并请求
0
合并请求
0
Pages
DevOps
DevOps
流水线
流水线任务
计划
分析
分析
仓库分析
DevOps
Wiki
0
Wiki
成员
成员
收起侧边栏
关闭侧边栏
动态
分支图
创建新Issue
流水线任务
提交
Issue看板
提交
92ded721
编写于
6月 21, 2022
作者:
M
Megvii Engine Team
浏览文件
操作
浏览文件
下载
电子邮件补丁
差异文件
docs(api/lite): add global setting docstring for megenginelite
GitOrigin-RevId: 764a479cab61b99dde15ab4f012cabbdd408f325
上级
171c683d
变更
1
隐藏空白更改
内联
并排
Showing
1 changed file
with
97 addition
and
4 deletion
+97
-4
lite/pylite/megenginelite/global_setting.py
lite/pylite/megenginelite/global_setting.py
+97
-4
未找到文件。
lite/pylite/megenginelite/global_setting.py
浏览文件 @
92ded721
...
...
@@ -16,7 +16,7 @@ LiteDecryptionFunc = CFUNCTYPE(
class
_GlobalAPI
(
_LiteCObjBase
):
"""
get the api
from the lib
Get APIs
from the lib
"""
_api_
=
[
...
...
@@ -42,8 +42,12 @@ class _GlobalAPI(_LiteCObjBase):
def
decryption_func
(
func
):
"""the decryption function decorator
:type func: a function accept three array, in_arr, key_arr and out_arr, if out_arr is None, just query the out array lenght in byte
"""the decryption function decorator.
.. note::
The function accept three array: ``in_arr``, ``key_arr`` and ``out_arr``.
If ``out_arr`` is None, just query the out array length in byte.
"""
@
CFUNCTYPE
(
c_size_t
,
c_void_p
,
c_size_t
,
POINTER
(
c_uint8
),
c_size_t
,
c_void_p
)
...
...
@@ -63,13 +67,23 @@ def decryption_func(func):
class
LiteGlobal
(
object
):
"""
s
ome global config in lite
S
ome global config in lite
"""
_api
=
_GlobalAPI
().
_lib
@
staticmethod
def
register_decryption_and_key
(
decryption_name
,
decryption_func
,
key
):
"""Register a custom decryption method and key to lite
Args:
decryption_name: the name of the decryption, which will act as the hash
key to find the decryption method.
decryption_func: the decryption function, which will decrypt the model with
the registered key, then return the decrypted model.
See :py:func:`~.decryption_func` for more details.
key: the decryption key of the method.
"""
c_name
=
c_char_p
(
decryption_name
.
encode
(
"utf-8"
))
key_length
=
len
(
key
)
c_key
=
(
c_uint8
*
key_length
)(
*
key
)
...
...
@@ -79,6 +93,14 @@ class LiteGlobal(object):
@
staticmethod
def
update_decryption_key
(
decryption_name
,
key
):
"""Update decryption key of a custom decryption method.
Args:
decrypt_name: the name of the decryption,
which will act as the hash key to find the decryption method.
key: the decryption key of the method,
if the length of key is zero, the key will not be updated.
"""
c_name
=
c_char_p
(
decryption_name
.
encode
(
"utf-8"
))
key_length
=
len
(
key
)
c_key
=
(
c_uint8
*
key_length
)(
*
key
)
...
...
@@ -86,42 +108,92 @@ class LiteGlobal(object):
@
staticmethod
def
set_loader_lib_path
(
path
):
"""Set the loader path to be used in lite.
Args:
path: the file path which store the loader library.
"""
c_path
=
c_char_p
(
path
.
encode
(
"utf-8"
))
LiteGlobal
.
_api
.
LITE_set_loader_lib_path
(
c_path
)
@
staticmethod
def
set_persistent_cache
(
path
,
always_sync
=
False
):
"""Set the algo policy cache file for CPU/CUDA,
the algo policy cache is produced by MegEngine fast-run.
Args:
path: the file path which store the cache.
always_sync: always update the cache file when model runs.
"""
c_path
=
c_char_p
(
path
.
encode
(
"utf-8"
))
LiteGlobal
.
_api
.
LITE_set_persistent_cache
(
c_path
,
always_sync
)
@
staticmethod
def
set_tensorrt_cache
(
path
):
"""Set the TensorRT engine cache path for serialized prebuilt ICudaEngine.
Args:
path: the cache file path to set
"""
c_path
=
c_char_p
(
path
.
encode
(
"utf-8"
))
LiteGlobal
.
_api
.
LITE_set_tensorrt_cache
(
c_path
)
@
staticmethod
def
dump_persistent_cache
(
path
):
"""Dump the PersistentCache policy cache to the specific file.
If the network is set to profile when forward,
though this the algo policy will dump to file.
Args:
path: the cache file path to be dumped.
"""
c_path
=
c_char_p
(
path
.
encode
(
"utf-8"
))
LiteGlobal
.
_api
.
LITE_dump_persistent_cache
(
c_path
)
@
staticmethod
def
dump_tensorrt_cache
():
"""Dump the TensorRT cache to the file set in :py:func:`~.set_tensorrt_cache`."""
LiteGlobal
.
_api
.
LITE_dump_tensorrt_cache
()
@
staticmethod
def
get_device_count
(
device_type
):
"""Get the number of device of the given device type in current context.
Args:
device_type: the device type to be counted.
Returns:
the number of device.
"""
count
=
c_size_t
()
LiteGlobal
.
_api
.
LITE_get_device_count
(
device_type
,
byref
(
count
))
return
count
.
value
@
staticmethod
def
try_coalesce_all_free_memory
():
"""Try to coalesce all free memory in MegEngine.
When call it MegEnine Lite will try to free all the unused memory
thus decrease the runtime memory usage.
"""
LiteGlobal
.
_api
.
LITE_try_coalesce_all_free_memory
()
@
staticmethod
def
register_memory_pair
(
vir_ptr
,
phy_ptr
,
length
,
device
,
backend
=
LiteBackend
.
LITE_DEFAULT
):
"""Register the physical and virtual address pair to the MegEngine,
some device need the map from physical to virtual.
Args:
vir_ptr: the virtual ptr to set to MegEngine.
phy_ptr: the physical ptr to set to MegEngine.
length: the length of bytes to set pair memory.
device: the the device to set the pair memory.
backend: the backend to set the pair memory
Return:
Whether the register operation is successful.
"""
assert
isinstance
(
vir_ptr
,
c_void_p
)
and
isinstance
(
phy_ptr
,
c_void_p
),
"clear memory pair only accept c_void_p type."
...
...
@@ -131,6 +203,17 @@ class LiteGlobal(object):
@
staticmethod
def
clear_memory_pair
(
vir_ptr
,
phy_ptr
,
device
,
backend
=
LiteBackend
.
LITE_DEFAULT
):
"""Clear the physical and virtual address pair in MegEngine.
Args:
vir_ptr: the virtual ptr to set to MegEngine.
phy_ptr: the physical ptr to set to MegEngine.
device: the the device to set the pair memory.
backend: the backend to set the pair memory.
Return:
Whether the clear is operation successful.
"""
assert
isinstance
(
vir_ptr
,
c_void_p
)
and
isinstance
(
phy_ptr
,
c_void_p
),
"clear memory pair only accept c_void_p type."
...
...
@@ -138,6 +221,16 @@ class LiteGlobal(object):
@
staticmethod
def
lookup_physic_ptr
(
vir_ptr
,
device
,
backend
=
LiteBackend
.
LITE_DEFAULT
):
"""Get the physic address by the virtual address in MegEngine.
Args:
vir_ptr: the virtual ptr to set to MegEngine.
device: the the device to set the pair memory.
backend: the backend to set the pair memory.
Return:
The physic address to lookup.
"""
assert
isinstance
(
vir_ptr
,
c_void_p
),
"lookup physic ptr only accept c_void_p type."
...
...
编辑
预览
Markdown
is supported
0%
请重试
或
添加新附件
.
添加附件
取消
You are about to add
0
people
to the discussion. Proceed with caution.
先完成此消息的编辑!
取消
想要评论请
注册
或
登录