JAX中利用vmap并行化模型集成：理解PyTree与结构化数组模式

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

# ... (init_params, predict, batched_predict等函数定义)

def mse_loss(params, inputs, targets):
    preds = batched_predict(params, inputs)
    loss = jnp.mean((targets - preds) ** 2)
    return loss

# 假设 ensemble_params 是一个包含多个模型参数PyTree的列表
# ensemble_params = [params_model_0, params_model_1, ...]
ensemble_loss = jax.vmap(fun=mse_loss, in_axes=(0, None, None))
losses = ensemble_loss(ensemble_params, x, y)

# 假设 ensemble_params 是一个列表，每个元素是一个模型的参数PyTree
# ensemble_params = [params_model_0, params_model_1, ..., params_model_N-1]

# 使用 jax.tree_map 将所有模型的参数堆叠起来
# lambda *args: jnp.stack(args) 会对所有传入的PyTree中相同路径的叶子节点进行堆叠
batched_ensemble_params = 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,)) # 修正：使用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) # 修正：num为层数减一
    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)
    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] # 示例网络结构：输入2 -> 隐藏3 -> 输出4
    batch_size = 2 # 单次推理的输入数据批量大小

    key = random.PRNGKey(seed=1)
    key, subkey = random.split(key)

    # 1. 初始化模型集成参数 (列表结构)
    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 = {loss}")

    # 2. 尝试直接使用 vmap (会导致错误)
    print("\n--- 尝试直接对列表结构使用 vmap (预期会失败) ---")
    try:
        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"意外成功: {losses_fail}")
    except ValueError as e:
        print(f"捕获到预期错误: {e}")
        print("错误原因：vmap 尝试映射列表中的 PyTree 内部数组，而非列表本身。")

    # 3. 转换为结构化数组模式并使用 vmap (正确方法)
    print("\n--- 转换为结构化数组模式后使用 vmap (正确方法) ---")
    # 将列表结构 (list-of-structs) 转换为结构化数组 (struct-of-arrays)
    batched_ensemble_params = jax.tree_map(lambda *args: jnp.stack(args), *ensemble_params_list)

    # 现在 ensemble_loss 可以正确地映射了
    ensemble_loss_vmap_success = jax.vmap(fun=mse_loss, in_axes=(0, None, None))
    losses_success = ensemble_loss_vmap_success(batched_ensemble_params, x, y)
    print(f"正确计算的批量损失: {losses_success}")

    # 验证结果与for循环一致
    expected_losses = jnp.array([mse_loss(p, x, y) for p in ensemble_params_list])
    print(f"for循环计算的损失 (验证): {expected_losses}")
    print(f"两种方法结果是否一致: {jnp.allclose(losses_success, expected_losses)}")