JAX记录MLP激活值遇TracerArrayConversionError的解决求助

我在使用JAX官方教程的MLP训练MNIST数据集时，尝试添加代码记录除最后一层外所有层的激活模式，但运行训练代码时持续报错。

修改后的代码 #

from collections import defaultdict
# this is my activation pattern logger
class ActivationLogger:
    def __init__(self, epoch):
       self.reset(epoch)

    def __call__(self, layer, activations):
        D = activations.shape[0]
        for i in range(D):
            self.activations[(layer, i)].append(
                    jax.lax.stop_gradient(activations[i]))

    def reset(self, epoch):
        self.epoch = epoch
        self.activations = defaultdict(list)

activation_logger = ActivationLogger(epoch=1)

...

def predict(params, image):
    # per-example predictions
    activations = image
    for l, (w, b) in enumerate(params[:-1]):
        outputs = jnp.dot(w, activations) + b
        activations = jnp.maximum(0, outputs)
        activation_logger(l+1, activations) # <- this was added

    final_w, final_b = params[-1]
    logits = jnp.dot(final_w, activations) + final_b
    return logits - logsumexp(logits)

batched_predict = jax.vmap(
        predict, 
        in_axes=(None, 0), 
        out_axes=0)

@jax.jit
def loss(params, images, targets):
  preds = batched_predict(params, images)
  return -jnp.mean(preds * targets)

错误信息 #

TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on traced array with shape float32[].
This BatchTracer with object id 7541955024 was created on line:
  /var/folders/km/3nj8tmq56s16dsgc9_63530r0000gn/T/ipykernel_73750/3494160804.py:12:20 (ActivationLogger.__call__)

解决建议 #

错误原因

JAX的JIT编译会追踪数组操作，你在JIT编译的loss函数内部调用activation_logger，尝试将追踪中的Tracer数组存入Python列表，这违反了JAX的纯函数规则：JIT编译的函数不能有副作用（比如修改外部状态），且Tracer数组不能直接转换为NumPy数组或存入Python容器。

具体解决步骤

• 移出JIT范围记录激活
  不要在JIT编译的predict或loss函数内调用日志逻辑。可以在训练循环中，每轮训练完成后，单独运行一次不带JIT的前向传播来收集激活数据，避免干扰JIT编译的纯函数环境。

• 使用host_callback处理JIT内日志（调试用）
  如果需要在JIT运行时记录，可以用jax.experimental.host_callback将激活数据传递到主机端处理，避免直接修改外部状态：

  from jax.experimental import host_callback
  
  def log_activation(args):
      layer, activations = args
      # 这里可以将激活数据写入日志或存储
      print(f"Layer {layer} activation shape: {activations.shape}")
  
  # 在predict函数中替换原logger调用：
  host_callback.id_tap(log_activation, (l+1, activations))
  

  注意：该方法会带来性能开销，适合调试场景。

• 避免JIT函数内修改外部状态
  ActivationLogger是带状态的对象，在JIT函数内调用其方法修改内部列表属于副作用操作，JAX不允许此类行为。必须保证JIT编译的函数是纯函数，仅依赖输入参数，不修改外部状态。

• 调整日志收集时机
  在训练循环中选取少量样本，单独运行不带JIT的前向传播来收集激活：

  # 训练循环示例
  for epoch in range(num_epochs):
      # ... 原有的JIT训练步骤
      # 收集当前epoch的激活
      activation_logger.reset(epoch)
      # 定义不带JIT的预测函数用于日志
      def predict_with_log(params, image):
          activations = image
          for l, (w, b) in enumerate(params[:-1]):
              outputs = jnp.dot(w, activations) + b
              activations = jnp.maximum(0, outputs)
              activation_logger(l+1, activations)
          final_w, final_b = params[-1]
          logits = jnp.dot(final_w, activations) + final_b
          return logits - logsumexp(logits)
      batched_predict_with_log = jax.vmap(predict_with_log, in_axes=(None, 0), out_axes=0)
      # 取小批量样本避免内存过载
      sample_images = train_images[:100]
      batched_predict_with_log(params, sample_images)
      # 处理收集到的激活数据
      print(f"已收集第{epoch}轮激活数据")
  

内容的提问来源于stack exchange，提问作者MoneyBall