#pragma once #include "caffe2/core/common.h" #include "caffe2/core/tensor.h" #include "caffe2/onnx/helper.h" #include "caffe2/proto/caffe2_pb.h" #include "onnx/onnx_pb.h" #include #include #include namespace caffe2 { namespace onnx { namespace { using ::ONNX_NAMESPACE::AttributeProto; using ::ONNX_NAMESPACE::GraphProto; using ::ONNX_NAMESPACE::ModelProto; using ::ONNX_NAMESPACE::NodeProto; using ::ONNX_NAMESPACE::TensorProto; } // namespace using ConvertedResult = std::pair, std::vector>; // Rewrite Caffe2 nets into SSA forms. Notice that we will preserve the external // output names for predict net. CAFFE2_API std::unordered_map SsaRewrite( caffe2::NetDef* init_net, caffe2::NetDef* pred_net); ::ONNX_NAMESPACE::TensorProto::DataType Caffe2TypeToOnnxType( caffe2::TensorProto::DataType t); class CAFFE2_API OnnxExporter { using SpecialOpConverter = ConvertedResult (OnnxExporter::*)( const caffe2::OperatorDef&, const std::unordered_map&); public: OnnxExporter(DummyName* dummy = nullptr) { if (dummy) { dummy_ = std::shared_ptr(dummy, [](DummyName*) {}); } else { dummy_ = std::make_shared(); } } ConvertedResult Caffe2OpToOnnxNodes( const caffe2::OperatorDef& def, const std::unordered_map& shapes); void InitOpToTensorProto(const caffe2::OperatorDef& def, TensorProto* tensor); private: ConvertedResult CommonCaffe2OpToOnnxNodes(const caffe2::OperatorDef& def); ConvertedResult CreateArgMaxMinOpNodes( const caffe2::OperatorDef& def, const std::unordered_map& shapes); ConvertedResult CreateBinaryElementwiseOpNodes( const caffe2::OperatorDef& def, const std::unordered_map& shapes); ConvertedResult CreateCastNodes( const caffe2::OperatorDef& def, const std::unordered_map& shapes); ConvertedResult CreateElementwiseLinearNodes( const caffe2::OperatorDef& def, const std::unordered_map& shapes); ConvertedResult CreateConvPoolNodes( const caffe2::OperatorDef& def, const std::unordered_map& shapes); ConvertedResult CreateGemmNodes( const caffe2::OperatorDef& def, const std::unordered_map& shapes); ConvertedResult CreateReshapeNodes( const caffe2::OperatorDef& def, const std::unordered_map& shapes); ConvertedResult CreateSliceNodes( const caffe2::OperatorDef& def, const std::unordered_map& shapes); ConvertedResult CreateChannelShuffleNodes( const caffe2::OperatorDef& def, const std::unordered_map& shapes); ConvertedResult CreateReduceMeanNodes( const caffe2::OperatorDef& def, const std::unordered_map& shapes); ConvertedResult CreateConcatNodes( const caffe2::OperatorDef& def, const std::unordered_map& shapes); ConvertedResult CreateMergeDimNodes( const caffe2::OperatorDef& def, const std::unordered_map& shapes); ConvertedResult CreateLrnNodes( const caffe2::OperatorDef& def, const std::unordered_map& shapes); ConvertedResult CreateUpsampleNodes( const caffe2::OperatorDef& def, const std::unordered_map& shapes); // \brief Check black listed arguemnts where we won't pass down when // converting to ONNX node bool IsBlackListed(const caffe2::Argument& arg); // \brief Convert Caffe2 argument to Onnx attribute void CopyCaffe2ArgToOnnxAttr( AttributeProto* attr, const std::string& op_type, const caffe2::Argument& arg); // LUT getters const std::unordered_map& get_renamed_operators() const; const std::unordered_map& get_renamed_attrs() const; const std:: unordered_map>& get_per_op_renamed_attrs() const; const std::unordered_map& get_special_operators() const; // Dummy name generator std::shared_ptr dummy_; }; } // namespace onnx } // namespace caffe2