PyTorch 的 DataParallel 仅使用一个 GPU

Question

我正在尝试使用两个 GPU 在 PyTorch 中训练模型。我正在使用torch.nn.DataParallel，但由于某种原因nvidia-smi说我只使用一个 GPU。

该代码大致如下：

>>> import torch.nn as nn 
>>> model = SomeModel()
>>> model = nn.DataParallel(model)
>>> model.to('cuda')

当我运行程序并观察 的输出时nvidia-smi，我只看到 GPU 0 正在运行。有人知道问题是什么吗？

Answer 1

遇到了同样的问题，并通过意识到我使用错误来解决它。也许这会帮助某人节省一些时间：

如果您像这样包装模型，model = nn.DataParallel(model)它将仅包装对此模型的调用（调用意味着model(input)）。因此，如果您的模型有一个线性层，并且您现在直接使用它，就像model.my_linear(input)它不会被包装一样。

如果您在代码中传递对模型实例的引用，您最终可能会得到对未包装模型的引用，这也将不起作用。

最简单的测试是在调用模型之前立即包装模型并检查其是否有效。如果是这样，则您之前的包装代码可能存在上述问题之一。