Python — Может ли Caffe или Caffe2 получать входные данные напрямую из gpu?

Я прочитал руководства по caffe2 и попробовал предварительно обученные модели. Я знал, что caffe2 запустит GPU для запуска модели / сети. Но входные данные, кажется, всегда передаются из памяти процессора (т.е. хоста). Например, в Загрузка предварительно обученных моделей, после загрузки модели мы можем предсказать изображение

result = p.run([img])

Тем не менее, изображение «img» должно быть прочитано в объеме процессора. То, что я ищу, — это структура, которая может подвести изображения (которые декодируются из видео и все еще находится в памяти GPU) непосредственно в модель прогнозирования, вместо того, чтобы копировать ее из GPU в область CPU, а затем снова передавать в GPU для прогнозирования результат. Caffe или Caffe2 предоставляют такие функции или интерфейсы для python или C ++? Или я должен исправить патч Caffe, чтобы сделать это? Спасибо вообще.

Вот мое решение:

Я нашел в tensor.hфункция ShareExternalPointer() могу делать то, что хочу.

Поток данных GPU таким образом,

pInputTensor->ShareExternalPointer(pGpuInput, InputSize);

затем прогнать прогнозируемую сеть через

pPredictNet->Run();

где pInputTensor является входным тензором для прогнозируемой сети pPredictNet