saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-04 04:07:22 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add utility to get the gradient graph from GradientGraphBuilder (#10995)

Browse source 
* Add pybind method to get the gradient graph

* Fix segmentation fault because of logging for gradien building
This commit is contained in:
Baiju Meswani 2022-03-25 17:13:56 -07:00 committed by GitHub
parent dc72159105
commit 9c6cc018a9
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 9 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
orttraining/orttraining/python/orttraining_pybind_state.cc
View file

@ -782,16 +782,16 @@ void addObjectMethodsForTraining(py::module& m, ExecutionProviderRegistrationFn
const std::unordered_set<std::string>& x_node_arg_names,
const std::string loss_node_arg_name) {
std::shared_ptr<Model> model;
auto logger = logging::LoggingManager::DefaultLogger();
auto logger_ptr = std::make_unique<logging::Logger>(logging::LoggingManager::DefaultLogger());
logger_ptr->SetSeverity(logging::Severity::kINFO);
ONNX_NAMESPACE::ModelProto model_proto;
std::istringstream model_istream(serialized_model);
ORT_THROW_IF_ERROR(Model::Load(model_istream, &model_proto));
ORT_THROW_IF_ERROR(Model::Load(model_proto, model, nullptr, logger));
ORT_THROW_IF_ERROR(Model::Load(model_proto, model, nullptr, *logger_ptr));
GradientGraphConfiguration gradient_graph_config{};
gradient_graph_config.set_gradients_as_graph_outputs = true;
// Save some objects, otherwise they get lost.
auto gradient_graph_config_ptr = std::make_unique<GradientGraphConfiguration>(gradient_graph_config);
auto logger_ptr = std::make_unique<logging::Logger>(logger);
auto gradient_graph_config_ptr = std::make_unique<GradientGraphConfiguration>(gradient_graph_config);
auto builder = std::make_unique<GradientGraphBuilder>(
&model->MainGraph(),
@ -808,6 +808,11 @@ void addObjectMethodsForTraining(py::module& m, ExecutionProviderRegistrationFn
})
.def("save", [](PyGradientGraphBuilder* gradient_graph_builder, const std::string& path) {
ORT_THROW_IF_ERROR(Model::Save(*(gradient_graph_builder->model), path));
})
.def("get_model", [](PyGradientGraphBuilder* gradient_graph_builder) {
std::string model_str;
gradient_graph_builder->model->ToProto().SerializeToString(&model_str);
return py::bytes(model_str);
});
py::class_<GradientNodeAttributeDefinition> gradient_node_attribute_definition(