18#include "caffe2/proto/caffe2.pb.h"
26#include <google/protobuf/io/zero_copy_stream_impl.h>
27#include <google/protobuf/io/coded_stream.h>
47 explicit Caffe2Importer(std::string predict_net, std::string init_net,
48 const std::vector<std::vector<int>> &input_shapes);
51 std::unique_ptr<mir::Graph> importModel();
56 std::string _predictNet;
58 std::unique_ptr<mir::Graph> _graph;
59 std::unique_ptr<caffe2::NetDef> _predict_net;
60 std::unique_ptr<caffe2::NetDef> _init_net;
61 std::unique_ptr<Caffe2OpCreator> _opCreator;
62 std::vector<mir::Shape> _inputShapes;
64 static const std::map<std::string, SupportedCaffe2OpType> _operatorTypes;
67 std::unordered_map<std::string, mir::Operation::Output *> _blobNameToOutput;
70 std::unique_ptr<mir::Graph> createIR();
76 void collectUnsupportedOps();
81 void createMIRNodesFromOp(const ::caffe2::OperatorDef &op);
86 std::vector<mir::Operation::Output *> getInputMIROps(const ::caffe2::OperatorDef &op);
88 void setOutputForTensor(
const std::string &tensor_name, Operation::Output *output);
94 void setGraphOutputs();
97using namespace ::caffe2;
100Caffe2Importer::Caffe2Importer(std::string predict_net, std::string init_net,
101 const std::vector<std::vector<int>> &input_shapes)
102 : _predictNet(
std::move(predict_net)), _initNet(
std::move(init_net))
104 for (
auto &shape : input_shapes)
105 _inputShapes.emplace_back(shape);
107 _graph = std::make_unique<mir::Graph>();
108 _opCreator = std::make_unique<Caffe2OpCreator>(_graph.get());
111Caffe2Importer::~Caffe2Importer() =
default;
113static void loadModelFile(
const std::string &filename, caffe2::NetDef *net)
115 GOOGLE_PROTOBUF_VERIFY_VERSION;
117 int file_handle = open(filename.c_str(), O_RDONLY);
119 if (file_handle == -1)
120 throw std::runtime_error(
"Couldn't open file \"" + filename +
"\": " + std::strerror(errno) +
123 google::protobuf::io::FileInputStream file_stream(file_handle);
124 file_stream.SetCloseOnDelete(
true);
126 google::protobuf::io::CodedInputStream coded_stream(&file_stream);
127 coded_stream.SetTotalBytesLimit(INT_MAX, INT_MAX);
129 if (!net->ParseFromCodedStream(&coded_stream))
130 throw std::runtime_error(
"Couldn't parse file \"" + filename +
"\".");
133 if (!coded_stream.ConsumedEntireMessage())
134 throw std::runtime_error(
"File \"" + filename +
"\" has not been consumed entirely.");
137void Caffe2Importer::import()
139 _predict_net = std::make_unique<NetDef>();
140 loadModelFile(_predictNet, _predict_net.get());
142 _init_net = std::make_unique<NetDef>();
143 loadModelFile(_initNet, _init_net.get());
145 collectUnsupportedOps();
148std::unique_ptr<mir::Graph> Caffe2Importer::createIR()
151 for (
const auto &op : _init_net->op())
152 createMIRNodesFromOp(op);
160 const auto &input_name = _predict_net->op(0).input(0);
163 setOutputForTensor(input_name, input);
165 for (
const auto &op : _predict_net->op())
166 createMIRNodesFromOp(op);
170 return std::move(_graph);
173std::unique_ptr<mir::Graph> Caffe2Importer::importModel()
179void Caffe2Importer::collectUnsupportedOps()
181 std::set<std::string> unsupportedOps;
182 for (
const auto &op : _predict_net->op())
184 if (_operatorTypes.find(op.type()) == _operatorTypes.end())
185 unsupportedOps.insert(op.type());
188 if (!unsupportedOps.empty())
190 std::string exceptionMsg(
"Can't load model, unsupported operators:");
191 for (
const auto &op : unsupportedOps)
192 exceptionMsg.append(
"\n * " + op);
193 throw std::runtime_error(exceptionMsg);
197void Caffe2Importer::createMIRNodesFromOp(
const OperatorDef &op)
199 std::vector<mir::Operation::Output *> outputs;
201 auto inputs = getInputMIROps(op);
206 case SupportedCaffe2OpType::constantFill:
207 case SupportedCaffe2OpType::givenTensorFill:
208 case SupportedCaffe2OpType::givenTensorInt64Fill:
209 outputs = _opCreator->convertConstant(inputs, op);
211 case SupportedCaffe2OpType::add:
212 outputs = _opCreator->convertAdd(inputs, op);
214 case SupportedCaffe2OpType::averagePool:
215 outputs = _opCreator->convertAveragePool(inputs, op);
217 case SupportedCaffe2OpType::conv:
218 outputs = _opCreator->convertConv(inputs, op);
220 case SupportedCaffe2OpType::concat:
221 outputs = _opCreator->convertConcat(inputs, op);
223 case SupportedCaffe2OpType::dropout:
224 outputs = _opCreator->convertDropout(inputs, op);
226 case SupportedCaffe2OpType::FC:
227 outputs = _opCreator->convertFC(inputs, op);
229 case SupportedCaffe2OpType::maxPool:
230 outputs = _opCreator->convertMaxPool(inputs, op);
232 case SupportedCaffe2OpType::mul:
233 outputs = _opCreator->convertMul(inputs, op);
235 case SupportedCaffe2OpType::relu:
236 outputs = _opCreator->convertRelu(inputs);
238 case SupportedCaffe2OpType::resizeNearest:
239 outputs = _opCreator->convertResizeNearest(inputs, op);
241 case SupportedCaffe2OpType::sigmoid:
242 outputs = _opCreator->convertSigmoid(inputs);
244 case SupportedCaffe2OpType::softmax:
245 outputs = _opCreator->convertSoftmax(inputs, op);
247 case SupportedCaffe2OpType::spatialBN:
248 outputs = _opCreator->convertSpatialBN(inputs, op);
250 case SupportedCaffe2OpType::sum:
251 outputs = _opCreator->convertSum(inputs);
253 case SupportedCaffe2OpType::clip:
254 outputs = _opCreator->convertClip(inputs, op);
256 case SupportedCaffe2OpType::reshape:
257 outputs = _opCreator->convertReshape(inputs, op);
260 assert(
false &&
"All unsupported types should have been found before this pass.");
263 for (
size_t i = 0; i < outputs.size(); ++i)
265 setOutputForTensor(op.output(i), outputs[i]);
269std::vector<mir::Operation::Output *> Caffe2Importer::getInputMIROps(
const OperatorDef &op)
271 std::vector<mir::Operation::Output *>
inputs;
273 for (
const auto &input_name : op.
input())
275 inputs.push_back(getOutputForTensor(input_name));
281void Caffe2Importer::setOutputForTensor(
const std::string &tensor_name, Operation::Output *output)
283 auto it = _blobNameToOutput.find(tensor_name);
284 if (it != _blobNameToOutput.cend())
288 it->second->setName(
"");
290 output->setName(tensor_name);
296 return _blobNameToOutput.at(name);
299void Caffe2Importer::setGraphOutputs()
305 const auto &output_name = _predict_net->op().rbegin()->output(0);
306 auto output = getOutputForTensor(output_name);
310const std::map<std::string, SupportedCaffe2OpType> Caffe2Importer::_operatorTypes = {
311 {
"Add", SupportedCaffe2OpType::add},
312 {
"AveragePool", SupportedCaffe2OpType::averagePool},
313 {
"Conv", SupportedCaffe2OpType::conv},
314 {
"Concat", SupportedCaffe2OpType::concat},
315 {
"ConstantFill", SupportedCaffe2OpType::constantFill},
316 {
"Dropout", SupportedCaffe2OpType::dropout},
317 {
"FC", SupportedCaffe2OpType::FC},
318 {
"GivenTensorFill", SupportedCaffe2OpType::givenTensorFill},
319 {
"MaxPool", SupportedCaffe2OpType::maxPool},
320 {
"Mul", SupportedCaffe2OpType::mul},
321 {
"Relu", SupportedCaffe2OpType::relu},
322 {
"ResizeNearest", SupportedCaffe2OpType::resizeNearest},
323 {
"Sigmoid", SupportedCaffe2OpType::sigmoid},
324 {
"Softmax", SupportedCaffe2OpType::softmax},
325 {
"SpatialBN", SupportedCaffe2OpType::spatialBN},
326 {
"Sum", SupportedCaffe2OpType::sum},
327 {
"Clip", SupportedCaffe2OpType::clip},
328 {
"Reshape", SupportedCaffe2OpType::reshape},
329 {
"GivenTensorInt64Fill", SupportedCaffe2OpType::givenTensorInt64Fill},
336std::unique_ptr<mir::Graph>
loadModel(std::string predict_net, std::string init_net,
337 const std::vector<std::vector<int>> &input_shapes)
339 Caffe2Importer importer(std::move(predict_net), std::move(init_net), input_shapes);
340 return importer.importModel();
Represents an output of a node.
const char * tensor_name(const circle::Tensor *tensor)
std::unique_ptr< mir::Graph > loadModel(std::string predict_net, std::string init_net, const std::vector< std::vector< int > > &input_shapes)