python - 使用 PyTorch 在多个 GPU 上并行进行 Tensor Inverse

Question

我想在多个 GPU 上并行运行 torch.inverse()，但我不能。

我在多个 GPU 上看到了这篇文章 Matmul，它介绍了 matmul 的过程。它表明，如果您为每个 GPU 分配了多个张量，matmul 将并行运行。我能够为 matmul 复制这种行为，但是当我尝试对 torch.inverse() 做同样的事情时，当我检查“watch nvidia-smi”时，它似乎是按顺序运行的。此外，当我用 torch fft 函数替换 torch.inverse() 函数时，我得到了并行 GPU 使用。有任何想法吗？

import torch

ngpu = torch.cuda.device_count()
# This is the allocation to each GPU.
lis = []
for i in range(ngpu):
    lis.append(torch.rand(5000,5000,device = 'cuda:'+ str(i)))

# per the matmul on multiple GPUs post this should already be in parallel
# but doesnt seem to be based on watch nvidia-smi
C_ = []
for i in range(ngpu):
    C_.append(torch.inverse(lis[i]))

编辑：这可以与上面链接中的 FFT 代码（下面）和 Matmul 代码进行比较。

import torch

ngpu = torch.cuda.device_count()
# This is the allocation to each GPU.
lis = []
for i in range(ngpu):
    lis.append(torch.rand(5000,5000,2,device = 'cuda:'+ str(i)))

# per the matmul on multiple GPUs post this should already be in parallel
# but doesnt seem to be based on watch nvidia-smi
C_ = []
for i in range(ngpu):
    C_.append(torch.fft(lis[i],2))

标签： pythongpupytorch