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

Question

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

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

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

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

标签： pythongpupytorch

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 讨论