diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index c10413f676..ef207af880 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -56,8 +56,10 @@ #include "core/session/custom_ops.h" #endif +#ifdef ENABLE_TRAINING #include "core/providers/cpu/controlflow/event_pool.h" #include "core/providers/cpu/controlflow/message_queue.h" +#endif using namespace ONNX_NAMESPACE; using namespace onnxruntime::experimental; @@ -378,6 +380,7 @@ InferenceSession::~InferenceSession() { } } +#ifdef ENABLE_TRAINING // TODO: find a better way to terminate the background thread // backward is not completed yet, set terminate_flag to True if (task_.bg_thread_future_.valid()) { @@ -385,6 +388,7 @@ InferenceSession::~InferenceSession() { Status s = ContinueRunInBackground({}); ORT_UNUSED_PARAMETER(s); } +#endif #ifdef ONNXRUNTIME_ENABLE_INSTRUMENT if (session_activity_started_) @@ -1728,6 +1732,7 @@ common::Status InferenceSession::Run(IOBinding& io_binding) { return Run(run_options, io_binding); } +#ifdef ENABLE_TRAINING common::Status InferenceSession::RunInBackgroundAndWaitForYield(RunOptions& run_options, IOBinding& io_binding, std::vector& user_outputs) { const int64_t main_thread_event_id = 0; @@ -1781,6 +1786,7 @@ common::Status InferenceSession::ContinueRunInBackground(const std::vector void InferenceSession::StartProfiling(const std::basic_string& file_prefix) { diff --git a/onnxruntime/core/session/inference_session.h b/onnxruntime/core/session/inference_session.h index 5cf674803b..ee3e0fba38 100644 --- a/onnxruntime/core/session/inference_session.h +++ b/onnxruntime/core/session/inference_session.h @@ -5,8 +5,11 @@ #include #include + +#ifdef ENABLE_TRAINING #include #include +#endif #include "core/common/common.h" #include "core/common/logging/logging.h" @@ -299,12 +302,14 @@ class InferenceSession { virtual common::Status Run(const RunOptions& run_options, IOBinding& io_binding) ORT_MUST_USE_RESULT; common::Status Run(IOBinding& io_binding) ORT_MUST_USE_RESULT; +#ifdef ENABLE_TRAINING // For ORTModule.forward() virtual common::Status RunInBackgroundAndWaitForYield(RunOptions& run_options, IOBinding& io_binding, std::vector& user_outputs) ORT_MUST_USE_RESULT; // For ORTModule.backward() common::Status ContinueRunInBackground(const std::vector& backward_output_grads) ORT_MUST_USE_RESULT; +#endif /** * @return pair.first = OK; FAIL otherwise. pair.second is non-NULL when pair.first = OK. @@ -670,6 +675,7 @@ class InferenceSession { // those into new OrtValue instances, at which point we won't free them until the InferenceSession goes away. std::vector ort_format_model_bytes_; +#ifdef ENABLE_TRAINING // background thread for RunInBackgroundAndWaitForYield struct Task { std::thread bg_thread_; @@ -677,10 +683,11 @@ class InferenceSession { std::future bg_thread_future_; bool* terminate_flag_ = nullptr; } task_; - +#endif std::shared_ptr allocator_manager_; }; + struct SessionIOBinding { public: SessionIOBinding(InferenceSession* session); diff --git a/onnxruntime/python/dlpack_convertor.cc b/onnxruntime/python/dlpack_convertor.cc index be8bd53059..a44442f1bb 100644 --- a/onnxruntime/python/dlpack_convertor.cc +++ b/onnxruntime/python/dlpack_convertor.cc @@ -1,6 +1,6 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. - +#ifdef ENABLE_TRAINING #include "python/dlpack_convertor.h" namespace onnxruntime { @@ -122,3 +122,4 @@ DLManagedTensor* ort_value_to_dlpack(const OrtValue& ort_value) { } // namespace python } // namespace onnxruntime +#endif \ No newline at end of file diff --git a/onnxruntime/python/onnxruntime_pybind_state.cc b/onnxruntime/python/onnxruntime_pybind_state.cc index c959abf237..e7b998fdf5 100644 --- a/onnxruntime/python/onnxruntime_pybind_state.cc +++ b/onnxruntime/python/onnxruntime_pybind_state.cc @@ -25,7 +25,10 @@ #include "core/platform/env.h" #include "core/session/IOBinding.h" #include "core/session/abi_session_options_impl.h" + +#ifdef ENABLE_TRAINING #include "python/dlpack_convertor.h" +#endif // execution provider factory creator headers #include "core/providers/cpu/cpu_provider_factory_creator.h" @@ -1057,6 +1060,7 @@ void addOpSchemaSubmodule(py::module& m) { #endif //onnxruntime_PYBIND_EXPORT_OPSCHEMA +#ifdef ENABLE_TRAINING void dlpack_capsule_destructor(PyObject* data) { DLManagedTensor* dlmanged_tensor = (DLManagedTensor*)PyCapsule_GetPointer(data, "dltensor"); if (dlmanged_tensor) { @@ -1068,6 +1072,7 @@ void dlpack_capsule_destructor(PyObject* data) { PyErr_Clear(); } } +#endif void addObjectMethods(py::module& m, Environment& env) { py::enum_(m, "GraphOptimizationLevel") @@ -1221,6 +1226,7 @@ void addObjectMethods(py::module& m, Environment& env) { return ml_value; }) +#ifdef ENABLE_TRAINING .def_static("ortvalue_from_data_ptr", [](std::vector& shape, py::object& element_type, OrtDevice& device, int64_t data_ptr) { ORT_ENFORCE(data_ptr != 0, "Pointer to data memory is invalid"); @@ -1242,6 +1248,7 @@ void addObjectMethods(py::module& m, Environment& env) { return ort_value; }) +#endif .def("data_ptr", [](OrtValue* ml_value) -> int64_t { // TODO: Assumes that the OrtValue is a Tensor, make this generic to handle non-Tensors ORT_ENFORCE(ml_value->IsTensor(), "Only OrtValues that are Tensors are currently supported"); @@ -1304,11 +1311,14 @@ void addObjectMethods(py::module& m, Environment& env) { #endif return obj; }) +#ifdef ENABLE_TRAINING .def("to_dlpack", [](OrtValue* ort_value) -> py::object { DLManagedTensor* dlmanaged_tensor = ort_value_to_dlpack(*ort_value); return py::reinterpret_steal( PyCapsule_New(dlmanaged_tensor, "dltensor", dlpack_capsule_destructor)); - }); + }) +#endif +; py::class_ session_io_binding(m, "SessionIOBinding"); session_io_binding @@ -1831,6 +1841,7 @@ including arg name, arg type (contains both type and shape).)pbdoc") if (!status.IsOK()) throw std::runtime_error("Error in execution: " + status.ErrorMessage()); }) +#ifdef ENABLE_TRAINING .def("run_forward", [](PyInferenceSession* sess, SessionIOBinding& io_binding, RunOptions& run_options) -> std::vector { std::vector module_outputs; Status status = sess->GetSessionHandle()->RunInBackgroundAndWaitForYield(run_options, *io_binding.Get(), module_outputs); @@ -1844,7 +1855,9 @@ including arg name, arg type (contains both type and shape).)pbdoc") Status status = sess->GetSessionHandle()->ContinueRunInBackground(backward_output_grads); if (!status.IsOK()) throw std::runtime_error("Error in execution: " + status.ErrorMessage()); - }); + }) +#endif +; py::enum_(m, "ArenaExtendStrategy", py::arithmetic()) .value("kNextPowerOfTwo", onnxruntime::ArenaExtendStrategy::kNextPowerOfTwo)