在多个PyTorch张量运算符上运行GPU

Question

我有以下PyTorch张量：

X = np.array([[1, 3, 2, 3], [2, 3, 5, 6]])
X = torch.FloatTensor(X).cuda()

我想知道如果我在一行中运行以下多个PyTorch运算符，则下面的方案A或B之间是否存在任何差异（尤其是速度）？

方案A：

X_sq_sum = (X**2).cuda().sum(dim = 1).cuda()

场景B：

X_sq_sum = (X**2).sum(dim = 1).cuda()

即。方案A有两个.cuda()，而方案B只有一个.cuda()。

非常感谢。

Answer 1

它们的性能相同，因为CUDA转换只完成一次。

如docs中所述，如果对象已经在CUDA内存中，则重复的.cuda()调用将成为无操作。