pytorch 中的 model.cuda()

Question

如果我调用model.cuda()pytorch，其中 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()你可以检查你有多少台设备。