Question

如果我在pytorch中调用model.cuda()，其中model是nn.Module的子类，并说如果我有四个GPU，它将如何利用这四个GPU，以及如何知道正在使用的GPU ？

Answer 1

如果您有一个在nn.Module所有模型参数之后从model.cuda()派生的自定义模块，（model.parameters()迭代器可以向您显示这些参数）将在您的cuda上结束。

要检查您的参数在哪里，只需在我的情况下将它们打印出来（cuda：0）：

class M(nn.Module):
    'custom module'
    def __init__(self):
        super().__init__()
        self.lin = nn.Linear(784, 10)

m = M()
m.cuda()
for _ in m.parameters():
    print(_)

# Parameter containing:
# tensor([[-0.0201,  0.0282, -0.0258,  ...,  0.0056,  0.0146,  0.0220],
#         [ 0.0098, -0.0264,  0.0283,  ...,  0.0286, -0.0052,  0.0007],
#         [-0.0036, -0.0045, -0.0227,  ..., -0.0048, -0.0003, -0.0330],
#         ...,
#         [ 0.0217, -0.0008,  0.0029,  ..., -0.0213,  0.0005,  0.0050],
#         [-0.0050,  0.0320,  0.0013,  ..., -0.0057, -0.0213,  0.0045],
#         [-0.0302,  0.0315,  0.0356,  ...,  0.0259,  0.0166, -0.0114]],
#        device='cuda:0', requires_grad=True)
# Parameter containing:
# tensor([-0.0027, -0.0353, -0.0349, -0.0236, -0.0230,  0.0176, -0.0156,  0.0037,
#          0.0222, -0.0332], device='cuda:0', requires_grad=True)

您还可以这样指定设备：

m.cuda('cuda:0')

使用torch.cuda.device_count()，您可以检查您拥有多少台设备。

Answer 2

要在prosti's answer上进行扩展以在多个GPU之间划分计算，应使用torch.nn.DataParallel或DistributedDataParallel。

pytorch中的model.cuda（）

2 个答案: