Caffe2 - C++ API
A deep learning, cross platform ML framework
print_registered_core_operators.cc
1 
17 #include <iostream>
18 #include <string>
19 
20 #include "caffe2/core/init.h"
21 #include "caffe2/core/operator.h"
22 #include "caffe2/core/operator_schema.h"
23 
24 C10_DEFINE_string(schema, "", "Print doc and schema of a particular operator");
25 
26 static bool HasSchema(const std::string& str) {
27  return caffe2::OpSchemaRegistry::Schema(str);
28 }
29 
30 static bool HasDoc(const std::string& str) {
31  const auto* schema = caffe2::OpSchemaRegistry::Schema(str);
32  return (schema != nullptr) && (schema->doc() != nullptr);
33 }
34 
35 int main(int argc, char** argv) {
36  caffe2::GlobalInit(&argc, &argv);
37 
38  if (!FLAGS_schema.empty()) {
39  const auto* schema = caffe2::OpSchemaRegistry::Schema(FLAGS_schema);
40  if (!schema) {
41  std::cerr << "Operator " << FLAGS_schema << " doesn't have a schema"
42  << std::endl;
43  return 1;
44  }
45  std::cout << "Operator " << FLAGS_schema << ": " << std::endl << *schema;
46  return 0;
47  }
48 
49  for (const auto& pair : *caffe2::gDeviceTypeRegistry()) {
50  std::cout << "Device type " << pair.first
51 #ifndef CAFFE2_USE_LITE_PROTO
52  << " ("
53  << at::DeviceTypeName(static_cast<caffe2::DeviceType>(pair.first))
54  << ")"
55 #endif
56  << std::endl;
57  for (const auto& key : pair.second->Keys()) {
58  std::cout << "\t(schema: " << HasSchema(key) << ", doc: " << HasDoc(key)
59  << ")\t" << key << std::endl;
60  }
61  }
62 
63  std::cout << "Operators that have gradients registered:" << std::endl;
64  for (const auto& key : caffe2::GradientRegistry()->Keys()) {
65  std::cout << "\t(schema: " << HasSchema(key) << ", doc: "
66  << HasDoc(key) << ")\t"
67  << key << std::endl;
68  }
69  return 0;
70 }
bool GlobalInit(int *pargc, char ***pargv)
Initialize the global environment of caffe2.
Definition: init.cc:44