Caffe或Caffe2可以直接从gpu获得输入数据吗？

Question

我读过caffe2的教程，并尝试过预先培训过的模型。我知道caffe2会利用GPU来运行模型/网络。但是输入数据似乎总是从CPU（即主机）内存中给出的。例如，在Loading Pre-Trained Models中，加载模型后，我们可以通过

result = p.run([img])

但是，图像“img”应该在CPU范围内读取。我所寻找的是一个框架，它可以将图像（从视频解码后仍然驻留在GPU内存中）直接传输到预测模型，而不是将其从GPU复制到CPU范围，然后再传输到GPU来预测结果。咖啡因还是咖啡因2提供了Python或C++的功能或接口？或者我应该给咖啡馆打补丁吗？谢谢。在

---

我的解决方案是：

我发现在tensor.h中，ShareExternalPointer()函数可以做我想要的。在

以这种方式传送gpu数据

^{pr2}$

然后运行预测网络

pPredictNet->Run();

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

Answer 1

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