Caffe2 - C++ API
A deep learning, cross platform ML framework
blob_serialization_gpu.cc
1 
17 #include "caffe2/core/blob.h"
18 #include "caffe2/core/blob_serialization.h"
19 #include "caffe2/core/context_gpu.h"
20 
21 namespace caffe2 {
22 
23 template <>
24 void TensorSerializer<CUDAContext>::StoreDeviceDetail(
25  const Tensor<CUDAContext>& input, TensorProto* proto) {
26  auto* device_detail = proto->mutable_device_detail();
27  device_detail->set_device_type(CUDA);
28  device_detail->set_cuda_gpu_id(
29  GetGPUIDForPointer(input.raw_data()));
30 }
31 
32 namespace {
33 REGISTER_BLOB_SERIALIZER(
34  (TypeMeta::Id<TensorCUDA>()),
35  TensorSerializer<CUDAContext>);
36 REGISTER_BLOB_DESERIALIZER(TensorCUDA, TensorDeserializer<CUDAContext>);
37 }
38 } // namespace caffe2
39 
Copyright (c) 2016-present, Facebook, Inc.
int GetGPUIDForPointer(const void *ptr)
Gets the GPU id that the current pointer is located at.
Definition: common_gpu.cc:149
Copyright (c) 2016-present, Facebook, Inc.