可以直接从 gpu 给 Caffe 或 Caffe2 输入数据吗?

Can Caffe or Caffe2 be given input data directly from gpu?

我已经阅读了 caffe2 教程并尝试了预训练模型。我知道 caffe2 会利用 GPU 运行 model/net。但是输入数据似乎总是从CPU(即主机)内存中给出。例如,在Loading Pre-Trained Models中,加载模型后,我们可以通过

预测图像
result = p.run([img])

但是,图像 "img" 应该在 CPU 范围内读取。我寻找的是一个框架,它可以将图像(从视频解码并仍然驻留在 GPU 内存中)直接传输到预测模型,而不是将其从 GPU 复制到 CPU 范围,然后传输到GPU 再次预测结果。 Caffe或Caffe2是否为python或C++提供了这样的功能或接口?还是我需要修补 Caffe 才能这样做?非常感谢。


这是我的解决方案:

我在 tensor.h 中发现,函数 ShareExternalPointer() 可以完全满足我的要求。

以这种方式提供 gpu 数据,

pInputTensor->ShareExternalPointer(pGpuInput, InputSize);

然后运行预测网通过

pPredictNet->Run();

其中 pInputTensor 是预测网的入口张量 pPredictNet

我不认为你可以在 和 python 界面中做到这一点。
但我认为它可以使用 c++ 来完成:在 c++ 中,您可以访问 Blobmutable_gpu_data()。您可以在设备上编写 运行 代码,并直接从 gpu "fill" 输入 Blob 的 mutable_gpu_data()。进行此更新后,caffe 应该能够从那里继续其 net->forward()

更新
2017 年 9 月 19 日 PR #5904 被并入 master。此 PR 通过 python 接口公开 blob 的 GPU 指针。
您可以直接从 python 访问 blob._gpu_data_ptrblob._gpu_diff_ptr,风险自负.

如您所述,使用 Python 层会强制数据进出 GPU,这会对性能造成巨大影响。这不仅适用于 Caffe,也适用于其他框架。要详细说明 Shai 的回答,您可以查看 this step-by-step tutorial on adding C++ layers to Caffe。给出的示例应该涉及处理层实现的大多数问题。披露:我是作者。