自动向量化

自动向量化#

在上一节中，我们讨论了通过 jax.jit() 函数进行的 JIT 编译。本笔记本讨论了 JAX 的另一个转换：通过 jax.vmap() 进行向量化。

手动向量化#

考虑以下计算两个一维向量卷积的简单代码

import jax
import jax.numpy as jnp

x = jnp.arange(5)
w = jnp.array([2., 3., 4.])

def convolve(x, w):
  output = []
  for i in range(1, len(x)-1):
    output.append(jnp.dot(x[i-1:i+2], w))
  return jnp.array(output)

convolve(x, w)

Array([11., 20., 29.], dtype=float32)

假设我们想将此函数应用于一批权重 w 到一批向量 x。

xs = jnp.stack([x, x])
ws = jnp.stack([w, w])

最简单的方法是简单地在 Python 中循环遍历批次

def manually_batched_convolve(xs, ws):
  output = []
  for i in range(xs.shape[0]):
    output.append(convolve(xs[i], ws[i]))
  return jnp.stack(output)

manually_batched_convolve(xs, ws)

Array([[11., 20., 29.],
       [11., 20., 29.]], dtype=float32)

这会产生正确的结果，但效率不高。

为了有效地批处理计算，您通常必须手动重写该函数以确保以向量化形式完成。这实现起来并不特别困难，但确实涉及到更改函数处理索引、轴和输入其他部分的方式。

例如，我们可以手动重写 convolve() 以支持跨批处理维度的向量化计算，如下所示

def manually_vectorized_convolve(xs, ws):
  output = []
  for i in range(1, xs.shape[-1] -1):
    output.append(jnp.sum(xs[:, i-1:i+2] * ws, axis=1))
  return jnp.stack(output, axis=1)

manually_vectorized_convolve(xs, ws)

Array([[11., 20., 29.],
       [11., 20., 29.]], dtype=float32)

随着函数复杂性的增加，这种重新实现可能会很混乱且容易出错；幸运的是，JAX 提供了另一种方法。

自动向量化#

在 JAX 中，jax.vmap() 变换旨在自动生成函数的这种向量化实现。

auto_batch_convolve = jax.vmap(convolve)

auto_batch_convolve(xs, ws)

Array([[11., 20., 29.],
       [11., 20., 29.]], dtype=float32)

它通过类似于 jax.jit() 的方式跟踪函数，并在每个输入的开头自动添加批次轴。

如果批次维度不是第一个，你可以使用 in_axes 和 out_axes 参数来指定输入和输出中批次维度的位置。如果所有输入和输出的批次轴都相同，则这些参数可以是整数；否则，可以是列表。

auto_batch_convolve_v2 = jax.vmap(convolve, in_axes=1, out_axes=1)

xst = jnp.transpose(xs)
wst = jnp.transpose(ws)

auto_batch_convolve_v2(xst, wst)

Array([[11., 11.],
       [20., 20.],
       [29., 29.]], dtype=float32)

jax.vmap() 还支持仅对一个参数进行批处理的情况：例如，如果你想用一批向量 x 对单个权重集 w 进行卷积；在这种情况下，可以将 in_axes 参数设置为 None。

batch_convolve_v3 = jax.vmap(convolve, in_axes=[0, None])

batch_convolve_v3(xs, w)

Array([[11., 20., 29.],
       [11., 20., 29.]], dtype=float32)

组合变换#

与所有 JAX 变换一样，jax.jit() 和 jax.vmap() 被设计为可组合的，这意味着你可以用 jit 包裹一个 vmapped 函数，或者用 vmap 包裹一个 jitted 函数，并且一切都会正常工作。

jitted_batch_convolve = jax.jit(auto_batch_convolve)

jitted_batch_convolve(xs, ws)

Array([[11., 20., 29.],
       [11., 20., 29.]], dtype=float32)

自动向量化

目录

自动向量化#

手动向量化#

自动向量化#

组合变换#