PyTorch：使用多个 GPU 对单个非常大的图像进行推断？

Question

我想在非常大的卫星图像上执行推理（即语义分割）而不将其分成几部分。 我可以使用 4 个 GPU（每个 GPU 有 15 GB 的内存）并且想知道是否有可能以某种方式使用这些 GPU 的所有内存组合（即 60 GB）来推断PyTorch的图像？

Answer 1

您正在寻找模型并行工作模式。
基本上，您可以将模型的不同部分分配给不同的 GPU，然后您应该处理“簿记”。
此解决方案非常特定于模型和特定于任务，因此没有“通用”包装器（与数据并行相反）。

例如：

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