Python 输出中带有南的亚麻神经网络

发布于04月16日

我正在使用Flax训练神经网络.我的训练数据在输出中有大量的nan.我想忽略这些，只使用非nan值进行训练.为了实现这一点，我try 使用jnp.nanmean来计算损失，即:

def nanloss(params, inputs, targets):
    pred = model.apply(params, inputs)
    return jnp.nanmean((pred - targets) ** 2)

def train_step(state, inputs, targets):
    loss, grads = jax.value_and_grad(nanloss)(state.params, inputs, targets)
    state = state.apply_gradients(grads=grads)
    return state, loss

然而，经过一个训练步骤，损失就没有了.

我正在努力实现的目标是否可能？如果是这样，我该如何解决这个问题？

推荐答案

我怀疑您正在解决这里讨论的问题:JAX FAQ: gradients contain NaN where using where.您已经在计算本身中处理了NaN，但由于autodiff的实现方式，它们仍然潜入梯度中.

如果这确实是问题所在，您可以通过在计算损失之前过滤值来解决这个问题;例如:

def nanloss(params, inputs, targets):
    pred = model.apply(params, inputs)
    mask = jnp.isnan(pred) | jnp.isnan(targets)
    pred = jnp.where(mask, 0, pred)
    targets = jnp.where(mask, 0, targets)
    return jnp.mean((pred - targets) ** 2, where=~mask)