PyTorch: вывод на одном очень большом изображении с использованием нескольких графических процессоров?

Я хочу сделать вывод (то есть семантическую сегментацию) на очень большом спутниковом изображении, не разбивая его на части. У меня есть доступ к 4 графическим процессорам (каждый из которых имеет 15 ГБ памяти), и мне было интересно, можно ли каким-то образом использовать всю память этих графических процессоров вместе взятых (т.е. 60 ГБ) для вывода изображения в PyTorch?


person Federico    schedule 16.09.2020    source источник


Ответы (1)


Вы ищете режим работы параллельного моделирования.
По сути, вы можете назначьте разные части вашей модели разным графическим процессорам, а затем вы должны позаботиться о бухгалтерском учете.
Это решение очень специфично для модели и задачи, поэтому для него не существует общих оболочек. (в отличие от параллельных данных).

Например:

class MyModelParallelNetwork(nn.Module):
  def __init__(self, ...):
    # define the network
    self.part_one = ... # some nn.Module
    self.part_two = ... # additional nn.Module 
    self.part_three = ... 
    self.part_four = ...

    # important part - "send" the different parts to different GPUs
    self.part_one.to(torch.device('gpu:0'))
    self.part_two.to(torch.device('gpu:1'))
    self.part_three.to(torch.device('gpu:2'))
    self.part_four.to(torch.device('gpu:3'))

  def forward(self, x):
    # forward through model parts and GPUs:
    p1 = self.part_one(x.to(torch.device('gpu:0')))
    p2 = self.part_two(p1.to(torch.device('gpu:1')))
    p3 = self.part_three(p2.to(torch.device('gpu:2')))
    y = self.part_four(p3.to(torch.device('gpu:3')))
    return y  # result is on cuda:3 device
person Shai    schedule 07.04.2021
comment
Это предполагает, что x подходит для одного графического процессора, верно? Я думаю, что ОП спрашивает, можно ли использовать изображения, которые не помещаются в один графический процессор. - person xro7; 07.04.2021
comment
@xro7 этот подход предполагает, что x, модель и все промежуточные карты объектов не подходят для одного графического процессора, но могут поместиться на 4. Если вы не можете поместить саму x в графический процессор памяти, чем вы собираетесь обрабатывать x, не разбивая ее на части??? - person Shai; 07.04.2021
comment
Полностью понял, что вы говорите. Но я думаю, что проблема ОП именно в этом, и простой ответ может быть: вы не можете: D. Может быть, я неправильно понял. - person xro7; 07.04.2021