技术背景
这篇文章来源于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
复现代码结构
debug/
├── debug
│ ├── __init__.py
│ └── test_vmap.py
├── example
│ └── test.py
└── setup.py
代码内容
setup.py
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
from mindspore import ops
func = lambda ts: ts
batch_func = ops.vmap(func, in_axes=(0, ), out_axes=(0, ))
test.py
# 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)
复现步骤
- 按照目录树结构构建好本地测试复现环境；
- 取消test.py文件中的两行注释，运行python3 test.py；
- 通过setup.py构建安装python3 setup.py install，然后把test.py文件的头两行注释掉，再次运行python3 test.py。
预期结果
两次运行test.py都输出[0 1]
报错信息
第一次测试，为了方便代码阅读，使用了lambda函数，后来发现这个报错可能是即时编译跟lambda函数不兼容所导致的：
后来修改成了普通的函数，才复现了一个RuntimeError：
test_vmap.py
from mindspore import ops, jit
@jit
def func(ts):
return ts
batch_func = ops.vmap(func, in_axes=(0, ), out_axes=(0, ))
得到的报错信息为：
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包的内容识别为第三方仓库。从第三方仓库来的函数，有可能出现无法入图的问题：
解决的方法就是，把相应的仓库添加到环境变量中去：
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