python - JAX:jit 函数的时间随着函数访问的内存而超线性增长
问题描述
这是一个简单的例子,它对两个高斯 pdf 的乘积进行数值积分。其中一个高斯是固定的,均值始终为 0。另一个高斯的均值不同:
import time
import jax.numpy as np
from jax import jit
from jax.scipy.stats.norm import pdf
# set up evaluation points for numerical integration
integr_resolution = 6400
lower_bound = -100
upper_bound = 100
integr_grid = np.linspace(lower_bound, upper_bound, integr_resolution)
proba = pdf(integr_grid)
integration_weight = (upper_bound - lower_bound) / integr_resolution
# integrate with new mean
def integrate(mu_new):
x_new = integr_grid - mu_new
proba_new = pdf(x_new)
total_proba = sum(proba * proba_new * integration_weight)
return total_proba
print('starting jit')
start = time.perf_counter()
integrate = jit(integrate)
integrate(1)
stop = time.perf_counter()
print('took: ', stop - start)
该函数看起来很简单,但它根本无法扩展。以下列表包含成对的(integr_resolution 的值,运行代码所用的时间):
- 100 | 0.107s
- 200 | 0.23s
- 400 | 0.537s
- 800 | 1.52s
- 1600 | 5.2s
- 3200 | 19 岁
- 6400 | 134s
作为参考,应用到的 unjitted 函数integr_resolution=6400
需要 0.02s。
我认为这可能与函数正在访问全局变量这一事实有关。但是移动代码以在函数内部设置积分点对时序没有显着影响。以下代码需要 5.36 秒才能运行。它对应于先前花费 5.2 秒的 1600 表条目:
# integrate with new mean
def integrate(mu_new):
# set up evaluation points for numerical integration
integr_resolution = 1600
lower_bound = -100
upper_bound = 100
integr_grid = np.linspace(lower_bound, upper_bound, integr_resolution)
proba = pdf(integr_grid)
integration_weight = (upper_bound - lower_bound) / integr_resolution
x_new = integr_grid - mu_new
proba_new = pdf(x_new)
total_proba = sum(proba * proba_new * integration_weight)
return total_proba
这里发生了什么?
解决方案
我也在https://github.com/google/jax/issues/1776回答了这个问题,但也在这里添加了答案。
这是因为代码使用sum
了它应该使用的地方np.sum
。
sum
是一个 Python 内置程序,它提取序列的每个元素并使用+
运算符将它们一一相加。这具有构建大型、展开的添加链的效果,XLA 需要很长时间才能编译。
如果使用np.sum
,那么 JAX 会构建一个 XLA 归约运算符,它的编译速度要快得多。
并且只是为了说明我是如何解决这个问题的:我使用了jax.make_jaxpr
,它转储了 JAX 的函数的内部跟踪表示。在这里,它显示:
In [3]: import jax
In [4]: jax.make_jaxpr(integrate)(1)
Out[4]:
{ lambda b c ; ; a.
let d = convert_element_type[ new_dtype=float32
old_dtype=int32 ] a
e = sub c d
f = sub e 0.0
g = pow f 2.0
h = div g 1.0
i = add 1.8378770351409912 h
j = neg i
k = div j 2.0
l = exp k
m = mul b l
n = mul m 2.0
o = slice[ start_indices=(0,)
limit_indices=(1,)
strides=(1,)
operand_shape=(100,) ] n
p = reshape[ new_sizes=()
dimensions=None
old_sizes=(1,) ] o
q = add p 0.0
r = slice[ start_indices=(1,)
limit_indices=(2,)
strides=(1,)
operand_shape=(100,) ] n
s = reshape[ new_sizes=()
dimensions=None
old_sizes=(1,) ] r
t = add q s
u = slice[ start_indices=(2,)
limit_indices=(3,)
strides=(1,)
operand_shape=(100,) ] n
v = reshape[ new_sizes=()
dimensions=None
old_sizes=(1,) ] u
w = add t v
x = slice[ start_indices=(3,)
limit_indices=(4,)
strides=(1,)
operand_shape=(100,) ] n
y = reshape[ new_sizes=()
dimensions=None
old_sizes=(1,) ] x
z = add w y
... similarly ...
然后很明显为什么这很慢:程序非常大。
对比np.sum
版本:
In [5]: def integrate(mu_new):
...: x_new = integr_grid - mu_new
...:
...: proba_new = pdf(x_new)
...: total_proba = np.sum(proba * proba_new * integration_weight)
...:
...: return total_proba
...:
In [6]: jax.make_jaxpr(integrate)(1)
Out[6]:
{ lambda b c ; ; a.
let d = convert_element_type[ new_dtype=float32
old_dtype=int32 ] a
e = sub c d
f = sub e 0.0
g = pow f 2.0
h = div g 1.0
i = add 1.8378770351409912 h
j = neg i
k = div j 2.0
l = exp k
m = mul b l
n = mul m 2.0
o = reduce_sum[ axes=(0,)
input_shape=(100,) ] n
in [o] }
希望有帮助!
推荐阅读
- python - 如何使用 PANDAS 数据框中的值作为正则表达式中的过滤器参数?
- javascript - 写入文件而不用 fs 覆盖
- python - Python:如何使用从列计算的乘法字段扩展 DataFrame
- groovy - 如何使用脚本运行器创建自定义字段?
- haskell - 递归类型是haskell,它不重复同一个构造函数
- javascript - 在复选框更改时,找到并选中另一个复选框并修改一些文本输入的值
- c++ - C++ 无锁队列
- django - MultiValueDictKeyError,当我按下注册按钮时
- java - 如何在 xpath 中传递变量字符串以识别 Selenium 和 Java 中的元素
- python - 如何使用 Matplotlib 更改 Y 轴的比例距离?