在 GPU 上与 torch.FloatTensor 进行比较的优雅方式

Question

我试着比较位于 GPU 上的两个 torch.FloatTensor（只有一个条目）：

if (FloatTensor_A > FloatTensor_B): do something

问题是，(FloatTensor_A > FloatTensor_B) 回馈了 ByteTensor。有没有办法在这两个标量 FloatTensors 之间进行布尔比较，而不用在 CPU 上加载张量并将它们转换回 numpy 或传统浮点数？

Answer 1

是的，很简单。

示例：

In [24]: import os

# select `GPU 0` for the whole session
In [25]: os.environ['CUDA_VISIBLE_DEVICES'] = '0'

# required `data type` (for GPU) 
In [26]: dtype = torch.cuda.FloatTensor

# define `x` & `y` directly on GPU
In [27]: x = torch.randn(100, 100).type(dtype)
In [28]: y = torch.randn(100, 100).type(dtype)

# stay on GPU with desired `dtype`
In [31]: x.gt(y).type(dtype)
Out[31]: 

    0     1     1  ...      0     0     0
    1     0     0  ...      1     0     1
    1     1     1  ...      0     0     0
       ...          ⋱          ...       
    1     1     1  ...      0     0     0
    0     1     1  ...      1     1     1
    1     0     1  ...      1     0     1
[torch.cuda.FloatTensor of size 100x100 (GPU 0)]


# sanity check :)
In [33]: x.gt(y).type(dtype).is_cuda
Out[33]: True

Answer 2

PyTorch 中的比较操作return ByteTensors（参见docs）。为了将结果转换回浮点数据类型，您可以对结果调用 .float()。例如：

(t1 > t2).float()

(t1 > t2) 会 return 一个 ByteTensor.

操作的输入必须在同一内存上（CPU 或 GPU）。 return 结果将在同一内存中。当然，任何Tensor都可以通过调用.cpu()或.cuda()移动到各自的内存中。

在 GPU 上与 torch.FloatTensor 进行比较的优雅方式

Elegant way to compare to torch.FloatTensor on GPU

python

numpy

boolean

pytorch

tensor