根据我的理解，内置的pytorch操作都可以通过隐式矢量化自动处理批处理，从而允许跨多个GPU进行并行处理 .
但是，当根据Documentation在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 .
How would the allocation of the output Tensor work in a multi-GPU scenario?
当然，可以在启动适当的内核之前在每个GPU上创建中间Tensors，但是将输入数据复制到每个GPU并且再次返回的开销将是有问题的 .
Is there a simpler way to launch the kernels without first probing the environment for GPU information (# GPU's etc)?
最终目标是使用与torch.nn.DataParallel一起使用的CUDA操作 .