从 PyTorch N 维张量中过滤出 NaN 值

Question

这个问题与在 -Dimensional 张量中过滤np.nanpytorch 中的值非常相似。不同之处在于我想将相同的概念应用于 2 维或更高维度的张量。

我有一个如下所示的张量：

import torch

tensor = torch.Tensor(
[[1, 1, 1, 1, 1],
 [float('nan'), float('nan'), float('nan'), float('nan'), float('nan')],
 [2, 2, 2, 2, 2]]
)

>>> tensor.shape
>>> [3, 5]

我想找到最Pythonic/PyTorch的方法来过滤（删除）张量的行nan。tensor通过沿第一个（第轴）过滤它，0我想获得filtered_tensor如下所示的：

>>> print(filtered_tensor)
>>> torch.Tensor(
[[1, 1, 1, 1, 1],
 [2, 2, 2, 2, 2]]
)
>>> filtered_tensor.shape
>>> [2, 5]

Answer 1

将 PyTorchisnan()与 一起使用，使用获得的布尔掩码对 的行any()进行切片，如下所示：tensor

filtered_tensor = tensor[~torch.any(tensor.isnan(),dim=1)]

nan请注意，这将删除其中包含值的任何行。如果您只想删除所有值都nan替换torch.any为 的行torch.all。

对于 N 维张量，您可以展平除第一个维度之外的所有维度，并应用与上面相同的过程：

#Flatten:
shape = tensor.shape
tensor_reshaped = tensor.reshape(shape[0],-1)
#Drop all rows containing any nan:
tensor_reshaped = tensor_reshaped[~torch.any(tensor_reshaped.isnan(),dim=1)]
#Reshape back:
tensor = tensor_reshaped.reshape(tensor_reshaped.shape[0],*shape[1:])