python - 是否可以将 GPU 的所有内存用于一种型号？

Question

有一个模型和两个GPU。我使用 model.cuda() 将 model 放在 GPU 上。如果我将大图像传递给 model，它会分配 GPU0 的所有内存，然后引发 CUDA out of memory 错误，而不会分配 GPU1 的任何内存。

因为每个 forward() 只有一个图像，我不能使用这样的 torch.nn.DataParallel 东西来分割输入。

在将图像传递给模型时，有什么方法可以使用 GPU 的所有内存吗？

我正在使用 Python3.7 和 Pytorch1.1。

Answer 1

您可以将您的模型拆分为两个子模块。像这样:

class MyModel(nn.Module):
    def __init__(self, split_gpus):
        self.large_submodule1 = ...
        self.large_submodule2 = ...

        self.split_gpus = split_gpus
        if split_gpus:
            self.large_submodule1.cuda(0)
            self.large_submodule1.cuda(1)

    def forward(self, x):
        x = self.large_submodule1(x)
        if split_gpus:
            x = x.cuda(1) # P2P GPU transfer
        return self.large_submodule2(x)

来自 pytorch discuss