Skip to content
体验新版
项目
组织
正在加载...
登录
切换导航
打开侧边栏
MegEngine 天元
MegEngine
提交
76119020
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看板
提交
76119020
编写于
7月 10, 2023
作者:
M
Megvii Engine Team
浏览文件
操作
浏览文件
下载
电子邮件补丁
差异文件
docs(imperative): add dlpack docstring
GitOrigin-RevId: 4898d4a4db9053d70633ae8ec8cc4c10ca9dc0c9
上级
15637898
变更
2
显示空白变更内容
内联
并排
Showing
2 changed file
with
60 addition
and
3 deletion
+60
-3
imperative/python/megengine/utils/__init__.py
imperative/python/megengine/utils/__init__.py
+1
-0
imperative/python/megengine/utils/dlpack.py
imperative/python/megengine/utils/dlpack.py
+59
-3
未找到文件。
imperative/python/megengine/utils/__init__.py
浏览文件 @
76119020
# -*- coding: utf-8 -*-
# -*- coding: utf-8 -*-
from
.
import
dlpack
imperative/python/megengine/utils/dlpack.py
浏览文件 @
76119020
from
typing
import
Any
from
..core._imperative_rt.core2
import
_from_dlpack
from
..core._imperative_rt.core2
import
_from_dlpack
__all__
=
[
"to_dlpack"
,
"from_dlpack"
,
]
def
to_dlpack
(
tensor
,
stream
=
None
):
def
to_dlpack
(
tensor
,
stream
=
None
):
"""
Encodes a megengine tensor to DLPack.
Args:
tensor (Tensor): The input tensor, and the data type can be `float16`, `float32`,
`int8`, `int16`, `int32`, `uint8`, `uint16`, `complex64`.
stream (Integer or None): An optional Python integer representing a CUDA stream,
The current stream is synchronized with this stream before the capsule is created.
If None or -1 is passed then no synchronization is performed.
Returns:
dltensor, and the data type is PyCapsule.
Examples:
.. code-block:: python
import megengine as mge
# x is a tensor with shape [2, 3]
x = mge.tensor([[0.2, 0.3, 0.5],
[0.1, 0.2, 0.6]])
dlpack = mge.utils.dlpack.to_dlpack(x)
print(dlpack)
# <capsule object "dltensor" at 0x7ff04b69cc00>
"""
if
stream
is
not
None
and
stream
!=
-
1
:
if
stream
is
not
None
and
stream
!=
-
1
:
return
tensor
.
__dlpack__
(
stream
)
return
tensor
.
__dlpack__
(
stream
)
else
:
else
:
return
tensor
.
__dlpack__
()
return
tensor
.
__dlpack__
()
def
from_dlpack
(
ext_tensor
:
Any
,
stream
=
None
):
def
from_dlpack
(
ext_tensor
,
stream
=
None
):
"""
Decodes a DLPack to a megengine tensor.
Args:
ext_tensor (PyCapsule): a PyCapsule object with the dltensor.
stream (Integer or None): An optional Python integer representing a CUDA stream,
user needs to know which stream the dlpack is generated.
If None then represent producers and consumers on the same stream.
Returns:
tensor (Tensor): a megengine tensor decoded from DLPack.
Examples:
.. code-block:: python
import megengine as mge
# x is a tensor with shape [2, 3]
x = mge.tensor([[0.2, 0.3, 0.5],
[0.1, 0.2, 0.6]])
dlpack = mge.utils.dlpack.to_dlpack(x)
x = mge.utils.dlpack.from_dlpack(dlpack)
print(x)
# Tensor([[0.2, 0.3, 0.5],
# [0.1, 0.2, 0.6]], device=gpu0:0)
"""
if
isinstance
(
stream
,
int
):
if
isinstance
(
stream
,
int
):
assert
stream
>=
0
,
"device stream should be a positive value"
assert
stream
>=
0
,
"device stream should be a positive value"
stream
=
0
if
stream
is
None
else
stream
stream
=
0
if
stream
is
None
else
stream
...
...
编辑
预览
Markdown
is supported
0%
请重试
或
添加新附件
.
添加附件
取消
You are about to add
0
people
to the discussion. Proceed with caution.
先完成此消息的编辑!
取消
想要评论请
注册
或
登录