python - Pytorch 将操作应用于某些行但不是全部

Question

我正在尝试在张量的某些行上应用 softmax 函数，但问题是我的某些行具有所有 -inf 值。因此，这些行上的 softmax 会输出 NaN，这会导致模型后期出现问题。

因此，我想创建一个将 softmax 应用于一行的函数，除非 in 全部为 -inf。在这种情况下，它输出一个零向量。有什么简单的方法可以做到这一点？

Answer 1

是否将所有 nan在 softmax 之后的所有行都设置为 0 对您有用？这样，您可以确保不会覆盖任何意外的 nan。

import numpy as np
import torch
import torch.nn.functional as F

array = np.arange(25, dtype=np.float32).reshape((5, 5))
array[3, ...] *= -np.inf

# [[  0.   1.   2.   3.   4.]
#  [  5.   6.   7.   8.   9.]
#  [ 10.  11.  12.  13.  14.]
#  [-inf -inf -inf -inf -inf]
#  [ 20.  21.  22.  23.  24.]]

array = torch.tensor(array)
array = F.softmax(array, dim=1)
mask = array.isnan().all(dim=1)
array[mask, ...] = 0
print(array)

# tensor([[0.0117, 0.0317, 0.0861, 0.2341, 0.6364],
#         [0.0117, 0.0317, 0.0861, 0.2341, 0.6364],
#         [0.0117, 0.0317, 0.0861, 0.2341, 0.6364],
#         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
#         [0.0117, 0.0317, 0.0861, 0.2341, 0.6364]])