model.cuda () в pytorch

Question

Если я позвоню model.cuda() в pytorch, где модель является подклассом nn.Module, и скажу, что если у меня четыре GPU, как он будет использовать четыре GPU и как узнать, какие GPU используют?

Answer 1

Если у вас есть пользовательский модуль, полученный из nn.Module после model.cuda() всех параметров модели, (итератор model.parameters() может показать вам их) закончится на вашей куда.

Чтобы проверить, где находятся ваши параметры, просто напечатайте их (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

Чтобы расширить ответ прости , чтобы разделить вычисления между несколькими графическими процессорами, вы должны использовать torch.nn.DataParallel или DistributedDataParallel.