[caffe-plugin] Fix Tshape/TBlob. Remove setblob. Update config.mk. (apache#2941)

Author: Taco-W

make/config.mk

@@ -111,8 +111,9 @@ EXTRA_OPERATORS =
 # plugins
 #----------------------------
 
-# whether to use caffe integration. This requires including caffe submodule.
-# CAFFE_PATH = caffe-lite 
+# whether to use caffe integration. This requires installing caffe.
+# You also need to add CAFFE_PATH/build/lib to your LD_LIBRARY_PATH
+# CAFFE_PATH = $(HOME)/caffe
 # MXNET_PLUGINS += plugin/caffe/caffe.mk
 
 # whether to use torch integration. This requires installing torch.

plugin/caffe/caffe_blob.cc

@@ -12,7 +12,7 @@ namespace caffe {
 template<>
 void SetDataGradToBlob<mshadow::cpu, float>(caffeMemoryTypes memType,
                             std::vector<::caffe::Blob<float>*>::iterator blob,
-                            std::vector<mshadow::TBlob>::const_iterator itr) {
+                            std::vector<TBlob>::const_iterator itr) {
   float *data_ptr = reinterpret_cast<float*>((*itr).dptr_);
   if (memType == Data)
     (*blob)->set_cpu_data(data_ptr);
@@ -23,7 +23,7 @@ void SetDataGradToBlob<mshadow::cpu, float>(caffeMemoryTypes memType,
 template<>
 void SetDataGradToBlob<mshadow::cpu, double>(caffeMemoryTypes memType,
                             std::vector<::caffe::Blob<double>*>::iterator blob,
-                            std::vector<mshadow::TBlob>::const_iterator itr) {
+                            std::vector<TBlob>::const_iterator itr) {
   double *data_ptr = reinterpret_cast<double*>((*itr).dptr_);
   if (memType == Data)
     (*blob)->set_cpu_data(data_ptr);
@@ -34,7 +34,7 @@ void SetDataGradToBlob<mshadow::cpu, double>(caffeMemoryTypes memType,
 template<>
 void SetDataGradToBlob<mshadow::gpu, float>(caffeMemoryTypes memType,
                             std::vector<::caffe::Blob<float>*>::iterator blob,
-                            std::vector<mshadow::TBlob>::const_iterator itr) {
+                            std::vector<TBlob>::const_iterator itr) {
   float *data_ptr = reinterpret_cast<float*>((*itr).dptr_);
   if (memType == Data)
     (*blob)->set_gpu_data(data_ptr);
@@ -45,27 +45,25 @@ void SetDataGradToBlob<mshadow::gpu, float>(caffeMemoryTypes memType,
 template<>
 void SetDataGradToBlob<mshadow::gpu, double>(caffeMemoryTypes memType,
                             std::vector<::caffe::Blob<double>*>::iterator blob,
-                            std::vector<mshadow::TBlob>::const_iterator itr) {
+                            std::vector<TBlob>::const_iterator itr) {
   double *data_ptr = reinterpret_cast<double*>((*itr).dptr_);
   if (memType == Data)
     (*blob)->set_gpu_data(data_ptr);
   else
     (*blob)->set_gpu_diff(data_ptr);
 }
 
-mshadow::TShape Vector2TShape(const std::vector<int> &vec_int) {
-  mshadow::TShape res;
-  std::vector<mshadow::index_t> vec_indx;
+TShape Vector2TShape(const std::vector<int> &vec_int) {
+  std::vector<mshadow::index_t> vec;
   for (int i = 0; i < vec_int.size(); ++i)
-    vec_indx.push_back(vec_int[i]);
+    vec.push_back(vec_int[i]);
   // 0-dim represents scalar in caffe
   if (vec_int.size() == 0)
-    vec_indx.push_back(1);
-  res = vec_indx;
-  return res;
+    vec.push_back(1);
+  return {vec.begin(), vec.end()};
 }
 
-std::vector<int> TShape2Vector(const mshadow::TShape &tshape) {
+std::vector<int> TShape2Vector(const TShape &tshape) {
   std::vector<int> s;
   for (int i =0 ; i < tshape.ndim(); ++i)
     s.push_back(tshape[i]);

plugin/caffe/caffe_blob.h

@@ -7,10 +7,10 @@
 #ifndef PLUGIN_CAFFE_CAFFE_BLOB_H_
 #define PLUGIN_CAFFE_CAFFE_BLOB_H_
 
-#include<mshadow/tensor.h>
-#include<mshadow/tensor_blob.h>
-#include<vector>
-#include<caffe/blob.hpp>
+#include <mxnet/tensor_blob.h>
+#include <vector>
+#include <caffe/blob.hpp>
+#include <caffe/layer.hpp>
 
 namespace mxnet {
 namespace op {
@@ -20,14 +20,14 @@ namespace caffe {
 // Declare Memory Type for Caffe blob
 enum caffeMemoryTypes {Data, Grad, Non};
 
-mshadow::TShape Vector2TShape(const std::vector<int> &vec_int);
-std::vector<int> TShape2Vector(const mshadow::TShape &tshape);
+TShape Vector2TShape(const std::vector<int> &vec_int);
+std::vector<int> TShape2Vector(const TShape &tshape);
 
 // implementation of tensor to blob, called by TensorToBlob
 template<typename Device, typename Dtype>
 void SetDataGradToBlob(caffeMemoryTypes memType,
                        typename std::vector< ::caffe::Blob<Dtype>*>::iterator blob,
-                       typename std::vector<mshadow::TBlob>::const_iterator itr);
+                       typename std::vector<TBlob>::const_iterator itr);
 
 /**
  * \brief The interface to convert mxnet's tensor to caffe's blob
@@ -36,14 +36,22 @@ void SetDataGradToBlob(caffeMemoryTypes memType,
 template<typename Device, typename Dtype>
 void TBlob2CaffeBlob(caffeMemoryTypes memType,
                      typename std::vector< ::caffe::Blob<Dtype>*>::iterator blob,
-                     typename std::vector<mshadow::TBlob>::const_iterator tblob,
+                     typename std::vector<TBlob>::const_iterator tblob,
                      int n = 1) {
   for (int i = 0; i < n; ++i, ++blob, ++tblob) {
     (*blob)->Reshape(TShape2Vector((*tblob).shape_));
     SetDataGradToBlob<Device, Dtype>(memType, blob, tblob);
   }
 }
 
+template<typename Dtype>
+void SetOpBlobs(::caffe::Layer<Dtype> *caffeOp,
+                const std::vector< ::caffe::Blob<Dtype>*>& weights) {
+  CHECK_EQ(caffeOp->blobs().size(), weights.size());
+  for (int i = 0; i < weights.size(); ++i)
+    caffeOp->blobs()[i].reset(weights[i]);
+}
+
 }  // namespace caffe
 }  // namespace op
 }  // namespace mxnet

plugin/caffe/caffe_op-inl.h

@@ -110,7 +110,7 @@ class CaffeOp : public Operator {
                       wei_.begin(),
                       in_data.begin() + param_.num_data,
                       param_.num_weight);
-      caffeOp_->SetBlobs(wei_);
+      caffe::SetOpBlobs(caffeOp_, wei_);
     }
 
     caffeOp_->Forward(bot_, top_);