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

Question

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

由于每个forward()仅存在一张图像，因此我无法使用此类torch.nn.DataParallel来分割输入。

将图像传递到model时是否可以使用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