Add python api for graph optimization level (#882)

This commit is contained in:
Ashwini Khade 2019-04-23 11:11:35 -07:00 committed by GitHub
parent b0a37477db
commit fb3b63438d
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -261,13 +261,13 @@ void addGlobalMethods(py::module& m) {
#ifdef onnxruntime_PYBIND_EXPORT_OPSCHEMA
m.def(
"get_all_operator_schema",
"get_all_operator_schema",
[]() -> const std::vector<ONNX_NAMESPACE::OpSchema> {
return ONNX_NAMESPACE::OpSchemaRegistry::get_all_schemas_with_history();
},
"Return a vector of OpSchema all registed operators"
);
#endif
#endif
}
@ -380,7 +380,18 @@ Set this option to false if you don't want it. Default is True.)pbdoc")
R"pbdoc(Applies to session load, initialization, etc. Default is 0.)pbdoc")
.def_readwrite("session_thread_pool_size", &SessionOptions::session_thread_pool_size,
R"pbdoc(How many threads in the session thread pool. Default is 0 to let onnxruntime choose.
This parameter is unused unless *enable_sequential_execution* is false.)pbdoc");
This parameter is unused unless *enable_sequential_execution* is false.)pbdoc")
.def_property_readonly("graph_optimization_level",
[](const SessionOptions* options) -> uint32_t {
return static_cast<uint32_t>(options->graph_optimization_level);
},
R"pbdoc(Graph optimization level for this session.)pbdoc")
.def("set_graph_optimization_level",
[](SessionOptions* options, uint32_t level) -> void {
options->graph_optimization_level = static_cast<TransformerLevel>(level);
},
R"pbdoc(Graph optimization level for this session. 0 disables all optimizations.
Whereas 1 enables basic optimizations and 2 enables all optimizations.)pbdoc");
py::class_<RunOptions>(m, "RunOptions", R"pbdoc(Configuration information for a single Run.)pbdoc")
.def(py::init())
@ -410,52 +421,50 @@ including arg name, arg type (contains both type and shape).)pbdoc")
return *(na.Type());
},
"node type")
.def(
"__str__", [](const onnxruntime::NodeArg& na) -> std::string {
std::ostringstream res;
res << "NodeArg(name='" << na.Name() << "', type='" << *(na.Type()) << "', shape=";
auto shape = na.Shape();
std::vector<py::object> arr;
if (shape == nullptr || shape->dim_size() == 0) {
res << "[]";
} else {
res << "[";
for (int i = 0; i < shape->dim_size(); ++i) {
if (shape->dim(i).has_dim_value()) {
res << shape->dim(i).dim_value();
} else if (shape->dim(i).has_dim_param()) {
res << "None";
}
if (i < shape->dim_size() - 1) {
res << ", ";
}
}
res << "]";
.def("__str__", [](const onnxruntime::NodeArg& na) -> std::string {
std::ostringstream res;
res << "NodeArg(name='" << na.Name() << "', type='" << *(na.Type()) << "', shape=";
auto shape = na.Shape();
std::vector<py::object> arr;
if (shape == nullptr || shape->dim_size() == 0) {
res << "[]";
} else {
res << "[";
for (int i = 0; i < shape->dim_size(); ++i) {
if (shape->dim(i).has_dim_value()) {
res << shape->dim(i).dim_value();
} else if (shape->dim(i).has_dim_param()) {
res << "None";
}
res << ")";
if (i < shape->dim_size() - 1) {
res << ", ";
}
}
res << "]";
}
res << ")";
return std::string(res.str());
},
"converts the node into a readable string")
.def_property_readonly(
"shape", [](const onnxruntime::NodeArg& na) -> std::vector<py::object> {
auto shape = na.Shape();
std::vector<py::object> arr;
if (shape == nullptr || shape->dim_size() == 0) {
return arr;
}
return std::string(res.str());
},
"converts the node into a readable string")
.def_property_readonly("shape", [](const onnxruntime::NodeArg& na) -> std::vector<py::object> {
auto shape = na.Shape();
std::vector<py::object> arr;
if (shape == nullptr || shape->dim_size() == 0) {
return arr;
}
arr.resize(shape->dim_size());
for (int i = 0; i < shape->dim_size(); ++i) {
if (shape->dim(i).has_dim_value()) {
arr[i] = py::cast(shape->dim(i).dim_value());
} else if (shape->dim(i).has_dim_param()) {
arr[i] = py::none();
}
}
return arr;
},
"node shape (assuming the node holds a tensor)");
arr.resize(shape->dim_size());
for (int i = 0; i < shape->dim_size(); ++i) {
if (shape->dim(i).has_dim_value()) {
arr[i] = py::cast(shape->dim(i).dim_value());
} else if (shape->dim(i).has_dim_param()) {
arr[i] = py::none();
}
}
return arr;
},
"node shape (assuming the node holds a tensor)");
py::class_<SessionObjectInitializer>(m, "SessionObjectInitializer");
py::class_<InferenceSession>(m, "InferenceSession", R"pbdoc(This is the main class used to run a model.)pbdoc")
@ -470,16 +479,15 @@ including arg name, arg type (contains both type and shape).)pbdoc")
InitializeSession(sess);
},
R"pbdoc(Load a model saved in ONNX format.)pbdoc")
.def(
"read_bytes", [](InferenceSession* sess, const py::bytes& serializedModel) {
std::istringstream buffer(serializedModel);
auto status = sess->Load(buffer);
if (!status.IsOK()) {
throw std::runtime_error(status.ToString().c_str());
}
InitializeSession(sess);
},
R"pbdoc(Load a model serialized in ONNX format.)pbdoc")
.def("read_bytes", [](InferenceSession* sess, const py::bytes& serializedModel) {
std::istringstream buffer(serializedModel);
auto status = sess->Load(buffer);
if (!status.IsOK()) {
throw std::runtime_error(status.ToString().c_str());
}
InitializeSession(sess);
},
R"pbdoc(Load a model serialized in ONNX format.)pbdoc")
.def("run", [](InferenceSession* sess, std::vector<std::string> output_names, std::map<std::string, py::object> pyfeeds, RunOptions* run_options = nullptr) -> std::vector<py::object> {
NameMLValMap feeds;
for (auto _ : pyfeeds) {
@ -585,7 +593,7 @@ PYBIND11_MODULE(onnxruntime_pybind11_state, m) {
#ifdef onnxruntime_PYBIND_EXPORT_OPSCHEMA
addOpSchemaSubmodule(m);
#endif
}
} // namespace python