python - 将 numba 函数拆分为项目中的单独模块以进行打包

Question

我的项目中有几个模块，每个模块都包含一些 numba 函数。我知道在第一次导入时，函数会被编译。我现在才注意到的是，即使我只从模块中导入一个函数，似乎所有函数都会被编译，因为导入花费的时间相同。

我想为此实现一种更细粒度的方法，因为对于某些应用程序，您实际上只需要一个函数，因此编译所有函数是浪费时间。

为此，我将函数拆分为单独的模块，如下所示:

Project/
|--src/
|  |-- __init__.py
|  |-- fun1.py
|  |-- fun2.py
|  |-- fun3.py 
|  |-- fun4.py
|  |-- ...

__init__.py 包括

from .fun1 import fun1
from .fun2 import fun2
...

因此它们可以像 from src import fun1 一样导入。

这似乎工作正常，但在导入级别有一些重复，例如每个函数都需要 from numba import jit，其中一些需要 from numpy import zeros 等等。

所以我的问题是这是否是一个好的方法，或者是否有更好的方法来打包许多 numba 函数。

编辑:

将所有导入语句放入 __init__.py 显然意味着一旦导入一个函数，所有函数都会被编译 - 所以根本没有任何好处。

我仍然可以导入类似的功能

from src.fun1 import fun1

这似乎有效。但是语法有点笨拙。

Answer 1

I know that on first import the functions get compiled.

这不一定是真的。 Eager compilation仅当使用显式签名声明函数时才会发生。

声明带有显式签名的函数有两个后果:

• 函数由装饰器编译(在声明或导入时)。
• 不允许其他功能特化。

如果您不需要禁止额外的特化，您可以简单地删除函数签名，它们将在每个特化的第一次执行时编译。

您可以通过在 numba.core.dispatcher.Dispatcher.compile() 方法中放置一个断点来检查函数何时被编译。