[SOLVED] How to change blob value in Caffe2 while using GPU (C++)

How to change blob value in Caffe2 while using GPU (C++)

I am trying to manually access blobs value and change it.
I have a blob called "1conv1_w" and I access it by:

auto 1conv1_w = caffe2::TensorCPU((*workspace.GetBlob("1conv1_w")).Get<caffe2::TensorCPU>()).data<float>();

this will return a pointer float* to 1conv1_w. In CPU mode, I can use

std::cout << *1conv1_w << std::endl

to access the first value within Blob "1conv1_w" as well as modify the value. However, when turn into GPU mode, this will return error because there is no value within the pointer. If I use

auto 1conv1_w = caffe2::TensorCPU((*workspace.GetBlob("1conv1_w")).Get<caffe2::TensorCUDA>()).data<float>()[0];

Then I can access the first value but still cannot access other value within the blob.
I guess the problem is because while using GPU the memory is actually a temporary memory. The value is copied between CPU and GPU (probably memcpy). When I use Get<caffe2::TensorCUDA>() it just copy the address or value I want to me. So even I change the value in this address it will not affect the actual value saved in somewhere.

Does anyone face the same problem and know how to change the actual value of the blob?

Solution

First of all you cannot access GPU memory straight forward from CPU context. You may consider writing CUDA kernel for your purpose. If you really need to do it on CPU then you can get data from GPU to CPU with:

CPUContext context;
TensorCPU output;
auto& input = (*workspace.GetBlob("1conv1_w")).Get<TensorCUDA>();
output.ResizeLike(input);
context.CopyItems<CUDAContext, CPUContext>(
    input.meta(),
    input.size(),
    input.raw_data(),
    output.raw_mutable_data(input.meta()));

Then you can modify CPU version and put it back to GPU in the same analogical way.