在机器学习中，模型集成（ensemble learning）是一种常用的提升模型性能和鲁棒性的技术。当需要对一个包含多个神经网络的模型集成进行推理时，顺序执行每个模型的推理过程（如通过for循环）效率低下，尤其是在gpu或tpu等并行计算设备上。jax的jax.vmap函数是解决这类问题的强大工具，它能够自动向量化函数，将批量操作转换为单个函数调用，从而实现高效并行计算。
考虑一个简单的神经网络集成，其损失计算通常如下：
for params in ensemble_params:
loss = mse_loss(params, inputs=x, targets=y)
def mse_loss(params, inputs, targets):
preds = batched_predict(params, inputs)
loss = jnp.mean((targets - preds) ** 2)
return loss这里，ensemble_params是一个包含多个网络参数（PyTree结构）的Python列表。为了避免for循环，我们自然会想到使用jax.vmap：
ensemble_loss = jax.vmap(fun=mse_loss, in_axes=(0, None, None)) losses = ensemble_loss(ensemble_params, x, y)
然而，这种尝试通常会导致如下ValueError：
ValueError: vmap got inconsistent sizes for array axes to be mapped: * most axes (8 of them) had size 3, e.g. axis 0 of argument params[0][0][0] of type float32[3,2]; * some axes (8 of them) had size 4, e.g. axis 0 of argument params[0][1][0] of type float32[4,3]
这个错误信息表明vmap在尝试对输入参数进行映射时，发现某些数组的特定轴尺寸不一致。这揭示了对vmap如何处理复杂数据结构（如PyTree列表）的常见误解。
jax.vmap 的核心功能是在函数的输入参数上引入一个新的“批处理”维度，并沿着这个维度并行执行函数。它期望接收一个单一的PyTree结构作为输入（对于被映射的参数），其中这个批处理维度作为所有叶子数组（即实际的JAX数组）的第一个轴。
当我们将一个list（例如ensemble_params，它是一个包含多个网络参数PyTree的列表）传递给vmap并指定in_axes=0时，vmap不会自动将这个list的每个元素视为一个独立的批次项。相反，JAX将Python list视为一种特殊的PyTree节点。当vmap尝试对这个PyTree（即我们的ensemble_params列表）应用in_axes=0时，它会深入到这个PyTree的叶子节点。
具体来说，ensemble_params的结构是 [network1_params_pytree, network2_params_pytree, ...]。每个network_params_pytree内部又包含多个层的权重和偏置，例如 [[w1, b1], [w2, b2], ...]。当vmap试图对ensemble_params的“0轴”进行映射时，它会遍历其所有叶子节点（即每个w和b数组），并期望这些叶子数组的第0轴具有一致的尺寸，以便沿着这个轴进行批处理。
然而，在神经网络的参数PyTree中，不同层的权重矩阵（例如w1和w2）的形状通常是不同的。例如，w1可能是 (dim_out_layer1, dim_in_layer1)，而w2可能是 (dim_out_layer2, dim_in_layer2)。这意味着它们的第0轴（dim_out_layer1和dim_out_layer2）尺寸很可能不一致。vmap尝试将这些不一致的第0轴视为要映射的批处理维度，从而导致ValueError。
问题的根本在于：vmap期望的是一个“结构化数组”（Struct-of-Arrays）模式，而不是“结构列表”（List-of-Structs）模式。
要解决这个问题，我们需要将ensemble_params从“结构列表”转换为“结构化数组”。这意味着，对于集成中的所有模型，它们的相同层、相同类型的参数（例如所有模型的第一个隐藏层的权重）应该被堆叠到一个新的JAX数组中，形成一个新的批处理维度。
jax.tree_map函数是实现这一转换的理想工具。它可以在两个或多个PyTree上并行地应用一个函数，对它们对应的叶子节点进行操作。
ensemble_params_batched = jax.tree_map(lambda *args: jnp.stack(args), *ensemble_params)
这段代码的工作原理如下：
经过此转换后，ensemble_params_batched将是一个单一的PyTree，其内部的每个叶子数组都包含所有模型的对应参数，且批处理维度作为其第一个轴。现在，这个PyTree的结构就符合vmap的期望了。
让我们将解决方案整合到原始的最小可复现示例中。
import jax
from jax import Array
from jax import random
import jax.numpy as jnp
# 辅助函数：初始化单层参数
def layer_params(dim_in: int, dim_out: int, key: Array) -> tuple[Array]:
w_key, b_key = random.split(key=key)
weights = random.normal(key=w_key, shape=(dim_out, dim_in))
biases = random.normal(key=b_key, shape=(dim_out,)) # Changed w_key to b_key for bias init
return weights, biases
# 辅助函数：初始化单个网络参数PyTree
def init_params(layer_dims: list[int], key: Array) -> list[tuple[Array]]:
keys = random.split(key=key, num=len(layer_dims) - 1) # Split keys for each layer
params = []
for i, (dim_in, dim_out) in enumerate(zip(layer_dims[:-1], layer_dims[1:])):
params.append(layer_params(dim_in=dim_in, dim_out=dim_out, key=keys[i]))
return params
# 辅助函数：初始化网络集成（列表形式）
def init_ensemble(key: Array, num_models: int, layer_dims: list[int]) -> list:
keys = random.split(key=key, num=num_models)
# ensemble_params 是一个 PyTree 的列表 (list-of-structs)
models = [init_params(layer_dims=layer_dims, key=key) for key in keys]
return models
# 激活函数
def relu(x):
return jnp.maximum(0, x)
# 单个网络的前向传播
def predict(params, image):
activations = image
for w, b in params[:-1]:
outputs = jnp.dot(w, activations) + b
activations = relu(outputs)
final_w, final_b = params[-1]
logits = jnp.dot(final_w, activations) + final_b
return logits
# 对输入数据进行批处理预测
batched_predict = jax.vmap(predict, in_axes=(None, 0))
# 均方误差损失函数
def mse_loss(params, inputs, targets):
preds = batched_predict(params, inputs)
loss = jnp.mean((targets - preds) ** 2)
return loss
if __name__ == "__main__":
num_models = 4
dim_in = 2
dim_out = 4
layer_dims = [dim_in, 3, dim_out] # 示例网络结构
batch_size = 2
key = random.PRNGKey(seed=1)
key, subkey = random.split(key)
# 初始化模型集成，得到 list-of-structs 形式的参数
ensemble_params_list = init_ensemble(key=subkey, num_models=num_models, layer_dims=layer_dims)
key_x, key_y = random.split(key)
x = random.normal(key=key_x, shape=(batch_size, dim_in))
y = random.normal(key=key_y, shape=(batch_size, dim_out))
print("--- 传统 for 循环计算损失 ---")
for params in ensemble_params_list:
loss = mse_loss(params, inputs=x, targets=y)
print(f"{loss = }")
print("\n--- 尝试直接 vmap (会报错) ---")
try:
# 这一行会引发 ValueError，因为 ensemble_params_list 是 list-of-structs
ensemble_loss_vmap_fail = jax.vmap(fun=mse_loss, in_axes=(0, None, None))
losses_fail = ensemble_loss_vmap_fail(ensemble_params_list, x, y)
print(f"Failed vmap losses = {losses_fail}")
except ValueError as e:
print(f"捕获到预期错误: {e}")
print("\n--- 转换为 Struct-of-Arrays 并使用 vmap ---")
# 关键步骤：将 list-of-structs 转换为 struct-of-arrays
ensemble_params_batched = jax.tree_map(lambda *args: jnp.stack(args), *ensemble_params_list)
# 现在可以正确使用 vmap
ensemble_loss_vmap_success = jax.vmap(fun=mse_loss, in_axes=(0, None, None))
losses_success = ensemble_loss_vmap_success(ensemble_params_batched, x, y)
print(f"{losses_success = }")
# 预期与 for 循环的结果相同运行结果示例：
--- 传统 for 循环计算损失 --- loss = Array(3.762451, dtype=float32) loss = Array(4.39846, dtype=float32) loss = Array(4.1425314, dtype=float32) loss = Array(6.045669, dtype=float32) --- 尝试直接 vmap (会报错) --- 捕获到预期错误: vmap got inconsistent sizes for array axes to be mapped: * most axes (8 of them) had size 3, e.g. axis 0 of argument params[0][0][0] of type float32[3,2]; * some axes (8 of them) had size 4, e.g. axis 0 of argument params[0][1][0] of type float32[4,3] --- 转换为 Struct-of-Arrays 并使用 vmap --- losses_success = Array([3.762451 , 4.39846 , 4.1425314, 6.045669 ], dtype=float32)
可以看到，经过jax.tree_map转换后，vmap能够成功并行计算所有模型的损失，并且结果与for循环得到的结果一致。
以上就是JAX vmap 高效并行化模型集成推理：从列表到结构化数组的转换的详细内容，更多请关注php中文网其它相关文章！
每个人都需要一台速度更快、更稳定的 PC。随着时间的推移，垃圾文件、旧注册表数据和不必要的后台进程会占用资源并降低性能。幸运的是，许多工具可以让 Windows 保持平稳运行。
广告Copyright 2014-2025 https://www.php.cn/ All Rights Reserved | php.cn | 湘ICP备2023035733号
633
收藏
