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

提问者：小点点

我有以下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()。

事先多谢。

共1个答案

匿名用户

由于CUDA转换只进行一次，因此它们的性能相同。

如文档中所述，如果对象已经在CUDA内存中，则重复的.CUDA()调用将是无效的。