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

时间:2017-11-03 13:15:24

标签: python numpy boolean pytorch tensor

我尝试比较两个位于GPU上的torch.FloatTensor(只有一个条目):

  

if(FloatTensor_A> FloatTensor_B):做点什么

问题是,(FloatTensor_A > FloatTensor_B)回复了ByteTensor。有没有办法在这两个标量FloatTensors之间进行布尔比较,而无需在CPU上加载张量并将它们转换回numpy或常规浮点数?

2 个答案:

答案 0 :(得分:4)

PyTorch中的比较操作返回ByteTensors(参见docs)。要将结果转换回float数据类型,可以在结果上调用.float()。例如:

(t1 > t2).float()

(t1 > t2)将返回ByteTensor

操作的输入必须位于同一内存(CPU或GPU)上。返回结果将在同一内存中。当然,任何Tensor都可以通过callin .cpu().cuda()移动到相应的内存中。

答案 1 :(得分: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