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

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

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

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

std::vector 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 <<>>(input[i],
                                                         elementSpecificParam[i], 
                                                         output[i]);
    }

    return {output};
}

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

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

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

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

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