我尝试比较两个位于GPU上的torch.FloatTensor(只有一个条目):
if(FloatTensor_A> FloatTensor_B):做点什么
问题是,(FloatTensor_A > FloatTensor_B)
回复了ByteTensor
。有没有办法在这两个标量FloatTensors之间进行布尔比较,而无需在CPU上加载张量并将它们转换回numpy或常规浮点数?
答案 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