Question

如果我的模型仅包含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。追溯（最近一次通话）：文件 “ /opt/conda/lib/python3.6/site-packages/torch/nn/parallel/parallel_apply.py”， _worker中的第60行输出= module（* input，** kwargs）文件“ /opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py”，第541行，在致电中结果= self.forward（* input，** kwargs）文件“”，第7行，向前返回self.layer（x）@ self.weight RuntimeError：参数位于以下位置的不同GPU上 /pytorch/aten/src/THC/generic/THCTensorMathBlas.cu:277

当模型中有一些张量运算时，如何避免此错误？

Answer 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

Answer 2

由于OP也想冻结该层，因此在@Elgar de Groot中添加了答案。为此，您仍然可以使用torch.nn.Parameter，但是您可以像下面这样显式将require_grad设置为false：

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

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

2 个答案: