PyTorch训练时loss突然变NaN，换TensorFlow就没问题，是我哪里写错了吗？

我之前也遇到过类似的情况，PyTorch这边loss爆炸但换TensorFlow就没事，后来发现是Keras默认在优化器里加了比较激进的epsilon，而PyTorch的Adam默认eps是1e-8，某些层梯度稍微大一点就容易数值不稳。你可以先试试把优化器换成SGD看看还炸不炸，如果SGD没事那基本就是自适应优化器的锅。另外图像分割里如果用了BatchNorm，PyTorch默认momentum是0.1，跟Keras的0.99差别挺大的，小batch下统计量波动大也会把loss带崩。还有一点容易忽略的是混合精度，PyTorch的autocast在某些op上会偷偷转fp16，你如果开了AMP建议先把loss scale调保守点或者干脆关掉排查。其实可以逐层打印grad norm，看看到底是哪一层先炸的，一般能定位到具体是norm层还是attention那块。TensorFlow那边数值稳不代表你PyTorch写错了，两家默认超参和算子实现差异真的挺坑的。

我之前也碰到过类似情况，PyTorch这边确实容易在几个地方翻车。你可以先查下是不是某个BN层在特定batch上方差特别小，PyTorch默认eps是1e-5，有时候不够用。另外优化器那边，Adam的eps默认也是1e-8，跟Keras可能不一样，调大一点试试。混合精度的话，PyTorch的GradScaler没配好也会炸，可以先关掉AMP跑一遍看看。

PyTorch这边确实容易踩到数值稳定性的坑，我遇到过类似的。你可以重点看看是不是在softmax或者log操作里出现了log(0)或者除以零的情况，TF那边默认会给一些epsilon兜底，PyTorch就得自己加。另外检查下优化器里eps参数，Adam默认1e-8有时候不够用，调到1e-6或1e-5试试。还有个隐蔽的点是AMP混合精度，如果开了的话某些算子会溢出，换回fp32跑几个epoch对比一下。

我前段时间也遇到过类似的情况，PyTorch跑到一半loss直接飞了，换成Keras同样结构就稳得很，当时也怀疑人生。后来发现一个挺隐蔽的点，PyTorch默认的算子实现和TF在数值上确实有细微差异，尤其在归一化层和softmax那块，FP16混合精度下特别容易放大。你可以先试试关掉AMP用纯FP32跑一遍，如果稳了那大概率就是精度问题，而不是你代码逻辑写错了。另外检查下损失函数里有没有log或除法操作，PyTorch对接近零的值处理比TF敏感，加个eps往往能救。还有优化器这块，Adam的eps默认值两边好像不一样，PyTorch是1e-8，有些TF实现是1e-7，这种小差别在训练后期梯度很小时会突然放大。梯度裁剪你加了，但裁剪的时机和位置也值得看看，是在backward之后optimizer.step之前吗？我之前就是顺序搞错了白裁了。实在不行可以在loss变NaN前打点看是哪一层先出的问题，比盲目调参高效多了。