前往小程序,Get更优阅读体验!
立即前往
首页
学习
活动
专区
工具
TVP
发布
社区首页 >专栏 >MindSpore编译构建后Vmap模块的RuntimeError问题

MindSpore编译构建后Vmap模块的RuntimeError问题

作者头像
DechinPhy
发布2024-04-12 08:36:54
620
发布2024-04-12 08:36:54
举报
文章被收录于专栏:Dechin的专栏Dechin的专栏

技术背景

这篇文章来源于MindSpore仓库中的一个Issue,简单描述问题就是,如果你用MindSpore开发了一个python软件供别人使用,那么很有可能涉及到编译构建的问题。但是如果直接使用编译好的whl包去运行的话,就有可能出现一个跟Jit即时编译有关的报错,这里Jit在其他的一些模块中也会被使用到,比如Vmap函数和Grad函数等。

问题复现

我们通过一个简单的测试案例来复现一下这个RuntimeError。

软件环境

-- MindSpore version: 2.2.11 -- Python version: 3.9.16 -- OS platform and distribution: Linux Ubuntu 20.04

执行模式

GRAPH_MODE

复现代码结构

代码语言:javascript
复制
debug/
├── debug
│   ├── __init__.py
│   └── test_vmap.py
├── example
│   └── test.py
└── setup.py

代码内容

setup.py

代码语言:javascript
复制
from setuptools import setup, find_packages
from pathlib import Path
this_directory = Path(__file__).parent

setup(
    name="debugs",
    version="1.0",
    description="Test Vmap",
    license="Apache 2.0 Licence",
    packages=find_packages(),
    platforms="any",
    scripts=[],
    include_package_data=True
)

test_vmap.py

代码语言:javascript
复制
from mindspore import ops

func = lambda ts: ts
batch_func = ops.vmap(func, in_axes=(0, ), out_axes=(0, ))

test.py

代码语言:javascript
复制
# import sys
# sys.path.insert(0, '../')

from debug.test_vmap import batch_func
from mindspore import Tensor

a = Tensor([0, 1])
batch_a = batch_func(a)
print (batch_a)

复现步骤

  1. 按照目录树结构构建好本地测试复现环境;
  2. 取消test.py文件中的两行注释,运行python3 test.py;
  3. 通过setup.py构建安装python3 setup.py install,然后把test.py文件的头两行注释掉,再次运行python3 test.py。

预期结果

两次运行test.py都输出[0 1]

报错信息

第一次测试,为了方便代码阅读,使用了lambda函数,后来发现这个报错可能是即时编译跟lambda函数不兼容所导致的:

后来修改成了普通的函数,才复现了一个RuntimeError:

test_vmap.py

代码语言:javascript
复制
from mindspore import ops, jit

@jit
def func(ts):
    return ts

batch_func = ops.vmap(func, in_axes=(0, ), out_axes=(0, ))

得到的报错信息为:

代码语言:javascript
复制
Traceback (most recent call last):
  File "test.py", line 8, in <module>
    batch_a = batch_func(a)
  File "/home/dechin/anaconda3/envs/mindspore-latest/lib/python3.7/site-packages/mindspore/common/api.py", line 718, in staging_specialize
    out = _MindsporeFunctionExecutor(func, hash_obj, input_signature, process_obj, jit_config)(*args, **kwargs)
  File "/home/dechin/anaconda3/envs/mindspore-latest/lib/python3.7/site-packages/mindspore/common/api.py", line 121, in wrapper
    results = fn(*arg, **kwargs)
  File "/home/dechin/anaconda3/envs/mindspore-latest/lib/python3.7/site-packages/mindspore/common/api.py", line 350, in __call__
    raise err
  File "/home/dechin/anaconda3/envs/mindspore-latest/lib/python3.7/site-packages/mindspore/common/api.py", line 344, in __call__
    phase = self.compile(self.fn.__name__, *args_list, **kwargs)
  File "/home/dechin/anaconda3/envs/mindspore-latest/lib/python3.7/site-packages/mindspore/common/api.py", line 435, in compile
    is_compile = self._graph_executor.compile(self.fn, compile_args, kwargs, phase, True)
RuntimeError: 'VmapOperation' arg0 must be a 'Function' or 'Cell', but got AbstractScalar(Type: External, Value: PythonObject(type: <class 'function'>, value: <function func at 0x7f75c0375ef0>), Shape: NoShape).

----------------------------------------------------
- C++ Call Stack: (For framework developers)
----------------------------------------------------
mindspore/ccsrc/frontend/operator/composite/composite.cc:1420 CheckVmapFunc

----------------------------------------------------
- The Traceback of Net Construct Code:
----------------------------------------------------
# 0 In file /home/dechin/anaconda3/envs/mindspore-latest/lib/python3.7/site-packages/mindspore/ops/composite/base.py:686
            return vmap_(fn, in_axes, out_axes)(*args, **kwargs)
                   ^
 (See file '/tmp/debug/example/rank_0/om/analyze_fail.ir' for more details. Get instructions about `analyze_fail.ir` at https://www.mindspore.cn/search?inputValue=analyze_fail.ir)

问题原因

按照开发人员的回复,这个问题是因为Vmap会使用到静态图模式。但是编译成whl包之后,静态图模式下会把whl包的内容识别为第三方仓库。从第三方仓库来的函数,有可能出现无法入图的问题:

解决的方法就是,把相应的仓库添加到环境变量中去:

代码语言:javascript
复制
export MS_JIT_MODULES=debug

然后再次运行test.py,问题就解决了。

总结概要

在MindSpore编译计算图的过程中,会把从编译构建好的whl包中引入的模块视为第三方库,也就没有办法在即时编译的阶段入图。普通的math和numpy等第三方库不入图也不会影响计算。但如果是基于MindSpore本身开发的一些函数,如果用到了Jit、Grad和Vmap,那么有可能出现无法入图的问题,就会出现RuntimeError报错。解决方法就是设置一个跟即时编译有关的环境变量,把相关的第三方包引用修改为内部引用。

版权声明

本文首发链接为:https://www.cnblogs.com/dechinphy/p/ms_jit_module.html

作者ID:DechinPhy

更多原著文章:https://www.cnblogs.com/dechinphy/

请博主喝咖啡:https://www.cnblogs.com/dechinphy/gallery/image/379634.html

本文参与 腾讯云自媒体同步曝光计划,分享自作者个人站点/博客。
原始发表:2024-04-11,如有侵权请联系 cloudcommunity@tencent.com 删除

本文分享自 作者个人站点/博客 前往查看

如有侵权,请联系 cloudcommunity@tencent.com 删除。

本文参与 腾讯云自媒体同步曝光计划  ,欢迎热爱写作的你一起参与!

评论
登录后参与评论
0 条评论
热度
最新
推荐阅读
目录
  • 技术背景
  • 问题复现
    • 软件环境
      • 执行模式
        • 复现代码结构
          • 代码内容
            • setup.py
            • test_vmap.py
            • test.py
          • 复现步骤
            • 预期结果
              • 报错信息
                • test_vmap.py
            • 问题原因
            • 总结概要
            • 版权声明
            相关产品与服务
            腾讯云服务器利旧
            云服务器(Cloud Virtual Machine,CVM)提供安全可靠的弹性计算服务。 您可以实时扩展或缩减计算资源,适应变化的业务需求,并只需按实际使用的资源计费。使用 CVM 可以极大降低您的软硬件采购成本,简化 IT 运维工作。
            领券
            问题归档专栏文章快讯文章归档关键词归档开发者手册归档开发者手册 Section 归档