JAX vmap并行化模型集成推理：从列表结构到数组结构的转换技巧

# 初始的for循环计算损失
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_loss = jax.vmap(fun=mse_loss, in_axes=(0, None, None))
losses = ensemble_loss(ensemble_params, x, y)

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]

import jax
import jax.numpy as jnp

# ... (init_ensemble, predict, mse_loss等函数定义保持不变) ...

# 假设 ensemble_params 是一个 list-of-structs
# 例如：ensemble_params = [model1_params_pytree, model2_params_pytree, ...]

# 核心转换：将 list-of-structs 转换为 struct-of-arrays
ensemble_params_batched = jax.tree_map(lambda *args: jnp.stack(args), *ensemble_params)

import jax
from jax import Array, random
import jax.numpy as jnp

# 辅助函数：初始化单层参数
def layer_params(dim_in: int, dim_out: int, key: Array) -> tuple[Array, 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 b_key to b_key
    return weights, biases

# 辅助函数：初始化单个模型参数
def init_params(layer_dims: list[int], key: Array) -> list[tuple[Array, 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

# 辅助函数：初始化集成模型参数 (list-of-structs)
def init_ensemble(key: Array, num_models: int, layer_dims: list[int]) -> list:
    keys = random.split(key=key, num=num_models)
    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)

    # 1. 初始化集成模型参数 (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("--- 使用Python for循环计算损失 ---")
    for i, params in enumerate(ensemble_params_list):
        loss = mse_loss(params, inputs=x, targets=y)
        print(f"Model {i} loss = {loss}")

    # 2. 将 list-of-structs 转换为 struct-of-arrays
    # 这一步是解决问题的关键！
    ensemble_params_batched = jax.tree_map(lambda *args: jnp.stack(args), *ensemble_params_list)

    # 3. 使用 vmap 对转换后的参数结构进行损失计算
    # in_axes=(0, None, None) 表示：
    # - params: 沿着第一个轴进行映射 (现在这个轴代表了不同的模型)
    # - inputs: 不进行映射 (所有模型使用相同的输入数据)
    # - targets: 不进行映射 (所有模型使用相同的目标数据)
    ensemble_loss_vmapped = jax.vmap(fun=mse_loss, in_axes=(0, None, None))
    losses_vmapped = ensemble_loss_vmapped(ensemble_params_batched, x, y)

    print("\n--- 使用JAX vmap并行计算损失 ---")
    print(f"{losses_vmapped = }")

    # 验证两种方法的结果是否一致
    # 期望输出：
    # Model 0 loss = 3.762451
    # Model 1 loss = 4.39846
    # Model 2 loss = 4.1425314
    # Model 3 loss = 6.045669
    # losses_vmapped = Array([3.762451 , 4.39846  , 4.1425314, 6.045669 ], dtype=float32)