caffe2 Tensor <CUDAContext>赋值,构造或复制

Question

这是一个很长的镜头,如果您认为问题太局部化,请投票结束.我搜索了caffe2 github存储库,打开了一个问题同样的问题,在caffe2_ccp_tutorials存储库中打开了另一个问题,因为它的作者似乎最了解它,阅读caffe2 :: Tensor和caffe2 :: CUDAContext上的doxygen文档,甚至通过caffe2 源代码,特别是tensor.h,context_gpu.h和context_gpu.cc.

我知道目前caffe2不允许将设备内存复制到张量.我愿意扩展库并执行拉取请求以实现此目的.我背后的原因是我使用cv::cuda::*在设备内存上运行的方法进行所有图像预处理,因此我认为在gpu上进行预处理显然是一个问题,只是将结果下载到主机上,然后让它从主机重新上传到网络.

看看Tensor<Context>我的构造函数,我可以看到

template<class SrcContext , class ContextForCopy > 
Tensor (const Tensor< SrcContext > &src, ContextForCopy *context)

可能会实现我想要的,但我不知道如何设置<ContextForCopy>然后用它来构建.

此外,我看到我可以用正确的尺寸构造Tensor,然后可以使用

template <typename T>
T* mutable_data()

我可以分配/复制数据.数据本身存储在std::vector<cv::cuda::GpuMat,因此我将不得不迭代它,然后使用cuda::PtrStepSz或cuda::PtrStep访问底层设备分配的数据.这是我需要复制/分配到的相同数据caffe2::Tensor<CUDAContext>.

我一直试图找出内部Tensor<CPUContext>被复制的内容,Tensor<CUDAContext>因为我已经看过它的例子,但我无法弄明白,尽管我认为使用的方法是CopyFrom.如上所述的常见示例,从CPU复制到GPU:

TensorCPU tensor_cpu(...);
TensorCUDA tensor_cuda = workspace.CreateBlob("input")->GetMutable<TensorCUDA>();
tensor_cuda->ResizeLike(tensor_cpu);
tensor_cuda->ShareData(tensor_cpu);

我很惊讶没有人遇到过这个任务,简短的搜索只会产生一个未解决的问题,即作者(@peterneher)或多或少地提出同样的问题.

Answer 1

我已经设法弄清楚了这一点。最简单的方法是告诉 OpenCV使用哪个内存位置。这可以通过使用构造函数的第七个和第八个重载cv::cuda::GpuMat来完成，如下所示：

cv::cuda::GpuMat::GpuMat(int    rows,
                         int    cols,
                         int    type,
                         void *     data,
                         size_t     step = Mat::AUTO_STEP 
                        )       

cv::cuda::GpuMat::GpuMat(Size   size,
                         int    type,
                         void *     data,
                         size_t     step = Mat::AUTO_STEP 
                        )       

这样做意味着caffe2::TensorCUDA已经预先声明并分配了：

std::vector<caffe2::TIndex> dims({1, 3, 224, 224});
caffe2::TensorCUDA tensor;
auto ptr = tensor.mutable_data<float>();
cv::cuda::GpuMat matrix(224, 224, CV_32F, ptr);

例如，使用以下方法处理 3 通道 BGR 浮点矩阵cv::cuda::split：

cv::cuda::GpuMat mfloat;
// TODO: put your BGR float data in `mfloat`
auto ptr = tensor.mutable_data<float>();
size_t width = mfloat.cols * mfloat.rows;
std::vector<cv::cuda::GpuMat> input_channels {
    cv::cuda::GpuMat(mfloat.rows, mfloat.cols, CV_32F, &ptr[0]),
    cv::cuda::GpuMat(mfloat.rows, mfloat.cols, CV_32F, &ptr[width]),
    cv::cuda::GpuMat(mfloat.rows, mfloat.cols, CV_32F, &ptr[width * 2])
};
cv::cuda::split(mfloat, input_channels);

希望这对熟悉 Caffe2 C++ 方面的人有所帮助。

注意，这caffe2::Predictor不适用于caffe2::TensorCUDA，您将不得不手动传播张量。有关这方面的更多信息，请参阅caffe2_cpp_tutorial mnist.cc。