主页 / coding / 问题 / 79256001
Accepted
Ben
Asked: 2024-12-06 03:29:25 +0800 CST

在 JAX 中高效自定义数组创建例程

  • 772

我仍在掌握最佳实践jax。我的主要问题如下：

实现自定义数组创建例程的最佳实践是什么jax？

例如，我想实现一个函数，创建一个矩阵，其中除给定列中的 1 外，其他列均为 0。我选择了这个（Jupyter 笔记本）：

import numpy as np
import jax.numpy as jnp

def ones_at_col(shape_mat, idx):
    idxs = jnp.arange(shape_mat[1])[None,:]
    mat = jnp.where(idx==idxs, 1, 0)
    mat = jnp.repeat(mat, shape_mat[0], axis=0)
    return mat

shape_mat = (5,10)

print(ones_at_col(shape_mat, 5))

%timeit np.zeros(shape_mat)

%timeit jnp.zeros(shape_mat)

%timeit ones_at_col(shape_mat, 5)

输出为

[[0 0 0 0 0 1 0 0 0 0]
 [0 0 0 0 0 1 0 0 0 0]
 [0 0 0 0 0 1 0 0 0 0]
 [0 0 0 0 0 1 0 0 0 0]
 [0 0 0 0 0 1 0 0 0 0]]
127 ns ± 0.717 ns per loop (mean ± std. dev. of 7 runs, 10,000,000 loops each)
31.3 µs ± 331 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
123 µs ± 1.79 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

我的功能比常规功能慢了 4 倍jnp.zeros()，这还不算太糟糕。这说明我做的事情并不疯狂。

但是这两个jax例程都比等效例程慢得多numpy。这些函数无法进行 jitted，因为它们将形状作为参数，因此无法跟踪。我猜这就是它们天生就慢的原因？我猜如果它们中的任何一个出现在另一个 jitted 函数的范围内，它们可以被跟踪并加速？

我能做得更好吗？或者我是否正在突破可能的极限jax？

numpy

1 个回答

  1. Best Answer

    最好的方法可能是这样的：

    mat = jnp.zeros(shape_mat).at[:, 5].set(1)

    关于与 NumPy 的时间比较，相关阅读材料是JAX FAQ：JAX 比 NumPy 快吗？摘要是，对于这种特殊情况（创建一个简单的数组），由于 JAX 的每个操作调度开销，您不会期望 JAX 在性能方面与 NumPy 匹敌。

    如果您希望在 JAX 中获得更快的性能，则应始终使用jax.jit即时编译函数。例如，此版本的函数应该是相当理想的（尽管同样，由于 FAQ 链接中讨论的原因，它不如 NumPy 快）：

    @partial(jax.jit, static_argnames=['shape_mat', 'idx'])
def ones_at_col(shape_mat, idx):
  return jnp.zeros(shape_mat).at[:, idx].set(1)

    idx如果您将使用不同的索引值多次调用此函数，则可以保留非静态，如果您在另一个函数中创建这些数组，则应该将代码内联并对该外部函数进行 JIT 编译。

    另一个补充说明：您的微基准测试可能没有测量您认为它们正在测量的内容：有关此方面的提示，请参阅JAX FAQ：对 JAX 代码进行基准测试。特别要注意编译时间和异步调度效果。

    • 1

相关问题