torch.where()函数的一些应用

我们在使用torch对tensor进行变换的时候难免会在tensor中发现NAN，使得模型不能正常运行。下面使用torch.where()来解决这一问题。

首先我们创建一个tensor。

import torch
import numpy as np
a = torch.Tensor([[1, 2], [np.nan, 4], [3,5]])
#tensor([[1., 2.],
#       [nan, 4.],
#       [3., 5.]])

接下来我们使用如下函数将nan替换成0：

a = torch.where(torch.isnan(a), torch.full_like(a, 0), a)
#tensor([[1., 2.],
#        [0., 4.],
#        [3., 5.]])

这样问题就解决了。接下来我们在举一个例子，假设tensor中的元素必须大于0，若等于0择置为1。

代码如下：

a = torch.tensor([1,2,3,4,0,1,2,3,4,0])
#tensor([1, 2, 3, 4, 0, 1, 2, 3, 4, 0])
a = torch.where(a==0, 1, a)
#tensor([1, 2, 3, 4, 1, 1, 2, 3, 4, 1])

如有误，还请大家多多指正！