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

Perform graph transformations during offline tooling (#12422)

Browse source
This commit is contained in:
Baiju Meswani 2022-08-03 11:27:12 -07:00 committed by GitHub
parent dc984a03d5
commit 7f58bd7236
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 61 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -26,6 +26,7 @@
#include "orttraining/core/graph/gradient_definition_registry.h"
#include "python/onnxruntime_pybind_mlvalue.h"
#include "orttraining/python/orttraining_pybind_common.h"
#include "orttraining/core/optimizer/graph_transformer_utils.h"
#ifdef ENABLE_TRAINING_TORCH_INTEROP
#include "orttraining/core/framework/torch/custom_function_register.h"
@ -859,6 +860,60 @@ void addObjectMethodsForTraining(py::module& m, ExecutionProviderRegistrationFn
return tensor_protos_pybytes;
});
m.def("get_optimized_model",
[](const py::bytes& serialized_model,
const std::unordered_set<std::string>& graph_entities_that_require_gradients) {
// Load the serialized model
std::istringstream buffer(serialized_model);
ONNX_NAMESPACE::ModelProto model_proto;
ORT_THROW_IF_ERROR(Model::Load(buffer, &model_proto));
// Get the ort model from ModelProto model
auto logger_ptr = std::make_unique<logging::Logger>(logging::LoggingManager::DefaultLogger());
logger_ptr->SetSeverity(logging::Severity::kINFO);
std::shared_ptr<onnxruntime::Model> ort_model;
ORT_THROW_IF_ERROR(Model::Load(model_proto, ort_model, nullptr, *logger_ptr));
Graph& graph = ort_model->MainGraph();
ORT_THROW_IF_ERROR(graph.Resolve());
// Register the pretraining graph transformations so that they are run twice
constexpr size_t NumSteps = 2;
GraphTransformerManager graph_transformation_mgr{NumSteps};
std::unique_ptr<CPUExecutionProvider> cpu_execution_provider =
std::make_unique<CPUExecutionProvider>(CPUExecutionProviderInfo());
const auto add_transformers = [&cpu_execution_provider,
&graph_transformation_mgr,
&graph_entities_that_require_gradients](TransformerLevel level) {
auto transformers_to_register = transformer_utils::GeneratePreTrainingTransformers(
level, graph_entities_that_require_gradients, TrainingGraphTransformerConfiguration(),
*cpu_execution_provider);
for (auto& entry : transformers_to_register) {
ORT_THROW_IF_ERROR(graph_transformation_mgr.Register(std::move(entry), level));
}
return Status::OK();
};
for (int i = static_cast<int>(TransformerLevel::Level1); i <= static_cast<int>(TransformerLevel::MaxLevel); i++) {
TransformerLevel level = static_cast<TransformerLevel>(i);
if (TransformerLevel::MaxLevel >= level) {
ORT_THROW_IF_ERROR(add_transformers(level));
}
}
// Run the graph transformations
for (int i = static_cast<int>(TransformerLevel::Level1); i <= static_cast<int>(TransformerLevel::MaxLevel); i++) {
ORT_THROW_IF_ERROR(
graph_transformation_mgr.ApplyTransformers(graph, static_cast<TransformerLevel>(i), *logger_ptr));
}
// Return the optimized model.
std::string model_str;
ort_model->ToProto().SerializeToString(&model_str);
return py::bytes(model_str);
});
#endif
}

7
orttraining/orttraining/python/training/onnxblock/_graph_utils.py
View file

@ -7,7 +7,7 @@ import random
import onnx
from onnxruntime.capi._pybind_state import GradientGraphBuilder
from onnxruntime.capi._pybind_state import GradientGraphBuilder, get_optimized_model
def get_output_from_output_name(onnx_model, output_name):
@ -85,6 +85,11 @@ def build_gradient_graph(accessor, user_args_requiring_grad, user_args_not_requi
for argument_name in user_args_requiring_grad:
all_args_requiring_gradient.append(argument_name)
# Run graph transformations to optimize the graph
model.CopyFrom(
onnx.load_from_string(get_optimized_model(model.SerializeToString(), set(all_args_requiring_gradient)))
)
# Assumption is that the first graph output is the loss output
if isinstance(output_names, str):
output_names = [output_names]