Caffe2 - C++ API
A deep learning, cross platform ML framework
onnx_exporter.cc
1 #include "caffe2/core/logging.h"
2 #include "caffe2/onnx/onnx_exporter.h"
3 #include "caffe2/onnx/helper.h"
4 #include "caffe2/proto/caffe2_legacy.pb.h"
5 #include "caffe2/utils/map_utils.h"
6 
7 #include <unordered_set>
8 
9 namespace caffe2 {
10 namespace onnx {
11 
12 namespace {
13 // rewrite padding attributes
14 void ApplyTrans(
15  std::unordered_map<std::string, AttributeProto>* attrs,
16  bool global,
17  const std::string& k,
18  int dim = 2,
19  const std::string& ks = "") {
20  std::string ks2 = ks.empty() ? (k + "s") : ks;
21  std::string k_h, k_w, k_t, k_l, k_b, k_r;
22  if (dim == 2) {
23  k_h = k + "_h";
24  k_w = k + "_w";
25  } else {
26  k_t = k + "_t";
27  k_l = k + "_l";
28  k_b = k + "_b";
29  k_r = k + "_r";
30  }
31 
32  std::vector<int64_t> vals;
33  if (dim == 2 && attrs->count(k_h) && attrs->count(k_w)) {
34  auto it = attrs->find(k_h);
35  vals.push_back(it->second.i());
36  attrs->erase(it);
37  it = attrs->find(k_w);
38  vals.push_back(it->second.i());
39  attrs->erase(it);
40  } else if (
41  dim == 4 && attrs->count(k_t) && attrs->count(k_b) && attrs->count(k_l) &&
42  attrs->count(k_r)) {
43  auto it = attrs->find(k_t);
44  vals.push_back(it->second.i());
45  attrs->erase(it);
46  it = attrs->find(k_l);
47  vals.push_back(it->second.i());
48  attrs->erase(it);
49  it = attrs->find(k_b);
50  vals.push_back(it->second.i());
51  attrs->erase(it);
52  it = attrs->find(k_r);
53  vals.push_back(it->second.i());
54  attrs->erase(it);
55  } else if (attrs->count(k)) {
56  auto it = attrs->find(k);
57  auto tmp = it->second.i();
58  for (int i = 0; i < dim; ++i) {
59  vals.push_back(tmp);
60  }
61  attrs->erase(it);
62  }
63 
64  if (!vals.empty() && !global) {
65  attrs->emplace(ks2, MakeAttribute(ks2, vals));
66  }
67 }
68 
69 int64_t DimProd(const caffe2::TensorShape& shape, int start, int end) {
70  int64_t acc = 1;
71  for (int i = start; i < end; ++i) {
72  acc *= shape.dims(i);
73  }
74  return acc;
75 }
76 
77 TensorProto CreateOnnxShapeTensor(const std::vector<int64_t>& shape) {
78  TensorProto tensor;
79  tensor.set_name(DummyName::NewDummyName());
80  tensor.set_data_type(TensorProto::INT64);
81  tensor.add_dims(shape.size());
82  tensor.mutable_raw_data()->assign(
83  reinterpret_cast<const char*>(shape.data()), sizeof(int64_t) * shape.size());
84  return tensor;
85 }
86 } // namespace
87 
88 const std::unordered_map<std::string, std::string>&
89 OnnxExporter::get_renamed_operators() const {
90  const static std::unordered_map<std::string, std::string> kRenamedOperators{
91  {"SpatialBN", "BatchNormalization"},
92  {"Conv1D", "Conv"},
93  {"Conv2D", "Conv"},
94  {"Conv3D", "Conv"},
95  {"ConvTranspose1D", "ConvTranspose"},
96  {"ConvTranspose2D", "ConvTranspose"},
97  {"ConvTranspose3D", "ConvTranspose"},
98  {"MaxPool1D", "MaxPool"},
99  {"MaxPool2D", "MaxPool"},
100  {"MaxPool3D", "MaxPool"},
101  {"AveragePool1D", "AveragePool"},
102  {"AveragePool2D", "AveragePool"},
103  {"AveragePool3D", "AveragePool"}};
104  return kRenamedOperators;
105 }
106 
107 const std::unordered_map<std::string, std::string>&
108 OnnxExporter::get_renamed_attrs() const {
109  const static std::unordered_map<std::string, std::string> kRenamedAttrs{
110  {"kernels", "kernel_shape"}};
111  return kRenamedAttrs;
112 }
113 
114 const std::
115  unordered_map<std::string, std::unordered_map<std::string, std::string>>&
116  OnnxExporter::get_per_op_renamed_attrs() const {
117  const static std::
118  unordered_map<std::string, std::unordered_map<std::string, std::string>>
119  kPerOpRenamedAttrs = {{"Squeeze", {{"dims", "axes"}}},
120  {"Unsqueeze", {{"dims", "axes"}}},
121  {"Transpose", {{"axes", "perm"}}},
122  {"ConvTranspose", {{"adjs", "output_padding"}}},
123  {"Selu", {{"scale", "gamma"}}}};
124 
125  return kPerOpRenamedAttrs;
126 }
127 
128 const std::unordered_map<std::string, OnnxExporter::SpecialOpConverter>&
129 OnnxExporter::get_special_operators() const {
130  const static std::unordered_map<std::string, OnnxExporter::SpecialOpConverter>
131  kSpecialOperators = {
132  {"Conv", &OnnxExporter::CreateConvPoolNodes},
133  {"ConvTranspose", &OnnxExporter::CreateConvPoolNodes},
134  {"MaxPool", &OnnxExporter::CreateConvPoolNodes},
135  {"AveragePool", &OnnxExporter::CreateConvPoolNodes},
136  {"FC", &OnnxExporter::CreateGemmNodes},
137  {"Concat", &OnnxExporter::CreateConcatNodes},
138  {"LRN", &OnnxExporter::CreateLrnNodes},
139  {"Reshape", &OnnxExporter::CreateReshapeNodes},
140  {"Slice", &OnnxExporter::CreateSliceNodes},
141  {"ChannelShuffle", &OnnxExporter::CreateChannelShuffleNodes}
142  };
143  return kSpecialOperators;
144 }
145 
146 void OnnxExporter::CopyCaffe2ArgToOnnxAttr(
147  AttributeProto* attr,
148  const std::string& op_type,
149  const caffe2::Argument& arg) {
150  std::string name;
151  const auto& per_op_renamed_attr_lut = get_per_op_renamed_attrs();
152  const auto it = per_op_renamed_attr_lut.find(op_type);
153  if (it != per_op_renamed_attr_lut.end()) {
154  name = caffe2::get_default(it->second, arg.name(), arg.name());
155  } else {
156  name = caffe2::get_default(get_renamed_attrs(), arg.name(), arg.name());
157  }
158  attr->set_name(name);
159 
160  if (arg.has_f()) {
161  attr->set_f(arg.f());
162  attr->set_type(AttributeProto::FLOAT);
163  } else if (arg.has_i()) {
164  attr->set_i(arg.i());
165  attr->set_type(AttributeProto::INT);
166  } else if (arg.has_s()) {
167  attr->set_s(arg.s());
168  attr->set_type(AttributeProto::STRING);
169  } else if (arg.floats_size()) {
170  attr->mutable_floats()->CopyFrom(arg.floats());
171  attr->set_type(AttributeProto::STRINGS);
172  } else if (arg.ints_size()) {
173  attr->mutable_ints()->CopyFrom(arg.ints());
174  attr->set_type(AttributeProto::INTS);
175  } else if (arg.strings_size()) {
176  attr->mutable_strings()->CopyFrom(arg.strings());
177  attr->set_type(AttributeProto::STRINGS);
178  } else {
179  CAFFE_THROW(
180  caffe2::MakeString("Unsupported Caffe2 argument: ", arg.name()));
181  }
182 }
183 
184 bool OnnxExporter::IsBlackListed(const caffe2::Argument& arg) {
185  const static std::unordered_map<std::string, std::unordered_set<std::string>>
186  kBlackListString = {{"order", {"NCHW"}}};
187  const static std::unordered_map<std::string, std::unordered_set<int64_t>>
188  kBlackListInt = {{"cudnn_exhaustive_search", {0, 1}},
189  {"use_cudnn", {0, 1}}};
190 
191  if (arg.has_i()) {
192  const auto it = kBlackListInt.find(arg.name());
193  if (it != kBlackListInt.end()) {
194  return it->second.count(arg.i());
195  }
196  } else if (arg.has_s()) {
197  const auto it = kBlackListString.find(arg.name());
198  if (it != kBlackListString.end()) {
199  return it->second.count(arg.s());
200  }
201  }
202 
203  return false;
204 }
205 
206 ConvertedResult OnnxExporter::Caffe2OpToOnnxNodes(
207  const caffe2::OperatorDef& def,
208  const std::unordered_map<std::string, caffe2::TensorShape>& shapes) {
209  std::string type = def.type();
210  const auto& renamed_op_lut = get_renamed_operators();
211  const auto it = renamed_op_lut.find(type);
212  if (it != renamed_op_lut.end()) {
213  type = it->second;
214  }
215  const auto& special_op_lut = get_special_operators();
216  const auto it_op = get_special_operators().find(type);
217  if (it_op != special_op_lut.end()) {
218  return (this->*(it_op->second))(def, shapes);
219  } else {
220  return CommonCaffe2OpToOnnxNodes(def);
221  }
222 }
223 
224 ConvertedResult OnnxExporter::CommonCaffe2OpToOnnxNodes(
225  const caffe2::OperatorDef& def) {
226  ConvertedResult result;
227  auto& nodes = result.first;
228  nodes.emplace_back();
229  NodeProto& node = nodes.back();
230  node.set_name(def.name());
231  node.set_op_type(
232  caffe2::get_default(get_renamed_operators(), def.type(), def.type()));
233  for (const auto& i : def.input()) {
234  node.add_input(i);
235  }
236  for (const auto& o : def.output()) {
237  node.add_output(o);
238  }
239  for (const auto& a : def.arg()) {
240  if (!IsBlackListed(a)) {
241  auto* attr = node.add_attribute();
242  CopyCaffe2ArgToOnnxAttr(attr, def.type(), a);
243  }
244  }
245  return result;
246 }
247 
248 ConvertedResult OnnxExporter::CreateConvPoolNodes(
249  const caffe2::OperatorDef& def,
250  const std::unordered_map<std::string, caffe2::TensorShape>& shapes) {
251  auto result = CommonCaffe2OpToOnnxNodes(def);
252  auto& nodes = result.first;
253  auto& node = nodes.back();
254 
255  std::unordered_map<std::string, AttributeProto> attrs;
256  for (const auto& attr : node.attribute()) {
257  attrs.emplace(attr.name(), attr);
258  }
259 
260  // Handle global pooling
261  bool global = false;
262  if (node.op_type() == "MaxPool" || node.op_type() == "AveragePool") {
263  auto it = attrs.find("global_pooling");
264  if (it != attrs.end() && it->second.has_i() && it->second.i()) {
265  node.set_op_type("Global" + node.op_type());
266  global = true;
267  attrs.erase(it);
268  }
269  }
270 
271  ApplyTrans(&attrs, global, "kernel", 2, "kernel_shape");
272  ApplyTrans(&attrs, global, "stride");
273  ApplyTrans(&attrs, global, "dilation");
274  ApplyTrans(&attrs, global, "adj");
275  ApplyTrans(&attrs, global, "pad", 4);
276 
277  // Fix legacy pad attr
278  auto it = attrs.find("legacy_pad");
279  if (it != attrs.end()) {
280  auto legacy_pad_attr = it->second;
281  attrs.erase(it);
282  CAFFE_ENFORCE(
283  node.op_type().size() >= 4 &&
284  (node.op_type().rfind("Pool") == node.op_type().size() - 4));
285  CAFFE_ENFORCE(!global);
286  const auto& input_size = shapes.at(node.input(0));
287  const auto& output_size = shapes.at(node.output(0));
288  CAFFE_ENFORCE(output_size.dims().size() == 4);
289  if (legacy_pad_attr.i() ==
290  static_cast<int64_t>(caffe2::LegacyPadding::VALID)) {
291  CAFFE_ENFORCE(!attrs.count("pads"));
292  attrs.emplace("auto_pad", MakeAttribute("auto_pad", "VALID"));
293  } else if (
294  legacy_pad_attr.i() ==
295  static_cast<int64_t>(caffe2::LegacyPadding::SAME)) {
296  CAFFE_ENFORCE(!attrs.count("pads"));
297  // default behavior in Caffe2 is SAME_UPPER
298  // https://github.com/pytorch/pytorch/blob/master/caffe2/operators/conv_pool_op_base.h#L39
299  attrs.emplace("auto_pad", MakeAttribute("auto_pad", "SAME_UPPER"));
300  } else if (
301  legacy_pad_attr.i() ==
302  static_cast<int64_t>(caffe2::LegacyPadding::CAFFE_LEGACY_POOLING)) {
303  // The problem here is that, Pool op in Caffe may add an additional pixel,
304  // if the last part is smaller than stride. So we use the explicit padding
305  // to replace legacy_pad. pad[end] = output_size[start + 2] *
306  // stride[start] - pad[start] - 1 + kernel[start] - input[start + 2] end =
307  // start + len(pad) / 2
308  LOG(WARNING) << "Converting legacy padding to explicit padding.";
309  auto* pads_attr = attrs.at("pads").mutable_ints();
310  auto& strides_attr = attrs.at("strides").ints();
311  auto& kernel_shape_attr = attrs.at("kernel_shape").ints();
312  for (int i = 0; i < 2; ++i) {
313  int64_t tmp_pad = output_size.dims(i + 2) * strides_attr.Get(i) -
314  pads_attr->Get(i) - 1 + kernel_shape_attr.Get(i) -
315  input_size.dims(i + 2);
316  pads_attr->Set(i + 2, tmp_pad);
317  }
318  } else if (
319  legacy_pad_attr.i() !=
320  static_cast<int64_t>(caffe2::LegacyPadding::NOTSET)) {
321  CAFFE_THROW(caffe2::MakeString(
322  "Don't know how to handle the legacy_pad, while processing operator: ",
323  def.type()));
324  }
325  }
326 
327  node.clear_attribute();
328  for (const auto& kv : attrs) {
329  auto* attr = node.add_attribute();
330  attr->CopyFrom(kv.second);
331  }
332 
333  return result;
334 }
335 
336 ConvertedResult OnnxExporter::CreateLrnNodes(
337  const caffe2::OperatorDef& def,
338  const std::unordered_map<std::string, caffe2::TensorShape>& shapes) {
339  auto result = CommonCaffe2OpToOnnxNodes(def);
340  auto& nodes = result.first;
341 
342  CAFFE_ENFORCE_EQ(nodes.size(), 1);
343  auto& node = nodes.back();
344  if (node.output_size() == 2) {
345  node.mutable_output()->RemoveLast();
346  }
347 
348  return result;
349 }
350 
351 ConvertedResult OnnxExporter::CreateConcatNodes(
352  const caffe2::OperatorDef& def,
353  const std::unordered_map<std::string, caffe2::TensorShape>& shapes) {
354  auto result = CommonCaffe2OpToOnnxNodes(def);
355  auto& nodes = result.first;
356 
357  CAFFE_ENFORCE_EQ(nodes.size(), 1);
358  auto& node = nodes.back();
359  if (node.output_size() == 2) {
360  node.mutable_output()->RemoveLast();
361  }
362 
363  bool explicit_axis = false;
364  for (const auto& a: def.arg()) {
365  if (a.name() == "axis") {
366  explicit_axis = true;
367  break;
368  }
369  }
370  if (!explicit_axis) {
371  node.add_attribute()->CopyFrom(MakeAttribute("axis", 1L));
372  }
373 
374  return result;
375 }
376 
377 ConvertedResult OnnxExporter::CreateChannelShuffleNodes(
378  const caffe2::OperatorDef& def,
379  const std::unordered_map<std::string, caffe2::TensorShape>& shapes) {
380  const auto& x = def.input(0);
381  const auto& y = def.output(0);
382  const auto& x_shape = shapes.at(x);
383  CAFFE_ENFORCE_EQ(
384  x_shape.dims().size(),
385  4,
386  "Input shape of ChannelShuffle needs to be in NCHW format");
387  auto n = x_shape.dims(0);
388  auto c = x_shape.dims(1);
389  auto h = x_shape.dims(2);
390  auto w = x_shape.dims(3);
391  int64_t g = 0;
392  for (const auto& arg: def.arg()) {
393  if (arg.name() == "group") {
394  g = arg.i();
395  break;
396  }
397  }
398  CAFFE_ENFORCE(g && c % g == 0);
399  ConvertedResult result;
400  auto& nodes = result.first;
401  auto& const_tensors = result.second;
402 
403  const auto reshape_output = DummyName::NewDummyName();
404  std::vector<int64_t> dims = {n, g, c / g, h, w};
405  const_tensors.emplace_back(CreateOnnxShapeTensor(dims));
406  nodes.emplace_back(
407  MakeNode("Reshape", {x, const_tensors.back().name()}, {reshape_output}));
408 
409  const auto transpose_output = DummyName::NewDummyName();
410  dims = {0, 2, 1, 3, 4};
411  nodes.emplace_back(MakeNode(
412  "Transpose",
413  {reshape_output},
414  {transpose_output},
415  {MakeAttribute("perm", dims)}));
416 
417  dims = {n, c, h, w};
418  const_tensors.emplace_back(CreateOnnxShapeTensor(dims));
419  nodes.emplace_back(MakeNode(
420  "Reshape", {transpose_output, const_tensors.back().name()}, {y}));
421 
422  return result;
423 }
424 
425 ConvertedResult OnnxExporter::CreateSliceNodes(
426  const caffe2::OperatorDef& def,
427  const std::unordered_map<std::string, caffe2::TensorShape>& shapes) {
428  CAFFE_ENFORCE_EQ(
429  def.input_size(),
430  1,
431  "ONNX Slice operator does not support dynamic slice.");
432  auto result = CommonCaffe2OpToOnnxNodes(def);
433  auto& nodes = result.first;
434  CAFFE_ENFORCE_EQ(nodes.size(), 1);
435  auto& node = nodes.back();
436  const auto& shape = shapes.at(node.input(0));
437 
438  std::vector<int64_t> dims;
439  for (auto& attr: *node.mutable_attribute()) {
440  if (attr.name() == "starts") {
441  auto len = attr.ints_size();
442  if (len) {
443  dims.resize(len);
444  std::iota(dims.begin(), dims.end(), 0);
445  }
446  } else if (attr.name() == "ends") {
447  for (int i = 0; i < attr.ints_size(); ++i) {
448  auto end = attr.ints(i);
449  if (end >=0) {
450  continue;
451  }
452  if (end == -1) {
453  end = shape.dims(i);
454  } else {
455  ++end;
456  }
457  attr.set_ints(i, end);
458  }
459  }
460  }
461  if (!dims.empty()) {
462  node.add_attribute()->CopyFrom(MakeAttribute("axes", dims));
463  }
464 
465  return result;
466 }
467 
468 ConvertedResult OnnxExporter::CreateReshapeNodes(
469  const caffe2::OperatorDef& def,
470  const std::unordered_map<std::string, caffe2::TensorShape>& shapes) {
471  auto result = CommonCaffe2OpToOnnxNodes(def);
472  auto& nodes = result.first;
473  auto& const_tensors = result.second;
474  CAFFE_ENFORCE_EQ(nodes.size(), 1);
475  auto& node = nodes.back();
476 
477  int i = 0;
478  int attr_size = node.attribute_size();
479  for (; i < attr_size; ++i) {
480  const auto& attr = node.attribute(i);
481  if (attr.name() == "shape") {
482  std::vector<int64_t> shape;
483  for (const auto k: attr.ints()) {
484  shape.push_back(k);
485  }
486  const_tensors.emplace_back(CreateOnnxShapeTensor(shape));
487  node.add_input(const_tensors.back().name());
488  break;
489  }
490  }
491  if (i != attr_size) {
492  if (i != attr_size - 1) {
493  node.mutable_attribute()->SwapElements(i, attr_size - 1);
494  }
495  node.mutable_attribute()->RemoveLast();
496  }
497 
498  if (node.output_size() == 2) {
499  node.mutable_output()->RemoveLast();
500  }
501 
502  return result;
503 }
504 
505 ConvertedResult OnnxExporter::CreateGemmNodes(
506  const caffe2::OperatorDef& def,
507  const std::unordered_map<std::string, caffe2::TensorShape>& shapes) {
508  CAFFE_ENFORCE_EQ(def.input_size(), 3);
509  CAFFE_ENFORCE_GE(def.output_size(), 1);
510  auto x = def.input(0);
511  auto w = def.input(1);
512  const auto& b = def.input(2);
513  const auto& y = def.output(0);
514  const auto& x_shape = shapes.at(x);
515 
516  ConvertedResult result;
517  auto& nodes = result.first;
518  auto& const_tensors = result.second;
519  std::unordered_map<std::string, const caffe2::Argument*> args;
520  for (const auto& a : def.arg()) {
521  args.emplace(a.name(), &a);
522  }
523 
524  auto it = args.find("axis");
525  bool has_axis = (it != args.end());
526  int64_t axis = 0;
527  if (has_axis) {
528  axis = it->second->i();
529  auto outer = DimProd(x_shape, 0, axis);
530  auto inner = DimProd(x_shape, axis, x_shape.dims().size());
531  std::vector<int64_t> dims = {outer, inner};
532  auto reshaped_x = DummyName::NewDummyName();
533  const_tensors.emplace_back(CreateOnnxShapeTensor(dims));
534  nodes.emplace_back(
535  MakeNode("Reshape", {x, const_tensors.back().name()}, {reshaped_x}));
536  x = reshaped_x;
537  }
538 
539  it = args.find("axis_w");
540  if (it != args.end()) {
541  auto axis_w = it->second->i();
542  const auto& w_shape = shapes.at(w);
543  auto outer = DimProd(w_shape, 0, axis_w);
544  auto inner = DimProd(w_shape, axis_w, w_shape.dims().size());
545  std::vector<int64_t> dims = {outer, inner};
546  auto reshaped_w = DummyName::NewDummyName();
547  const_tensors.emplace_back(CreateOnnxShapeTensor(dims));
548  nodes.emplace_back(
549  MakeNode("Reshape", {w, const_tensors.back().name()}, {reshaped_w}));
550  w = reshaped_w;
551  }
552 
553  auto gemm_y_output = (has_axis) ? DummyName::NewDummyName() : y;
554  nodes.emplace_back(MakeNode(
555  "Gemm",
556  {x, w, b},
557  {gemm_y_output},
558  {MakeAttribute("transB", 1L), MakeAttribute("broadcast", 1)},
559  def.name()));
560 
561  if (has_axis) {
562  std::vector<int64_t> dims;
563  for (int i = 0; i < axis; ++i) {
564  dims.push_back(x_shape.dims(i));
565  }
566  dims.push_back(-1);
567  const_tensors.emplace_back(CreateOnnxShapeTensor(dims));
568  nodes.emplace_back(
569  MakeNode("Reshape", {gemm_y_output, const_tensors.back().name()}, {y}));
570  }
571 
572  return result;
573 }
574 } // namespace onnx
575 } // namespace caffe2
576 
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...