c++ - caffe2 Tensor<CUDAContext> 赋值、构造或复制

Question

这是一个远景，如果您认为问题过于局部化，请投票关闭。我在 caffe2 github 上搜索过存储库，已打开 an issue问同样的问题，在 caffe2_ccp_tutorials 打开了另一个问题存储库，因为它的作者似乎最了解它，请阅读 caffe2::Tensor 上的 doxygen 文档和 caffe2::CUDAContext , 甚至浏览了 caffe2 source code ，特别是 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>然后用于构建。

此外，我看到我可以用正确的尺寸构造张量，然后也许使用

template <typename T>
T* mutable_data()

我可以分配/复制数据。 数据本身存储在 std::vector<cv::cuda::GpuMat 中，所以我将不得不迭代它，然后使用 cuda::PtrStepSz or 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);

我很惊讶还没有人遇到这个任务，简短的搜索只产生一个 open issue作者 (@peterneher) 或多或少在问同样的事情。

Answer 1

我已经设法解决了这个问题。 最简单的方法是告诉 OpenCV 使用哪个内存位置。 这可以通过使用 7th and 8th overload of the cv::cuda::GpuMat constructor 来完成。如下所示:

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);

例如，使用 cv::cuda::split 处理 3 channel BGR 浮点矩阵:

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 一起使用，您必须手动传播张量。 有关这方面的更多信息，the caffe2_cpp_tutorial mnist.cc .