| /*! |
| * Copyright (c) 2016 by Contributors |
| * \file caffe_op-inl.h |
| * \brief Caffe Operator |
| * \author Haoran Wang |
| */ |
| #ifndef PLUGIN_CAFFE_CAFFE_OP_INL_H_ |
| #define PLUGIN_CAFFE_CAFFE_OP_INL_H_ |
| |
| #include <dmlc/logging.h> |
| #include <dmlc/parameter.h> |
| #include <mxnet/operator.h> |
| #include <caffe/proto/caffe.pb.h> |
| |
| #include <map> |
| #include <vector> |
| #include <string> |
| #include <utility> |
| |
| #include "../../src/operator/operator_common.h" |
| #include "caffe_common.h" |
| #include "caffe_stream.h" |
| #include "caffe_fieldentry.h" |
| #include "caffe_blob.h" |
| |
| namespace mxnet { |
| namespace op { |
| |
| struct CaffeOpParam : public dmlc::Parameter<CaffeOpParam> { |
| ::caffe::LayerParameter prototxt; |
| int num_data, num_weight, num_out; |
| |
| DMLC_DECLARE_PARAMETER(CaffeOpParam) { DMLC_DECLARE_FIELD(prototxt).set_default("layer{}") |
| .describe("Caffe's layer parameter"); |
| DMLC_DECLARE_FIELD(num_data).set_default(1) |
| .describe("Operator input number"); |
| DMLC_DECLARE_FIELD(num_weight).set_default(0) |
| .describe("Weight number"); |
| DMLC_DECLARE_FIELD(num_out).set_default(1) |
| .describe("Operator output number"); |
| } |
| }; |
| |
| |
| /** |
| * \brief this is the implementation of caffe operator in caffe. |
| * \tparam xpu the device that the op will be executed on. |
| */ |
| template<typename xpu, typename Dtype> |
| class CaffeOp : public Operator { |
| public: |
| explicit CaffeOp(CaffeOpParam p):param_(p), |
| init_w_(false), |
| init_wd_(false), |
| setup_(false) { |
| std::string type = param_.prototxt.type(); |
| caffeOp_ = caffe::LayerRegistry<Dtype>::CreateLayer(param_.prototxt); |
| |
| caffe::InitCaffeBlobs<Dtype>(&bot_, param_.num_data); |
| caffe::InitCaffeBlobs<Dtype>(&top_, param_.num_out); |
| caffe::InitCaffeBlobs<Dtype>(&wei_, param_.num_weight); |
| flags_.resize(param_.num_data); |
| } |
| |
| ~CaffeOp() { |
| caffe::DelCaffeBlobs(&bot_, param_.num_data); |
| caffe::DelCaffeBlobs(&top_, param_.num_out); |
| caffe::DelCaffeBlobs(&wei_, param_.num_weight); |
| } |
| |
| virtual void Forward(const OpContext &ctx, |
| const std::vector<TBlob> &in_data, |
| const std::vector<OpReqType> &req, |
| const std::vector<TBlob> &out_data, |
| const std::vector<TBlob> &aux_args) { |
| // Set mode before forward |
| caffe::CaffeMode::SetMode<xpu>(); |
| using ::caffe::Blob; |
| using std::vector; |
| using namespace mshadow; |
| using namespace mshadow::expr; |
| for (uint32_t i = 0; i < req.size(); ++i) |
| CHECK_EQ(req[i], kWriteTo); |
| int expected_num_data = param_.num_weight + param_.num_data; |
| CHECK_EQ(in_data.size(), expected_num_data); |
| CHECK_EQ(out_data.size(), param_.num_out); |
| |
| #if defined(__CUDACC__) |
| Stream<xpu> *s = ctx.get_stream<xpu>(); |
| // TODO(Haoran): when need cublas handle in stream? |
| CHECK_EQ(s->blas_handle_ownership_, Stream<xpu>::OwnHandle) |
| << "Must init CuBLAS handle in stream"; |
| #endif // __CUDACC__ |
| |
| caffe::TBlob2CaffeBlob<xpu, Dtype>(caffe::Data, |
| bot_.begin(), |
| in_data.begin(), |
| param_.num_data); |
| caffe::TBlob2CaffeBlob<xpu, Dtype>(caffe::Data, |
| top_.begin(), |
| out_data.begin(), |
| param_.num_out); |
| CaffeOpSetup(); |
| // Init caffe's weight pointer |
| if (!init_w_) { |
| init_w_ = true; |
| caffe::TBlob2CaffeBlob<xpu, Dtype>(caffe::Data, |
| wei_.begin(), |
| in_data.begin() + param_.num_data, |
| param_.num_weight); |
| caffe::SetOpBlobs(caffeOp_, wei_); |
| } |
| if (ctx.is_train) |
| MXCAFFELAYER(caffeOp_, Dtype)->SetPhase(::caffe::TRAIN); |
| else |
| MXCAFFELAYER(caffeOp_, Dtype)->SetPhase(::caffe::TEST); |
| caffeOp_->Forward(bot_, top_); |
| |
| #if defined(__CUDACC__) |
| // Sync cpu data to gpu data |
| for (uint32_t i = 0; i < top_.size(); ++i) |
| top_[i]->gpu_data(); |
| |
| CHECK_EQ(cudaStreamSynchronize(NULL), cudaSuccess); |
| #endif // __CUDACC__ |
| } |
| |
| // Set up caffe op with real data |
| void CaffeOpSetup() { |
| if (!setup_) { |
| setup_ = true; |
| caffeOp_->SetUp(bot_, top_); |
| } |
| } |
| |
| virtual void Backward(const OpContext &ctx, |
| const std::vector<TBlob> &out_grad, |
| const std::vector<TBlob> &in_data, |
| const std::vector<TBlob> &out_data, |
| const std::vector<OpReqType> &req, |
| const std::vector<TBlob> &in_grad, |
| const std::vector<TBlob> &aux_args) { |
| // Set mode before backward |
| caffe::CaffeMode::SetMode<xpu>(); |
| using namespace mshadow; |
| using namespace mshadow::expr; |
| CHECK_EQ(out_grad.size(), param_.num_out); |
| for (int i = 0; i < param_.num_data; ++i) |
| CHECK(req[i] != kAddTo) << "caffe doesn't accm diff on bottom data"; |
| |
| int expected_num_data = param_.num_weight + param_.num_data; |
| CHECK(in_data.size() == expected_num_data && in_grad.size() == expected_num_data); |
| CHECK_EQ(req.size(), expected_num_data); |
| |
| Stream<xpu> *s = ctx.get_stream<xpu>(); |
| #if defined(__CUDACC__) |
| // TODO(Haoran): when need cublas handle in stream? |
| CHECK_EQ(s->blas_handle_ownership_, Stream<xpu>::OwnHandle) |
| << "Must init CuBLAS handle in stream"; |
| #endif // __CUDACC__ |
| |
| caffe::TBlob2CaffeBlob<xpu, Dtype>(caffe::Grad, |
| bot_.begin(), |
| in_grad.begin(), |
| param_.num_data); |
| caffe::TBlob2CaffeBlob<xpu, Dtype>(caffe::Grad, |
| top_.begin(), |
| out_grad.begin(), |
| param_.num_out); |
| |
| // Init caffe's gradient pointer |
| if (!init_wd_) { |
| init_wd_ = true; |
| caffe::TBlob2CaffeBlob<xpu, Dtype>(caffe::Grad, |
| wei_.begin(), |
| in_grad.begin() + param_.num_data, |
| param_.num_weight); |
| } |
| |
| // Handle OpReqType of weights |
| for (int i = param_.num_data; i < expected_num_data; ++i) |
| HandleOpReq(s, req[i], in_grad[i]); |
| |
| // Set BP flag |
| for (int i = 0; i < param_.num_data; ++i) |
| flags_[i] = req[i] != kNullOp; |
| |
| caffeOp_->Backward(top_, flags_, bot_); |
| |
| #if defined(__CUDACC__) |
| // Sync cpu diff to gpu diff |
| for (uint32_t i = 0; i < bot_.size(); ++i) |
| bot_[i]->gpu_diff(); |
| |
| CHECK_EQ(cudaStreamSynchronize(NULL), cudaSuccess); |
| #endif // __CUDACC__ |
| } |
| |
| void HandleOpReq(mshadow::Stream<xpu>*s, OpReqType req, const TBlob& in_g) { |
| if ((req == kWriteInplace) || (req == kWriteTo)) { |
| mshadow::Tensor<xpu, 2, Dtype> grad = in_g.FlatTo2D<xpu, Dtype>(s); |
| grad = 0; |
| } |
| } |
| |
| private: |
| CaffeOpParam param_; |
| ::caffe::Layer<Dtype> *caffeOp_; |
| std::vector< ::caffe::Blob<Dtype> *> bot_, top_, wei_; |
| std::vector<bool> flags_; |
| bool init_w_, init_wd_, setup_; |
| }; // class CaffeOp |
| |
| // Decalre Factory function, used for dispatch specialization |
| template<typename xpu> |
| Operator* CreateOp(CaffeOpParam param, int); |
| |
| #if DMLC_USE_CXX11 |
| class CaffeOpProp : public OperatorProperty { |
| public: |
| std::vector<std::string> ListArguments() const override { |
| std::vector<std::string> res; |
| for (int i = 0; i < param_.num_data; ++i) |
| res.push_back(std::string("data_") + std::to_string(i)); |
| |
| for (int i = 0; i < param_.num_weight; ++i) { |
| if (i == 0) |
| res.push_back(std::to_string(i) + "_weight"); |
| else |
| res.push_back(std::to_string(i) + "_bias"); |
| } |
| return res; |
| } |
| |
| std::vector<std::string> ListOutputs() const override { |
| if (param_.num_out > 1) { |
| std::vector<std::string> ret; |
| for (int i = 0; i < param_.num_out; ++i) |
| ret.push_back("output" + std::to_string(i)); |
| return ret; |
| } else { |
| return {"output"}; |
| } |
| } |
| |
| void Init(const std::vector<std::pair<std::string, std::string> >& kwargs) override { |
| param_.Init(kwargs); |
| } |
| |
| std::map<std::string, std::string> GetParams() const override { |
| return param_.__DICT__(); |
| } |
| |
| /* |
| * \brief Set up caffeOp_ to infer weights & output shape |
| * \brief Initialize param_'s in & out dims |
| */ |
| bool InferShape(std::vector<TShape> *in_shape, |
| std::vector<TShape> *out_shape, |
| std::vector<TShape> *aux_shape) const override { |
| if (caffeOp_ == NULL) |
| caffeOp_ = caffe::LayerRegistry<float>::CreateLayer(param_.prototxt); |
| using namespace mshadow; |
| using ::caffe::Blob; |
| using std::vector; |
| CHECK_GE(in_shape->size(), param_.num_data); |
| // Initialize emtryp bottom & top blobs for caffeop |
| vector<Blob<float> *> bot_blobs, top_blobs; |
| |
| for (int i = 0; i < param_.num_data; ++i) { |
| TShape tshape = (*in_shape)[i]; |
| if (tshape.ndim() == 0) return false; |
| auto blob_ptr = new Blob<float>(); |
| blob_ptr->Reshape(caffe::TShape2Vector(tshape)); |
| bot_blobs.push_back(blob_ptr); |
| } |
| |
| for (int i = 0; i < param_.num_out; ++i) |
| top_blobs.push_back(new Blob<float>()); |
| |
| caffeOp_->SetUp(bot_blobs, top_blobs); |
| CHECK_EQ(in_shape->size(), caffeOp_->blobs().size() + param_.num_data); |
| // Set weight shape |
| CHECK_EQ(param_.num_weight, caffeOp_->blobs().size()); |
| for (int i = 0; i < param_.num_weight ; ++i) { |
| TShape tshape = caffe::Vector2TShape(caffeOp_->blobs()[i]->shape()); |
| SHAPE_ASSIGN_CHECK(*in_shape, i + param_.num_data, tshape); |
| } |
| // Initialize out shapes |
| out_shape->clear(); |
| for (auto blob : top_blobs) { |
| TShape tshape = caffe::Vector2TShape(blob->shape()); |
| out_shape->push_back(tshape); |
| } |
| |
| for (auto blob_ptr : bot_blobs) |
| delete blob_ptr; |
| for (auto blob_ptr : top_blobs) |
| delete blob_ptr; |
| return true; |
| } |
| |
| OperatorProperty* Copy() const override { |
| auto copy_prop = new CaffeOpProp(); |
| copy_prop->param_ = this->param_; |
| return copy_prop; |
| } |
| |
| std::string TypeString() const override { |
| return "CaffeOp"; |
| } |
| |
| Operator* CreateOperator(Context ctx) const override { |
| LOG(FATAL) << "Not Implemented."; |
| return NULL; |
| } |
| |
| Operator* CreateOperatorEx(Context ctx, std::vector<TShape> *in_shape, |
| std::vector<int> *in_type) const override; |
| |
| private: |
| mutable CaffeOpParam param_; |
| mutable ::caffe::Layer<float> *caffeOp_; |
| }; // class CaffeOpSymbol |
| #endif |
| |
| } // namespace op |
| } // namespace mxnet |
| #endif // PLUGIN_CAFFE_CAFFE_OP_INL_H_ |