将torch.nn.DataParallel与自定义CUDA扩展一起使用

Question

根据我的理解,内置的pytorch操作都可以通过隐式矢量化自动处理批处理,允许跨多个GPU进行并行处理.

但是,根据文档在CUDA中编写自定义操作时,给定的LLTM示例执行批处理不变的操作,例如,计算元素的Sigmoid函数的梯度.

但是,我有一个不是批处理元素不变的用例而不是vectorisable.在单个GPU上运行,我目前(效率低下)遍历批处理中的每个元素,为每个元素执行内核启动,就像这样(用浏览器编写,只是为了演示):

std::vector<at::Tensor> op_cuda_forward(at::Tensor input, 
                                        at::Tensor elementSpecificParam) {

    auto output = at::zeros(torch::CUDA(/* TYPE */), {/* DIMENSIONS */});

    const size_t blockDim = //
    const size_t gridDim = //
    const size_t = numBatches = //

    for (size_t i = 0; i < numBatches; i++) {
        op_cuda_forward_kernel<T><<<gridDim, blockDim>>>(input[i],
                                                         elementSpecificParam[i], 
                                                         output[i]);
    }

    return {output};
}

但是,我希望通过批处理元素将此操作拆分为多个GPU.

如何output在多GPU场景中分配Tensor？

当然,可以在启动适当的内核之前在每个GPU上创建中间Tensors,但是将输入数据复制到每个GPU并且再次返回的开销将是有问题的.

有没有一种更简单的方法来启动内核而不首先探测GPU信息的环境(#GPU等)？

最终目标是使用与torch.nn.DataParallel一起使用的CUDA操作.

Answer 1

这有点不寻常，因为通常“批量”被准确定义为网络的所有操作沿该维度不变。例如，您可以引入另一个维度。因此，您拥有“前批次维度”，其中您的操作不是不变的。为此，请保留当前的实施。然后，在多个“实际批次”数据的新维度上并行化。

但是，为了更接近您提出的问题，我看到两个选择：

• 正如您所说，在您的实现中找出您正在操作的原始批次（取决于并行拆分的总数等）。这可能会变得毛茸茸的。
• 将您的参数视为输入的一部分！在外部调用中，将参数沿着输入数据传递到模型的前端。所以（类似Python的伪代码）：

Network(nn.Module):
  ...
  def forward(x, parameter):
    x=self.pre_modules(x)
    x=self.custom_module(x,parameter)
    return x

parameter=torch.zeros(16,requires_grad=True)
net=nn.DataParallel(model)
net(input,parameter)

如果您愿意接受这将是网络的泄漏抽象，并且主要对让事情正常工作感兴趣，我会首先尝试后一种方法。