python - JAX:jit 函数的时间随着函数访问的内存而超线性增长

Question

这是一个简单的示例，它对两个高斯 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.107秒
• 200 | 0.23秒
• >400 | 0.537秒
• 800 | 1.52秒
• >1600 | 5.2秒
• 3200 | 19秒
• 6400 | 134秒

仅供引用，应用于 integr_resolution=6400 的未抖动函数需要 0.02 秒。

我认为这可能与该函数正在访问全局变量有关。但是移动代码以在函数内部设置积分点对时序没有显着影响。以下代码运行需要 5.36 秒。它对应于之前需要 5.2s 的 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

这里发生了什么？

Answer 1

我也在https://github.com/google/jax/issues/1776回答了这个问题，但也在这里添加答案。

这是因为代码在应该使用 np.sum 的地方使用了 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] }

希望有帮助!