当模型包含张量操作时，Pytorch DataParallel 不起作用答案

【问题标题】：Pytorch DataParallel doesn't work when the model contain tensor operation当模型包含张量操作时，Pytorch DataParallel 不起作用
【发布时间】：2020-07-03 02:17:28
【问题描述】：

如果我的模型仅包含 nn.Module 层，例如 nn.Linear，则 nn.DataParallel 工作正常。

x = torch.randn(100,10)

class normal_model(torch.nn.Module):
    def __init__(self):
        super(normal_model, self).__init__()
        self.layer = torch.nn.Linear(10,1)

    def forward(self, x):
        return self.layer(x)

model = normal_model()
model = nn.DataParallel(model.to('cuda:0'))
model(x)

但是，当我的模型包含如下张量操作时

class custom_model(torch.nn.Module):
    def __init__(self):
        super(custom_model, self).__init__()
        self.layer = torch.nn.Linear(10,5)
        self.weight = torch.ones(5,1, device='cuda:0')
    def forward(self, x):
        return self.layer(x) @ self.weight

model = custom_model()
model = torch.nn.DataParallel(model.to('cuda:0'))
model(x)

它给了我以下错误

RuntimeError：在设备 1 上的副本 1 中捕获 RuntimeError。原始 Traceback（最近一次通话最后一次）：文件 "/opt/conda/lib/python3.6/site-packages/torch/nn/parallel/parallel_apply.py", 第 60 行，在 _worker 输出=模块（*输入，**kwargs）文件“/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py”，第 541 行，在调用 result = self.forward(*input, **kwargs) File "", line 7, in forward return self.layer(x) @self.weight RuntimeError: 参数位于不同的 GPU 上 /pytorch/aten/src/THC/generic/THCTensorMathBlas.cu:277

当我们的模型中有一些张量操作时，如何避免这个错误？

【问题讨论】：

标签： python-3.x pytorch

【解决方案1】：

我没有使用DataParallel 的经验，但我认为这可能是因为您的张量不是模型参数的一部分。你可以这样写：

torch.nn.Parameter(torch.ones(5,1))

请注意，您不必在初始化时将其移动到 gpu，因为现在当您调用 model.to('cuda:0') 时，这是自动完成的。

我可以想象DataParallel 使用模型参数将它们移动到适当的 gpu。

有关火炬张量和torch.nn.Parameter 之间区别的更多信息，请参阅this answer。

如果您不希望在训练期间通过反向传播更新张量值，您可以添加requires_grad=False。

另一种可行的方法是覆盖to 方法，并在前向传递中初始化张量：

class custom_model(torch.nn.Module):
    def __init__(self):
        super(custom_model, self).__init__()
        self.layer = torch.nn.Linear(10,5)
    def forward(self, x):
        return self.layer(x) @ torch.ones(5,1, device=self.device)
    def to(self, device: str):
        new_self = super(custom_model, self).to(device)
        new_self.device = device
        return new_self

或类似的东西：

class custom_model(torch.nn.Module):
    def __init__(self, device:str):
        super(custom_model, self).__init__()
        self.layer = torch.nn.Linear(10,5)
        self.weight = torch.ones(5,1, device=device)
    def forward(self, x):
        return self.layer(x) @ self.weight
    def to(self, device: str):
        new_self = super(custom_model, self).to(device)
        new_self.device = device
        new_self.weight = torch.ones(5,1, device=device)
        return new_self

【讨论】：

如果我使用torch.nn.Parameter(torch.ones(5,1))，这个张量的梯度会被记录下来，并且会在梯度下降时更新，是吗？使用 torch.nn.Parameter() 确实有效，但我不希望它是可训练的。
Cna 有人建议如何使这个不可训练，这样它就不会训练用 nn.Prameter 包裹的那些？
只需将其.requires_grad 属性设置为False.

【解决方案2】：

添加到@Elgar de Groot 的答案，因为 OP 也想冻结该层。为此，您仍然可以使用 torch.nn.Parameter 但随后您将 requires_grad 显式设置为 false ，如下所示：

self.layer = torch.nn.Parameter(torch.ones(5,1))
self.layer.requires_grad = False

【讨论】：