diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index e5793057a2..0007f735a0 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -160,6 +160,9 @@ option(onnxruntime_ENABLE_BITCODE "Enable bitcode for iOS only" OFF) # build eager mode option(onnxruntime_ENABLE_EAGER_MODE "build ort eager mode") +# build on device training mode +option(onnxruntime_ENABLE_ON_DEVICE_TRAINING "build ort on device training") + # build separate library of schemas of (custom) ops used by ORT (for ONNX to MLIR translation) option(onnxruntime_BUILD_OPSCHEMA_LIB "Build op schema library" ON) @@ -1935,6 +1938,15 @@ if (onnxruntime_ENABLE_EAGER_MODE) add_compile_definitions(ENABLE_EAGER_MODE) list(APPEND ONNXRUNTIME_TARGETS onnxruntime_eager) endif() +if (onnxruntime_ENABLE_ON_DEVICE_TRAINING) + if (NOT onnxruntime_ENABLE_TRAINING) + message( + FATAL_ERROR + "Option onnxruntime_ENABLE_ON_DEVICE_TRAINING can only be used when onnxruntime_ENABLE_TRAINING is enabled") + endif() + add_compile_definitions(ENABLE_ON_DEVICE_TRAINING) + list(APPEND ONNXRUNTIME_TARGETS onnxruntime_on_device_training) +endif() foreach(target_name ${ONNXRUNTIME_TARGETS}) include(${target_name}.cmake) endforeach() diff --git a/cmake/onnxruntime_on_device_training.cmake b/cmake/onnxruntime_on_device_training.cmake new file mode 100644 index 0000000000..9aed1d6c24 --- /dev/null +++ b/cmake/onnxruntime_on_device_training.cmake @@ -0,0 +1,52 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +set(_onnxflow_pb_cpp_srcs + "${ORTTRAINING_ROOT}/orttraining/onnxflow/csrc/onnxflow.pb.cc" + "${ORTTRAINING_ROOT}/orttraining/onnxflow/csrc/onnxflow.pb.h" + ) + +if(EXISTS "${ONNX_CUSTOM_PROTOC_EXECUTABLE}") + set(PROTOC_EXECUTABLE ${ONNX_CUSTOM_PROTOC_EXECUTABLE}) +else() + set(PROTOC_EXECUTABLE $) + set(PROTOC_DEPS protobuf::protoc) +endif() + +add_custom_command( + OUTPUT ${_onnxflow_pb_cpp_srcs} + COMMAND ${PROTOC_EXECUTABLE} + ARGS --python_out=${ORTTRAINING_ROOT}/orttraining/onnxflow/onnxflow/ --cpp_out=${ORTTRAINING_ROOT}/orttraining/onnxflow/csrc/ --proto_path=${ORTTRAINING_ROOT}/orttraining/onnxflow --proto_path=${REPO_ROOT}/cmake/external/protobuf/src ${ORTTRAINING_ROOT}/orttraining/onnxflow/onnxflow.proto + DEPENDS ${ORTTRAINING_ROOT}/orttraining/onnxflow/onnxflow.proto ${PROTOC_DEPS} + COMMENT "Running cpp protocol buffer compiler on onnxflow.proto" + VERBATIM ) + +file(GLOB onnxruntime_on_device_training_srcs CONFIGURE_DEPENDS + "${ORTTRAINING_ROOT}/orttraining/onnxflow/csrc/*.h" + "${ORTTRAINING_ROOT}/orttraining/onnxflow/csrc/*.cpp" + ) +list(APPEND onnxruntime_on_device_training_srcs ${_onnxflow_pb_cpp_srcs}) + +source_group(TREE ${REPO_ROOT} FILES ${onnxruntime_on_device_training_srcs}) + +onnxruntime_add_static_library(onnxruntime_on_device_training ${onnxruntime_on_device_training_srcs}) + +onnxruntime_add_include_to_target(onnxruntime_on_device_training onnxruntime_common onnxruntime_framework onnxruntime_optimizer onnxruntime_graph onnx onnx_proto ${PROTOBUF_LIB} flatbuffers) +target_include_directories(onnxruntime_on_device_training PRIVATE ${ONNXRUNTIME_ROOT} ${eigen_INCLUDE_DIRS}) +add_dependencies(onnxruntime_on_device_training ${onnxruntime_EXTERNAL_DEPENDENCIES}) +set_target_properties(onnxruntime_on_device_training PROPERTIES FOLDER "ONNXRuntime") +if (onnxruntime_ENABLE_TRAINING) + target_include_directories(onnxruntime_session PRIVATE ${ORTTRAINING_ROOT}) +endif() + +# sample loading of file +file(GLOB orttraining_on_device_sample_src CONFIGURE_DEPENDS + "${ORTTRAINING_ROOT}/orttraining/onnxflow/sample.m.cpp" + ) +onnxruntime_add_executable(orttraining_on_device_sample ${orttraining_on_device_sample_src}) +onnxruntime_add_include_to_target(orttraining_on_device_sample onnxruntime_on_device_training onnxruntime_common onnx onnx_proto ${PROTOBUF_LIB} onnxruntime_training flatbuffers) +target_include_directories(orttraining_on_device_sample PUBLIC ${CMAKE_CURRENT_BINARY_DIR} ${ONNXRUNTIME_ROOT} ${ORTTRAINING_ROOT} ${eigen_INCLUDE_DIRS} ${CXXOPTS} ${extra_includes} ${onnxruntime_graph_header} ${onnxruntime_exec_src_dir} ${CMAKE_CURRENT_BINARY_DIR} ${CMAKE_CURRENT_BINARY_DIR}/onnx onnxruntime_training_runner ${PROTOBUF_LIB}) + +target_link_libraries(orttraining_on_device_sample PRIVATE onnxruntime_on_device_training onnx onnx_proto onnxruntime_training ${ONNXRUNTIME_LIBS} ${onnxruntime_EXTERNAL_LIBRARIES} libprotobuf) +# set_target_properties(onnxruntime_training_mnist PROPERTIES FOLDER "ONNXRuntimeTest") + diff --git a/orttraining/orttraining/onnxflow/README.md b/orttraining/orttraining/onnxflow/README.md new file mode 100644 index 0000000000..04f96a7b18 --- /dev/null +++ b/orttraining/orttraining/onnxflow/README.md @@ -0,0 +1,28 @@ +# onnxflow + + +1. Build onnxruntime with on device training flag: +```sh +./build.sh --config RelWithDebInfo --enable_training --use_cuda --cuda_home /usr/local/cuda/ --cudnn_home /usr/local/cuda/ --build_wheel --parallel --cuda_version=11.3 --skip_tests --build_wheel --build_on_device_training +``` + +This will generate the protobuf files needed for serialization and deserialization of the parameters: +- orttraining/orttraining/onnxflow/csrc/onnxflow.pb.h +- orttraining/orttraining/onnxflow/csrc/onnxflow.pb.cc +- orttraining/orttraining/onnxflow/onnxflow/onnxflow_pb2.py + +2. Compose the model with the necessary loss and optimizer by running this from `orttraining/orttraining/onnxflow`: +```py +python sample.py +``` + +This will create the following: +- Forward+Loss+Backward training onnx graph +- Optimizer onnx graph +- Serialized parameters (saved as `parameters.of`) + +3. Use the saved onnx files and the parameters to perform training. To load the serialized parameters, see example utility `orttraining/orttraining/onnxflow/sample.m.cpp`. And run it by executing +```sh +orttraining_on_device_sample +``` +Pass in the absolute path of the `parameters.of` file when prompted. diff --git a/orttraining/orttraining/onnxflow/csrc/load_parameters.cpp b/orttraining/orttraining/onnxflow/csrc/load_parameters.cpp new file mode 100644 index 0000000000..23ebeffde0 --- /dev/null +++ b/orttraining/orttraining/onnxflow/csrc/load_parameters.cpp @@ -0,0 +1,26 @@ + +#include "load_parameters.h" +#include +#include +#include +#include +#include +#include + + +namespace onnxflow { + +OnnxFlowParameters load_parameters(const std::string& path_to_file) +{ + GOOGLE_PROTOBUF_VERIFY_VERSION; + + OnnxFlowParameters params; + std::ifstream t(path_to_file); + std::stringstream buffer; + buffer << t.rdbuf(); + params.ParseFromString(buffer.str()); + + return params; +} + +} // end namespace onnxflow diff --git a/orttraining/orttraining/onnxflow/csrc/load_parameters.h b/orttraining/orttraining/onnxflow/csrc/load_parameters.h new file mode 100644 index 0000000000..849677806f --- /dev/null +++ b/orttraining/orttraining/onnxflow/csrc/load_parameters.h @@ -0,0 +1,9 @@ + + +#include "onnxflow.pb.h" + +namespace onnxflow { + +OnnxFlowParameters load_parameters(const std::string& path_to_file); + +} // end namespace onnxflow diff --git a/orttraining/orttraining/onnxflow/csrc/onnxflow.pb.cc b/orttraining/orttraining/onnxflow/csrc/onnxflow.pb.cc new file mode 100644 index 0000000000..ff989845a1 --- /dev/null +++ b/orttraining/orttraining/onnxflow/csrc/onnxflow.pb.cc @@ -0,0 +1,543 @@ +// Generated by the protocol buffer compiler. DO NOT EDIT! +// source: onnxflow.proto + +#include "onnxflow.pb.h" + +#include + +#include +#include +#include +#include +// @@protoc_insertion_point(includes) +#include + +PROTOBUF_PRAGMA_INIT_SEG +namespace onnxflow { +constexpr OnnxFlowParameter::OnnxFlowParameter( + ::PROTOBUF_NAMESPACE_ID::internal::ConstantInitialized) + : data_(nullptr) + , requires_grad_(false) + , is_parameter_(false){} +struct OnnxFlowParameterDefaultTypeInternal { + constexpr OnnxFlowParameterDefaultTypeInternal() + : _instance(::PROTOBUF_NAMESPACE_ID::internal::ConstantInitialized{}) {} + ~OnnxFlowParameterDefaultTypeInternal() {} + union { + OnnxFlowParameter _instance; + }; +}; +PROTOBUF_ATTRIBUTE_NO_DESTROY PROTOBUF_CONSTINIT OnnxFlowParameterDefaultTypeInternal _OnnxFlowParameter_default_instance_; +constexpr OnnxFlowParameters::OnnxFlowParameters( + ::PROTOBUF_NAMESPACE_ID::internal::ConstantInitialized) + : parameters_(){} +struct OnnxFlowParametersDefaultTypeInternal { + constexpr OnnxFlowParametersDefaultTypeInternal() + : _instance(::PROTOBUF_NAMESPACE_ID::internal::ConstantInitialized{}) {} + ~OnnxFlowParametersDefaultTypeInternal() {} + union { + OnnxFlowParameters _instance; + }; +}; +PROTOBUF_ATTRIBUTE_NO_DESTROY PROTOBUF_CONSTINIT OnnxFlowParametersDefaultTypeInternal _OnnxFlowParameters_default_instance_; +} // namespace onnxflow +namespace onnxflow { + +// =================================================================== + +class OnnxFlowParameter::_Internal { + public: + using HasBits = decltype(std::declval()._has_bits_); + static const ::PROTOBUF_NAMESPACE_ID::Any& data(const OnnxFlowParameter* msg); + static void set_has_data(HasBits* has_bits) { + (*has_bits)[0] |= 1u; + } + static void set_has_requires_grad(HasBits* has_bits) { + (*has_bits)[0] |= 2u; + } + static void set_has_is_parameter(HasBits* has_bits) { + (*has_bits)[0] |= 4u; + } + static bool MissingRequiredFields(const HasBits& has_bits) { + return ((has_bits[0] & 0x00000007) ^ 0x00000007) != 0; + } +}; + +const ::PROTOBUF_NAMESPACE_ID::Any& +OnnxFlowParameter::_Internal::data(const OnnxFlowParameter* msg) { + return *msg->data_; +} +void OnnxFlowParameter::clear_data() { + if (data_ != nullptr) data_->Clear(); + _has_bits_[0] &= ~0x00000001u; +} +OnnxFlowParameter::OnnxFlowParameter(::PROTOBUF_NAMESPACE_ID::Arena* arena, + bool is_message_owned) + : ::PROTOBUF_NAMESPACE_ID::MessageLite(arena, is_message_owned) { + SharedCtor(); + if (!is_message_owned) { + RegisterArenaDtor(arena); + } + // @@protoc_insertion_point(arena_constructor:onnxflow.OnnxFlowParameter) +} +OnnxFlowParameter::OnnxFlowParameter(const OnnxFlowParameter& from) + : ::PROTOBUF_NAMESPACE_ID::MessageLite(), + _has_bits_(from._has_bits_) { + _internal_metadata_.MergeFrom(from._internal_metadata_); + if (from._internal_has_data()) { + data_ = new ::PROTOBUF_NAMESPACE_ID::Any(*from.data_); + } else { + data_ = nullptr; + } + ::memcpy(&requires_grad_, &from.requires_grad_, + static_cast(reinterpret_cast(&is_parameter_) - + reinterpret_cast(&requires_grad_)) + sizeof(is_parameter_)); + // @@protoc_insertion_point(copy_constructor:onnxflow.OnnxFlowParameter) +} + +void OnnxFlowParameter::SharedCtor() { +::memset(reinterpret_cast(this) + static_cast( + reinterpret_cast(&data_) - reinterpret_cast(this)), + 0, static_cast(reinterpret_cast(&is_parameter_) - + reinterpret_cast(&data_)) + sizeof(is_parameter_)); +} + +OnnxFlowParameter::~OnnxFlowParameter() { + // @@protoc_insertion_point(destructor:onnxflow.OnnxFlowParameter) + if (GetArenaForAllocation() != nullptr) return; + SharedDtor(); + _internal_metadata_.Delete(); +} + +inline void OnnxFlowParameter::SharedDtor() { + GOOGLE_DCHECK(GetArenaForAllocation() == nullptr); + if (this != internal_default_instance()) delete data_; +} + +void OnnxFlowParameter::ArenaDtor(void* object) { + OnnxFlowParameter* _this = reinterpret_cast< OnnxFlowParameter* >(object); + (void)_this; +} +void OnnxFlowParameter::RegisterArenaDtor(::PROTOBUF_NAMESPACE_ID::Arena*) { +} +void OnnxFlowParameter::SetCachedSize(int size) const { + _cached_size_.Set(size); +} + +void OnnxFlowParameter::Clear() { +// @@protoc_insertion_point(message_clear_start:onnxflow.OnnxFlowParameter) + ::PROTOBUF_NAMESPACE_ID::uint32 cached_has_bits = 0; + // Prevent compiler warnings about cached_has_bits being unused + (void) cached_has_bits; + + cached_has_bits = _has_bits_[0]; + if (cached_has_bits & 0x00000001u) { + GOOGLE_DCHECK(data_ != nullptr); + data_->Clear(); + } + ::memset(&requires_grad_, 0, static_cast( + reinterpret_cast(&is_parameter_) - + reinterpret_cast(&requires_grad_)) + sizeof(is_parameter_)); + _has_bits_.Clear(); + _internal_metadata_.Clear(); +} + +const char* OnnxFlowParameter::_InternalParse(const char* ptr, ::PROTOBUF_NAMESPACE_ID::internal::ParseContext* ctx) { +#define CHK_(x) if (PROTOBUF_PREDICT_FALSE(!(x))) goto failure + _Internal::HasBits has_bits{}; + while (!ctx->Done(&ptr)) { + ::PROTOBUF_NAMESPACE_ID::uint32 tag; + ptr = ::PROTOBUF_NAMESPACE_ID::internal::ReadTag(ptr, &tag); + switch (tag >> 3) { + // required .google.protobuf.Any data = 1; + case 1: + if (PROTOBUF_PREDICT_TRUE(static_cast<::PROTOBUF_NAMESPACE_ID::uint8>(tag) == 10)) { + ptr = ctx->ParseMessage(_internal_mutable_data(), ptr); + CHK_(ptr); + } else + goto handle_unusual; + continue; + // required bool requires_grad = 2; + case 2: + if (PROTOBUF_PREDICT_TRUE(static_cast<::PROTOBUF_NAMESPACE_ID::uint8>(tag) == 16)) { + _Internal::set_has_requires_grad(&has_bits); + requires_grad_ = ::PROTOBUF_NAMESPACE_ID::internal::ReadVarint64(&ptr); + CHK_(ptr); + } else + goto handle_unusual; + continue; + // required bool is_parameter = 3; + case 3: + if (PROTOBUF_PREDICT_TRUE(static_cast<::PROTOBUF_NAMESPACE_ID::uint8>(tag) == 24)) { + _Internal::set_has_is_parameter(&has_bits); + is_parameter_ = ::PROTOBUF_NAMESPACE_ID::internal::ReadVarint64(&ptr); + CHK_(ptr); + } else + goto handle_unusual; + continue; + default: + goto handle_unusual; + } // switch + handle_unusual: + if ((tag == 0) || ((tag & 7) == 4)) { + CHK_(ptr); + ctx->SetLastTag(tag); + goto message_done; + } + ptr = UnknownFieldParse( + tag, + _internal_metadata_.mutable_unknown_fields(), + ptr, ctx); + CHK_(ptr != nullptr); + } // while +message_done: + _has_bits_.Or(has_bits); + return ptr; +failure: + ptr = nullptr; + goto message_done; +#undef CHK_ +} + +::PROTOBUF_NAMESPACE_ID::uint8* OnnxFlowParameter::_InternalSerialize( + ::PROTOBUF_NAMESPACE_ID::uint8* target, ::PROTOBUF_NAMESPACE_ID::io::EpsCopyOutputStream* stream) const { + // @@protoc_insertion_point(serialize_to_array_start:onnxflow.OnnxFlowParameter) + ::PROTOBUF_NAMESPACE_ID::uint32 cached_has_bits = 0; + (void) cached_has_bits; + + cached_has_bits = _has_bits_[0]; + // required .google.protobuf.Any data = 1; + if (cached_has_bits & 0x00000001u) { + target = stream->EnsureSpace(target); + target = ::PROTOBUF_NAMESPACE_ID::internal::WireFormatLite:: + InternalWriteMessage( + 1, _Internal::data(this), target, stream); + } + + // required bool requires_grad = 2; + if (cached_has_bits & 0x00000002u) { + target = stream->EnsureSpace(target); + target = ::PROTOBUF_NAMESPACE_ID::internal::WireFormatLite::WriteBoolToArray(2, this->_internal_requires_grad(), target); + } + + // required bool is_parameter = 3; + if (cached_has_bits & 0x00000004u) { + target = stream->EnsureSpace(target); + target = ::PROTOBUF_NAMESPACE_ID::internal::WireFormatLite::WriteBoolToArray(3, this->_internal_is_parameter(), target); + } + + if (PROTOBUF_PREDICT_FALSE(_internal_metadata_.have_unknown_fields())) { + target = stream->WriteRaw(_internal_metadata_.unknown_fields(::PROTOBUF_NAMESPACE_ID::internal::GetEmptyString).data(), + static_cast(_internal_metadata_.unknown_fields(::PROTOBUF_NAMESPACE_ID::internal::GetEmptyString).size()), target); + } + // @@protoc_insertion_point(serialize_to_array_end:onnxflow.OnnxFlowParameter) + return target; +} + +size_t OnnxFlowParameter::RequiredFieldsByteSizeFallback() const { +// @@protoc_insertion_point(required_fields_byte_size_fallback_start:onnxflow.OnnxFlowParameter) + size_t total_size = 0; + + if (_internal_has_data()) { + // required .google.protobuf.Any data = 1; + total_size += 1 + + ::PROTOBUF_NAMESPACE_ID::internal::WireFormatLite::MessageSize( + *data_); + } + + if (_internal_has_requires_grad()) { + // required bool requires_grad = 2; + total_size += 1 + 1; + } + + if (_internal_has_is_parameter()) { + // required bool is_parameter = 3; + total_size += 1 + 1; + } + + return total_size; +} +size_t OnnxFlowParameter::ByteSizeLong() const { +// @@protoc_insertion_point(message_byte_size_start:onnxflow.OnnxFlowParameter) + size_t total_size = 0; + + if (((_has_bits_[0] & 0x00000007) ^ 0x00000007) == 0) { // All required fields are present. + // required .google.protobuf.Any data = 1; + total_size += 1 + + ::PROTOBUF_NAMESPACE_ID::internal::WireFormatLite::MessageSize( + *data_); + + // required bool requires_grad = 2; + total_size += 1 + 1; + + // required bool is_parameter = 3; + total_size += 1 + 1; + + } else { + total_size += RequiredFieldsByteSizeFallback(); + } + ::PROTOBUF_NAMESPACE_ID::uint32 cached_has_bits = 0; + // Prevent compiler warnings about cached_has_bits being unused + (void) cached_has_bits; + + if (PROTOBUF_PREDICT_FALSE(_internal_metadata_.have_unknown_fields())) { + total_size += _internal_metadata_.unknown_fields(::PROTOBUF_NAMESPACE_ID::internal::GetEmptyString).size(); + } + int cached_size = ::PROTOBUF_NAMESPACE_ID::internal::ToCachedSize(total_size); + SetCachedSize(cached_size); + return total_size; +} + +void OnnxFlowParameter::CheckTypeAndMergeFrom( + const ::PROTOBUF_NAMESPACE_ID::MessageLite& from) { + MergeFrom(*::PROTOBUF_NAMESPACE_ID::internal::DownCast( + &from)); +} + +void OnnxFlowParameter::MergeFrom(const OnnxFlowParameter& from) { +// @@protoc_insertion_point(class_specific_merge_from_start:onnxflow.OnnxFlowParameter) + GOOGLE_DCHECK_NE(&from, this); + ::PROTOBUF_NAMESPACE_ID::uint32 cached_has_bits = 0; + (void) cached_has_bits; + + cached_has_bits = from._has_bits_[0]; + if (cached_has_bits & 0x00000007u) { + if (cached_has_bits & 0x00000001u) { + _internal_mutable_data()->::PROTOBUF_NAMESPACE_ID::Any::MergeFrom(from._internal_data()); + } + if (cached_has_bits & 0x00000002u) { + requires_grad_ = from.requires_grad_; + } + if (cached_has_bits & 0x00000004u) { + is_parameter_ = from.is_parameter_; + } + _has_bits_[0] |= cached_has_bits; + } + _internal_metadata_.MergeFrom(from._internal_metadata_); +} + +void OnnxFlowParameter::CopyFrom(const OnnxFlowParameter& from) { +// @@protoc_insertion_point(class_specific_copy_from_start:onnxflow.OnnxFlowParameter) + if (&from == this) return; + Clear(); + MergeFrom(from); +} + +bool OnnxFlowParameter::IsInitialized() const { + if (_Internal::MissingRequiredFields(_has_bits_)) return false; + return true; +} + +void OnnxFlowParameter::InternalSwap(OnnxFlowParameter* other) { + using std::swap; + _internal_metadata_.InternalSwap(&other->_internal_metadata_); + swap(_has_bits_[0], other->_has_bits_[0]); + ::PROTOBUF_NAMESPACE_ID::internal::memswap< + PROTOBUF_FIELD_OFFSET(OnnxFlowParameter, is_parameter_) + + sizeof(OnnxFlowParameter::is_parameter_) + - PROTOBUF_FIELD_OFFSET(OnnxFlowParameter, data_)>( + reinterpret_cast(&data_), + reinterpret_cast(&other->data_)); +} + +std::string OnnxFlowParameter::GetTypeName() const { + return "onnxflow.OnnxFlowParameter"; +} + + +// =================================================================== + +class OnnxFlowParameters::_Internal { + public: +}; + +OnnxFlowParameters::OnnxFlowParameters(::PROTOBUF_NAMESPACE_ID::Arena* arena, + bool is_message_owned) + : ::PROTOBUF_NAMESPACE_ID::MessageLite(arena, is_message_owned), + parameters_(arena) { + SharedCtor(); + if (!is_message_owned) { + RegisterArenaDtor(arena); + } + // @@protoc_insertion_point(arena_constructor:onnxflow.OnnxFlowParameters) +} +OnnxFlowParameters::OnnxFlowParameters(const OnnxFlowParameters& from) + : ::PROTOBUF_NAMESPACE_ID::MessageLite(), + parameters_(from.parameters_) { + _internal_metadata_.MergeFrom(from._internal_metadata_); + // @@protoc_insertion_point(copy_constructor:onnxflow.OnnxFlowParameters) +} + +void OnnxFlowParameters::SharedCtor() { +} + +OnnxFlowParameters::~OnnxFlowParameters() { + // @@protoc_insertion_point(destructor:onnxflow.OnnxFlowParameters) + if (GetArenaForAllocation() != nullptr) return; + SharedDtor(); + _internal_metadata_.Delete(); +} + +inline void OnnxFlowParameters::SharedDtor() { + GOOGLE_DCHECK(GetArenaForAllocation() == nullptr); +} + +void OnnxFlowParameters::ArenaDtor(void* object) { + OnnxFlowParameters* _this = reinterpret_cast< OnnxFlowParameters* >(object); + (void)_this; +} +void OnnxFlowParameters::RegisterArenaDtor(::PROTOBUF_NAMESPACE_ID::Arena*) { +} +void OnnxFlowParameters::SetCachedSize(int size) const { + _cached_size_.Set(size); +} + +void OnnxFlowParameters::Clear() { +// @@protoc_insertion_point(message_clear_start:onnxflow.OnnxFlowParameters) + ::PROTOBUF_NAMESPACE_ID::uint32 cached_has_bits = 0; + // Prevent compiler warnings about cached_has_bits being unused + (void) cached_has_bits; + + parameters_.Clear(); + _internal_metadata_.Clear(); +} + +const char* OnnxFlowParameters::_InternalParse(const char* ptr, ::PROTOBUF_NAMESPACE_ID::internal::ParseContext* ctx) { +#define CHK_(x) if (PROTOBUF_PREDICT_FALSE(!(x))) goto failure + while (!ctx->Done(&ptr)) { + ::PROTOBUF_NAMESPACE_ID::uint32 tag; + ptr = ::PROTOBUF_NAMESPACE_ID::internal::ReadTag(ptr, &tag); + switch (tag >> 3) { + // repeated .onnxflow.OnnxFlowParameter parameters = 1; + case 1: + if (PROTOBUF_PREDICT_TRUE(static_cast<::PROTOBUF_NAMESPACE_ID::uint8>(tag) == 10)) { + ptr -= 1; + do { + ptr += 1; + ptr = ctx->ParseMessage(_internal_add_parameters(), ptr); + CHK_(ptr); + if (!ctx->DataAvailable(ptr)) break; + } while (::PROTOBUF_NAMESPACE_ID::internal::ExpectTag<10>(ptr)); + } else + goto handle_unusual; + continue; + default: + goto handle_unusual; + } // switch + handle_unusual: + if ((tag == 0) || ((tag & 7) == 4)) { + CHK_(ptr); + ctx->SetLastTag(tag); + goto message_done; + } + ptr = UnknownFieldParse( + tag, + _internal_metadata_.mutable_unknown_fields(), + ptr, ctx); + CHK_(ptr != nullptr); + } // while +message_done: + return ptr; +failure: + ptr = nullptr; + goto message_done; +#undef CHK_ +} + +::PROTOBUF_NAMESPACE_ID::uint8* OnnxFlowParameters::_InternalSerialize( + ::PROTOBUF_NAMESPACE_ID::uint8* target, ::PROTOBUF_NAMESPACE_ID::io::EpsCopyOutputStream* stream) const { + // @@protoc_insertion_point(serialize_to_array_start:onnxflow.OnnxFlowParameters) + ::PROTOBUF_NAMESPACE_ID::uint32 cached_has_bits = 0; + (void) cached_has_bits; + + // repeated .onnxflow.OnnxFlowParameter parameters = 1; + for (unsigned int i = 0, + n = static_cast(this->_internal_parameters_size()); i < n; i++) { + target = stream->EnsureSpace(target); + target = ::PROTOBUF_NAMESPACE_ID::internal::WireFormatLite:: + InternalWriteMessage(1, this->_internal_parameters(i), target, stream); + } + + if (PROTOBUF_PREDICT_FALSE(_internal_metadata_.have_unknown_fields())) { + target = stream->WriteRaw(_internal_metadata_.unknown_fields(::PROTOBUF_NAMESPACE_ID::internal::GetEmptyString).data(), + static_cast(_internal_metadata_.unknown_fields(::PROTOBUF_NAMESPACE_ID::internal::GetEmptyString).size()), target); + } + // @@protoc_insertion_point(serialize_to_array_end:onnxflow.OnnxFlowParameters) + return target; +} + +size_t OnnxFlowParameters::ByteSizeLong() const { +// @@protoc_insertion_point(message_byte_size_start:onnxflow.OnnxFlowParameters) + size_t total_size = 0; + + ::PROTOBUF_NAMESPACE_ID::uint32 cached_has_bits = 0; + // Prevent compiler warnings about cached_has_bits being unused + (void) cached_has_bits; + + // repeated .onnxflow.OnnxFlowParameter parameters = 1; + total_size += 1UL * this->_internal_parameters_size(); + for (const auto& msg : this->parameters_) { + total_size += + ::PROTOBUF_NAMESPACE_ID::internal::WireFormatLite::MessageSize(msg); + } + + if (PROTOBUF_PREDICT_FALSE(_internal_metadata_.have_unknown_fields())) { + total_size += _internal_metadata_.unknown_fields(::PROTOBUF_NAMESPACE_ID::internal::GetEmptyString).size(); + } + int cached_size = ::PROTOBUF_NAMESPACE_ID::internal::ToCachedSize(total_size); + SetCachedSize(cached_size); + return total_size; +} + +void OnnxFlowParameters::CheckTypeAndMergeFrom( + const ::PROTOBUF_NAMESPACE_ID::MessageLite& from) { + MergeFrom(*::PROTOBUF_NAMESPACE_ID::internal::DownCast( + &from)); +} + +void OnnxFlowParameters::MergeFrom(const OnnxFlowParameters& from) { +// @@protoc_insertion_point(class_specific_merge_from_start:onnxflow.OnnxFlowParameters) + GOOGLE_DCHECK_NE(&from, this); + ::PROTOBUF_NAMESPACE_ID::uint32 cached_has_bits = 0; + (void) cached_has_bits; + + parameters_.MergeFrom(from.parameters_); + _internal_metadata_.MergeFrom(from._internal_metadata_); +} + +void OnnxFlowParameters::CopyFrom(const OnnxFlowParameters& from) { +// @@protoc_insertion_point(class_specific_copy_from_start:onnxflow.OnnxFlowParameters) + if (&from == this) return; + Clear(); + MergeFrom(from); +} + +bool OnnxFlowParameters::IsInitialized() const { + if (!::PROTOBUF_NAMESPACE_ID::internal::AllAreInitialized(parameters_)) return false; + return true; +} + +void OnnxFlowParameters::InternalSwap(OnnxFlowParameters* other) { + using std::swap; + _internal_metadata_.InternalSwap(&other->_internal_metadata_); + parameters_.InternalSwap(&other->parameters_); +} + +std::string OnnxFlowParameters::GetTypeName() const { + return "onnxflow.OnnxFlowParameters"; +} + + +// @@protoc_insertion_point(namespace_scope) +} // namespace onnxflow +PROTOBUF_NAMESPACE_OPEN +template<> PROTOBUF_NOINLINE ::onnxflow::OnnxFlowParameter* Arena::CreateMaybeMessage< ::onnxflow::OnnxFlowParameter >(Arena* arena) { + return Arena::CreateMessageInternal< ::onnxflow::OnnxFlowParameter >(arena); +} +template<> PROTOBUF_NOINLINE ::onnxflow::OnnxFlowParameters* Arena::CreateMaybeMessage< ::onnxflow::OnnxFlowParameters >(Arena* arena) { + return Arena::CreateMessageInternal< ::onnxflow::OnnxFlowParameters >(arena); +} +PROTOBUF_NAMESPACE_CLOSE + +// @@protoc_insertion_point(global_scope) +#include diff --git a/orttraining/orttraining/onnxflow/csrc/onnxflow.pb.h b/orttraining/orttraining/onnxflow/csrc/onnxflow.pb.h new file mode 100644 index 0000000000..241239577f --- /dev/null +++ b/orttraining/orttraining/onnxflow/csrc/onnxflow.pb.h @@ -0,0 +1,602 @@ +// Generated by the protocol buffer compiler. DO NOT EDIT! +// source: onnxflow.proto + +#ifndef GOOGLE_PROTOBUF_INCLUDED_onnxflow_2eproto +#define GOOGLE_PROTOBUF_INCLUDED_onnxflow_2eproto + +#include +#include + +#include +#if PROTOBUF_VERSION < 3018000 +#error This file was generated by a newer version of protoc which is +#error incompatible with your Protocol Buffer headers. Please update +#error your headers. +#endif +#if 3018001 < PROTOBUF_MIN_PROTOC_VERSION +#error This file was generated by an older version of protoc which is +#error incompatible with your Protocol Buffer headers. Please +#error regenerate this file with a newer version of protoc. +#endif + +#include +#include +#include +#include +#include +#include +#include +#include +#include // IWYU pragma: export +#include // IWYU pragma: export +#include +// @@protoc_insertion_point(includes) +#include +#define PROTOBUF_INTERNAL_EXPORT_onnxflow_2eproto +PROTOBUF_NAMESPACE_OPEN +namespace internal { +class AnyMetadata; +} // namespace internal +PROTOBUF_NAMESPACE_CLOSE + +// Internal implementation detail -- do not use these members. +struct TableStruct_onnxflow_2eproto { + static const ::PROTOBUF_NAMESPACE_ID::internal::ParseTableField entries[] + PROTOBUF_SECTION_VARIABLE(protodesc_cold); + static const ::PROTOBUF_NAMESPACE_ID::internal::AuxiliaryParseTableField aux[] + PROTOBUF_SECTION_VARIABLE(protodesc_cold); + static const ::PROTOBUF_NAMESPACE_ID::internal::ParseTable schema[2] + PROTOBUF_SECTION_VARIABLE(protodesc_cold); + static const ::PROTOBUF_NAMESPACE_ID::internal::FieldMetadata field_metadata[]; + static const ::PROTOBUF_NAMESPACE_ID::internal::SerializationTable serialization_table[]; + static const ::PROTOBUF_NAMESPACE_ID::uint32 offsets[]; +}; +namespace onnxflow { +class OnnxFlowParameter; +struct OnnxFlowParameterDefaultTypeInternal; +extern OnnxFlowParameterDefaultTypeInternal _OnnxFlowParameter_default_instance_; +class OnnxFlowParameters; +struct OnnxFlowParametersDefaultTypeInternal; +extern OnnxFlowParametersDefaultTypeInternal _OnnxFlowParameters_default_instance_; +} // namespace onnxflow +PROTOBUF_NAMESPACE_OPEN +template<> ::onnxflow::OnnxFlowParameter* Arena::CreateMaybeMessage<::onnxflow::OnnxFlowParameter>(Arena*); +template<> ::onnxflow::OnnxFlowParameters* Arena::CreateMaybeMessage<::onnxflow::OnnxFlowParameters>(Arena*); +PROTOBUF_NAMESPACE_CLOSE +namespace onnxflow { + +// =================================================================== + +class OnnxFlowParameter final : + public ::PROTOBUF_NAMESPACE_ID::MessageLite /* @@protoc_insertion_point(class_definition:onnxflow.OnnxFlowParameter) */ { + public: + inline OnnxFlowParameter() : OnnxFlowParameter(nullptr) {} + ~OnnxFlowParameter() override; + explicit constexpr OnnxFlowParameter(::PROTOBUF_NAMESPACE_ID::internal::ConstantInitialized); + + OnnxFlowParameter(const OnnxFlowParameter& from); + OnnxFlowParameter(OnnxFlowParameter&& from) noexcept + : OnnxFlowParameter() { + *this = ::std::move(from); + } + + inline OnnxFlowParameter& operator=(const OnnxFlowParameter& from) { + CopyFrom(from); + return *this; + } + inline OnnxFlowParameter& operator=(OnnxFlowParameter&& from) noexcept { + if (this == &from) return *this; + if (GetOwningArena() == from.GetOwningArena() + #ifdef PROTOBUF_FORCE_COPY_IN_MOVE + && GetOwningArena() != nullptr + #endif // !PROTOBUF_FORCE_COPY_IN_MOVE + ) { + InternalSwap(&from); + } else { + CopyFrom(from); + } + return *this; + } + + inline const std::string& unknown_fields() const { + return _internal_metadata_.unknown_fields(::PROTOBUF_NAMESPACE_ID::internal::GetEmptyString); + } + inline std::string* mutable_unknown_fields() { + return _internal_metadata_.mutable_unknown_fields(); + } + + static const OnnxFlowParameter& default_instance() { + return *internal_default_instance(); + } + static inline const OnnxFlowParameter* internal_default_instance() { + return reinterpret_cast( + &_OnnxFlowParameter_default_instance_); + } + static constexpr int kIndexInFileMessages = + 0; + + friend void swap(OnnxFlowParameter& a, OnnxFlowParameter& b) { + a.Swap(&b); + } + inline void Swap(OnnxFlowParameter* other) { + if (other == this) return; + if (GetOwningArena() == other->GetOwningArena()) { + InternalSwap(other); + } else { + ::PROTOBUF_NAMESPACE_ID::internal::GenericSwap(this, other); + } + } + void UnsafeArenaSwap(OnnxFlowParameter* other) { + if (other == this) return; + GOOGLE_DCHECK(GetOwningArena() == other->GetOwningArena()); + InternalSwap(other); + } + + // implements Message ---------------------------------------------- + + inline OnnxFlowParameter* New() const final { + return new OnnxFlowParameter(); + } + + OnnxFlowParameter* New(::PROTOBUF_NAMESPACE_ID::Arena* arena) const final { + return CreateMaybeMessage(arena); + } + void CheckTypeAndMergeFrom(const ::PROTOBUF_NAMESPACE_ID::MessageLite& from) final; + void CopyFrom(const OnnxFlowParameter& from); + void MergeFrom(const OnnxFlowParameter& from); + PROTOBUF_ATTRIBUTE_REINITIALIZES void Clear() final; + bool IsInitialized() const final; + + size_t ByteSizeLong() const final; + const char* _InternalParse(const char* ptr, ::PROTOBUF_NAMESPACE_ID::internal::ParseContext* ctx) final; + ::PROTOBUF_NAMESPACE_ID::uint8* _InternalSerialize( + ::PROTOBUF_NAMESPACE_ID::uint8* target, ::PROTOBUF_NAMESPACE_ID::io::EpsCopyOutputStream* stream) const final; + void DiscardUnknownFields(); + int GetCachedSize() const final { return _cached_size_.Get(); } + + private: + void SharedCtor(); + void SharedDtor(); + void SetCachedSize(int size) const; + void InternalSwap(OnnxFlowParameter* other); + friend class ::PROTOBUF_NAMESPACE_ID::internal::AnyMetadata; + static ::PROTOBUF_NAMESPACE_ID::StringPiece FullMessageName() { + return "onnxflow.OnnxFlowParameter"; + } + protected: + explicit OnnxFlowParameter(::PROTOBUF_NAMESPACE_ID::Arena* arena, + bool is_message_owned = false); + private: + static void ArenaDtor(void* object); + inline void RegisterArenaDtor(::PROTOBUF_NAMESPACE_ID::Arena* arena); + public: + + std::string GetTypeName() const final; + + // nested types ---------------------------------------------------- + + // accessors ------------------------------------------------------- + + enum : int { + kDataFieldNumber = 1, + kRequiresGradFieldNumber = 2, + kIsParameterFieldNumber = 3, + }; + // required .google.protobuf.Any data = 1; + bool has_data() const; + private: + bool _internal_has_data() const; + public: + void clear_data(); + const ::PROTOBUF_NAMESPACE_ID::Any& data() const; + PROTOBUF_MUST_USE_RESULT ::PROTOBUF_NAMESPACE_ID::Any* release_data(); + ::PROTOBUF_NAMESPACE_ID::Any* mutable_data(); + void set_allocated_data(::PROTOBUF_NAMESPACE_ID::Any* data); + private: + const ::PROTOBUF_NAMESPACE_ID::Any& _internal_data() const; + ::PROTOBUF_NAMESPACE_ID::Any* _internal_mutable_data(); + public: + void unsafe_arena_set_allocated_data( + ::PROTOBUF_NAMESPACE_ID::Any* data); + ::PROTOBUF_NAMESPACE_ID::Any* unsafe_arena_release_data(); + + // required bool requires_grad = 2; + bool has_requires_grad() const; + private: + bool _internal_has_requires_grad() const; + public: + void clear_requires_grad(); + bool requires_grad() const; + void set_requires_grad(bool value); + private: + bool _internal_requires_grad() const; + void _internal_set_requires_grad(bool value); + public: + + // required bool is_parameter = 3; + bool has_is_parameter() const; + private: + bool _internal_has_is_parameter() const; + public: + void clear_is_parameter(); + bool is_parameter() const; + void set_is_parameter(bool value); + private: + bool _internal_is_parameter() const; + void _internal_set_is_parameter(bool value); + public: + + // @@protoc_insertion_point(class_scope:onnxflow.OnnxFlowParameter) + private: + class _Internal; + + // helper for ByteSizeLong() + size_t RequiredFieldsByteSizeFallback() const; + + template friend class ::PROTOBUF_NAMESPACE_ID::Arena::InternalHelper; + typedef void InternalArenaConstructable_; + typedef void DestructorSkippable_; + ::PROTOBUF_NAMESPACE_ID::internal::HasBits<1> _has_bits_; + mutable ::PROTOBUF_NAMESPACE_ID::internal::CachedSize _cached_size_; + ::PROTOBUF_NAMESPACE_ID::Any* data_; + bool requires_grad_; + bool is_parameter_; + friend struct ::TableStruct_onnxflow_2eproto; +}; +// ------------------------------------------------------------------- + +class OnnxFlowParameters final : + public ::PROTOBUF_NAMESPACE_ID::MessageLite /* @@protoc_insertion_point(class_definition:onnxflow.OnnxFlowParameters) */ { + public: + inline OnnxFlowParameters() : OnnxFlowParameters(nullptr) {} + ~OnnxFlowParameters() override; + explicit constexpr OnnxFlowParameters(::PROTOBUF_NAMESPACE_ID::internal::ConstantInitialized); + + OnnxFlowParameters(const OnnxFlowParameters& from); + OnnxFlowParameters(OnnxFlowParameters&& from) noexcept + : OnnxFlowParameters() { + *this = ::std::move(from); + } + + inline OnnxFlowParameters& operator=(const OnnxFlowParameters& from) { + CopyFrom(from); + return *this; + } + inline OnnxFlowParameters& operator=(OnnxFlowParameters&& from) noexcept { + if (this == &from) return *this; + if (GetOwningArena() == from.GetOwningArena() + #ifdef PROTOBUF_FORCE_COPY_IN_MOVE + && GetOwningArena() != nullptr + #endif // !PROTOBUF_FORCE_COPY_IN_MOVE + ) { + InternalSwap(&from); + } else { + CopyFrom(from); + } + return *this; + } + + inline const std::string& unknown_fields() const { + return _internal_metadata_.unknown_fields(::PROTOBUF_NAMESPACE_ID::internal::GetEmptyString); + } + inline std::string* mutable_unknown_fields() { + return _internal_metadata_.mutable_unknown_fields(); + } + + static const OnnxFlowParameters& default_instance() { + return *internal_default_instance(); + } + static inline const OnnxFlowParameters* internal_default_instance() { + return reinterpret_cast( + &_OnnxFlowParameters_default_instance_); + } + static constexpr int kIndexInFileMessages = + 1; + + friend void swap(OnnxFlowParameters& a, OnnxFlowParameters& b) { + a.Swap(&b); + } + inline void Swap(OnnxFlowParameters* other) { + if (other == this) return; + if (GetOwningArena() == other->GetOwningArena()) { + InternalSwap(other); + } else { + ::PROTOBUF_NAMESPACE_ID::internal::GenericSwap(this, other); + } + } + void UnsafeArenaSwap(OnnxFlowParameters* other) { + if (other == this) return; + GOOGLE_DCHECK(GetOwningArena() == other->GetOwningArena()); + InternalSwap(other); + } + + // implements Message ---------------------------------------------- + + inline OnnxFlowParameters* New() const final { + return new OnnxFlowParameters(); + } + + OnnxFlowParameters* New(::PROTOBUF_NAMESPACE_ID::Arena* arena) const final { + return CreateMaybeMessage(arena); + } + void CheckTypeAndMergeFrom(const ::PROTOBUF_NAMESPACE_ID::MessageLite& from) final; + void CopyFrom(const OnnxFlowParameters& from); + void MergeFrom(const OnnxFlowParameters& from); + PROTOBUF_ATTRIBUTE_REINITIALIZES void Clear() final; + bool IsInitialized() const final; + + size_t ByteSizeLong() const final; + const char* _InternalParse(const char* ptr, ::PROTOBUF_NAMESPACE_ID::internal::ParseContext* ctx) final; + ::PROTOBUF_NAMESPACE_ID::uint8* _InternalSerialize( + ::PROTOBUF_NAMESPACE_ID::uint8* target, ::PROTOBUF_NAMESPACE_ID::io::EpsCopyOutputStream* stream) const final; + void DiscardUnknownFields(); + int GetCachedSize() const final { return _cached_size_.Get(); } + + private: + void SharedCtor(); + void SharedDtor(); + void SetCachedSize(int size) const; + void InternalSwap(OnnxFlowParameters* other); + friend class ::PROTOBUF_NAMESPACE_ID::internal::AnyMetadata; + static ::PROTOBUF_NAMESPACE_ID::StringPiece FullMessageName() { + return "onnxflow.OnnxFlowParameters"; + } + protected: + explicit OnnxFlowParameters(::PROTOBUF_NAMESPACE_ID::Arena* arena, + bool is_message_owned = false); + private: + static void ArenaDtor(void* object); + inline void RegisterArenaDtor(::PROTOBUF_NAMESPACE_ID::Arena* arena); + public: + + std::string GetTypeName() const final; + + // nested types ---------------------------------------------------- + + // accessors ------------------------------------------------------- + + enum : int { + kParametersFieldNumber = 1, + }; + // repeated .onnxflow.OnnxFlowParameter parameters = 1; + int parameters_size() const; + private: + int _internal_parameters_size() const; + public: + void clear_parameters(); + ::onnxflow::OnnxFlowParameter* mutable_parameters(int index); + ::PROTOBUF_NAMESPACE_ID::RepeatedPtrField< ::onnxflow::OnnxFlowParameter >* + mutable_parameters(); + private: + const ::onnxflow::OnnxFlowParameter& _internal_parameters(int index) const; + ::onnxflow::OnnxFlowParameter* _internal_add_parameters(); + public: + const ::onnxflow::OnnxFlowParameter& parameters(int index) const; + ::onnxflow::OnnxFlowParameter* add_parameters(); + const ::PROTOBUF_NAMESPACE_ID::RepeatedPtrField< ::onnxflow::OnnxFlowParameter >& + parameters() const; + + // @@protoc_insertion_point(class_scope:onnxflow.OnnxFlowParameters) + private: + class _Internal; + + template friend class ::PROTOBUF_NAMESPACE_ID::Arena::InternalHelper; + typedef void InternalArenaConstructable_; + typedef void DestructorSkippable_; + ::PROTOBUF_NAMESPACE_ID::RepeatedPtrField< ::onnxflow::OnnxFlowParameter > parameters_; + mutable ::PROTOBUF_NAMESPACE_ID::internal::CachedSize _cached_size_; + friend struct ::TableStruct_onnxflow_2eproto; +}; +// =================================================================== + + +// =================================================================== + +#ifdef __GNUC__ + #pragma GCC diagnostic push + #pragma GCC diagnostic ignored "-Wstrict-aliasing" +#endif // __GNUC__ +// OnnxFlowParameter + +// required .google.protobuf.Any data = 1; +inline bool OnnxFlowParameter::_internal_has_data() const { + bool value = (_has_bits_[0] & 0x00000001u) != 0; + PROTOBUF_ASSUME(!value || data_ != nullptr); + return value; +} +inline bool OnnxFlowParameter::has_data() const { + return _internal_has_data(); +} +inline const ::PROTOBUF_NAMESPACE_ID::Any& OnnxFlowParameter::_internal_data() const { + const ::PROTOBUF_NAMESPACE_ID::Any* p = data_; + return p != nullptr ? *p : reinterpret_cast( + ::PROTOBUF_NAMESPACE_ID::_Any_default_instance_); +} +inline const ::PROTOBUF_NAMESPACE_ID::Any& OnnxFlowParameter::data() const { + // @@protoc_insertion_point(field_get:onnxflow.OnnxFlowParameter.data) + return _internal_data(); +} +inline void OnnxFlowParameter::unsafe_arena_set_allocated_data( + ::PROTOBUF_NAMESPACE_ID::Any* data) { + if (GetArenaForAllocation() == nullptr) { + delete reinterpret_cast<::PROTOBUF_NAMESPACE_ID::MessageLite*>(data_); + } + data_ = data; + if (data) { + _has_bits_[0] |= 0x00000001u; + } else { + _has_bits_[0] &= ~0x00000001u; + } + // @@protoc_insertion_point(field_unsafe_arena_set_allocated:onnxflow.OnnxFlowParameter.data) +} +inline ::PROTOBUF_NAMESPACE_ID::Any* OnnxFlowParameter::release_data() { + _has_bits_[0] &= ~0x00000001u; + ::PROTOBUF_NAMESPACE_ID::Any* temp = data_; + data_ = nullptr; +#ifdef PROTOBUF_FORCE_COPY_IN_RELEASE + auto* old = reinterpret_cast<::PROTOBUF_NAMESPACE_ID::MessageLite*>(temp); + temp = ::PROTOBUF_NAMESPACE_ID::internal::DuplicateIfNonNull(temp); + if (GetArenaForAllocation() == nullptr) { delete old; } +#else // PROTOBUF_FORCE_COPY_IN_RELEASE + if (GetArenaForAllocation() != nullptr) { + temp = ::PROTOBUF_NAMESPACE_ID::internal::DuplicateIfNonNull(temp); + } +#endif // !PROTOBUF_FORCE_COPY_IN_RELEASE + return temp; +} +inline ::PROTOBUF_NAMESPACE_ID::Any* OnnxFlowParameter::unsafe_arena_release_data() { + // @@protoc_insertion_point(field_release:onnxflow.OnnxFlowParameter.data) + _has_bits_[0] &= ~0x00000001u; + ::PROTOBUF_NAMESPACE_ID::Any* temp = data_; + data_ = nullptr; + return temp; +} +inline ::PROTOBUF_NAMESPACE_ID::Any* OnnxFlowParameter::_internal_mutable_data() { + _has_bits_[0] |= 0x00000001u; + if (data_ == nullptr) { + auto* p = CreateMaybeMessage<::PROTOBUF_NAMESPACE_ID::Any>(GetArenaForAllocation()); + data_ = p; + } + return data_; +} +inline ::PROTOBUF_NAMESPACE_ID::Any* OnnxFlowParameter::mutable_data() { + ::PROTOBUF_NAMESPACE_ID::Any* _msg = _internal_mutable_data(); + // @@protoc_insertion_point(field_mutable:onnxflow.OnnxFlowParameter.data) + return _msg; +} +inline void OnnxFlowParameter::set_allocated_data(::PROTOBUF_NAMESPACE_ID::Any* data) { + ::PROTOBUF_NAMESPACE_ID::Arena* message_arena = GetArenaForAllocation(); + if (message_arena == nullptr) { + delete reinterpret_cast< ::PROTOBUF_NAMESPACE_ID::MessageLite*>(data_); + } + if (data) { + ::PROTOBUF_NAMESPACE_ID::Arena* submessage_arena = + ::PROTOBUF_NAMESPACE_ID::Arena::InternalHelper< + ::PROTOBUF_NAMESPACE_ID::MessageLite>::GetOwningArena( + reinterpret_cast<::PROTOBUF_NAMESPACE_ID::MessageLite*>(data)); + if (message_arena != submessage_arena) { + data = ::PROTOBUF_NAMESPACE_ID::internal::GetOwnedMessage( + message_arena, data, submessage_arena); + } + _has_bits_[0] |= 0x00000001u; + } else { + _has_bits_[0] &= ~0x00000001u; + } + data_ = data; + // @@protoc_insertion_point(field_set_allocated:onnxflow.OnnxFlowParameter.data) +} + +// required bool requires_grad = 2; +inline bool OnnxFlowParameter::_internal_has_requires_grad() const { + bool value = (_has_bits_[0] & 0x00000002u) != 0; + return value; +} +inline bool OnnxFlowParameter::has_requires_grad() const { + return _internal_has_requires_grad(); +} +inline void OnnxFlowParameter::clear_requires_grad() { + requires_grad_ = false; + _has_bits_[0] &= ~0x00000002u; +} +inline bool OnnxFlowParameter::_internal_requires_grad() const { + return requires_grad_; +} +inline bool OnnxFlowParameter::requires_grad() const { + // @@protoc_insertion_point(field_get:onnxflow.OnnxFlowParameter.requires_grad) + return _internal_requires_grad(); +} +inline void OnnxFlowParameter::_internal_set_requires_grad(bool value) { + _has_bits_[0] |= 0x00000002u; + requires_grad_ = value; +} +inline void OnnxFlowParameter::set_requires_grad(bool value) { + _internal_set_requires_grad(value); + // @@protoc_insertion_point(field_set:onnxflow.OnnxFlowParameter.requires_grad) +} + +// required bool is_parameter = 3; +inline bool OnnxFlowParameter::_internal_has_is_parameter() const { + bool value = (_has_bits_[0] & 0x00000004u) != 0; + return value; +} +inline bool OnnxFlowParameter::has_is_parameter() const { + return _internal_has_is_parameter(); +} +inline void OnnxFlowParameter::clear_is_parameter() { + is_parameter_ = false; + _has_bits_[0] &= ~0x00000004u; +} +inline bool OnnxFlowParameter::_internal_is_parameter() const { + return is_parameter_; +} +inline bool OnnxFlowParameter::is_parameter() const { + // @@protoc_insertion_point(field_get:onnxflow.OnnxFlowParameter.is_parameter) + return _internal_is_parameter(); +} +inline void OnnxFlowParameter::_internal_set_is_parameter(bool value) { + _has_bits_[0] |= 0x00000004u; + is_parameter_ = value; +} +inline void OnnxFlowParameter::set_is_parameter(bool value) { + _internal_set_is_parameter(value); + // @@protoc_insertion_point(field_set:onnxflow.OnnxFlowParameter.is_parameter) +} + +// ------------------------------------------------------------------- + +// OnnxFlowParameters + +// repeated .onnxflow.OnnxFlowParameter parameters = 1; +inline int OnnxFlowParameters::_internal_parameters_size() const { + return parameters_.size(); +} +inline int OnnxFlowParameters::parameters_size() const { + return _internal_parameters_size(); +} +inline void OnnxFlowParameters::clear_parameters() { + parameters_.Clear(); +} +inline ::onnxflow::OnnxFlowParameter* OnnxFlowParameters::mutable_parameters(int index) { + // @@protoc_insertion_point(field_mutable:onnxflow.OnnxFlowParameters.parameters) + return parameters_.Mutable(index); +} +inline ::PROTOBUF_NAMESPACE_ID::RepeatedPtrField< ::onnxflow::OnnxFlowParameter >* +OnnxFlowParameters::mutable_parameters() { + // @@protoc_insertion_point(field_mutable_list:onnxflow.OnnxFlowParameters.parameters) + return ¶meters_; +} +inline const ::onnxflow::OnnxFlowParameter& OnnxFlowParameters::_internal_parameters(int index) const { + return parameters_.Get(index); +} +inline const ::onnxflow::OnnxFlowParameter& OnnxFlowParameters::parameters(int index) const { + // @@protoc_insertion_point(field_get:onnxflow.OnnxFlowParameters.parameters) + return _internal_parameters(index); +} +inline ::onnxflow::OnnxFlowParameter* OnnxFlowParameters::_internal_add_parameters() { + return parameters_.Add(); +} +inline ::onnxflow::OnnxFlowParameter* OnnxFlowParameters::add_parameters() { + ::onnxflow::OnnxFlowParameter* _add = _internal_add_parameters(); + // @@protoc_insertion_point(field_add:onnxflow.OnnxFlowParameters.parameters) + return _add; +} +inline const ::PROTOBUF_NAMESPACE_ID::RepeatedPtrField< ::onnxflow::OnnxFlowParameter >& +OnnxFlowParameters::parameters() const { + // @@protoc_insertion_point(field_list:onnxflow.OnnxFlowParameters.parameters) + return parameters_; +} + +#ifdef __GNUC__ + #pragma GCC diagnostic pop +#endif // __GNUC__ +// ------------------------------------------------------------------- + + +// @@protoc_insertion_point(namespace_scope) + +} // namespace onnxflow + +// @@protoc_insertion_point(global_scope) + +#include +#endif // GOOGLE_PROTOBUF_INCLUDED_GOOGLE_PROTOBUF_INCLUDED_onnxflow_2eproto diff --git a/orttraining/orttraining/onnxflow/models/simple_model.onnx b/orttraining/orttraining/onnxflow/models/simple_model.onnx new file mode 100644 index 0000000000..2a81d1fd28 Binary files /dev/null and b/orttraining/orttraining/onnxflow/models/simple_model.onnx differ diff --git a/orttraining/orttraining/onnxflow/onnxflow.proto b/orttraining/orttraining/onnxflow/onnxflow.proto new file mode 100644 index 0000000000..83791682d6 --- /dev/null +++ b/orttraining/orttraining/onnxflow/onnxflow.proto @@ -0,0 +1,17 @@ +syntax = "proto2"; + +package onnxflow; + +import "google/protobuf/any.proto"; + +message OnnxFlowParameter { + required google.protobuf.Any data = 1; + required bool requires_grad = 2; + required bool is_parameter = 3; +} + +message OnnxFlowParameters { + repeated OnnxFlowParameter parameters = 1; +} + +option optimize_for = LITE_RUNTIME; diff --git a/orttraining/orttraining/onnxflow/onnxflow/__init__.py b/orttraining/orttraining/onnxflow/onnxflow/__init__.py new file mode 100644 index 0000000000..5daf4c1039 --- /dev/null +++ b/orttraining/orttraining/onnxflow/onnxflow/__init__.py @@ -0,0 +1,7 @@ +from .graph import Graph, TrainingGraph +from .loss import MSELoss, CrossEntropyLoss +from .optim import AdamW + +def save(parameters, path_to_file): + with open(path_to_file, 'wb') as file_object: + file_object.write(parameters.SerializeToString()) diff --git a/orttraining/orttraining/onnxflow/onnxflow/graph.py b/orttraining/orttraining/onnxflow/onnxflow/graph.py new file mode 100644 index 0000000000..7a382fd92e --- /dev/null +++ b/orttraining/orttraining/onnxflow/onnxflow/graph.py @@ -0,0 +1,179 @@ + +from abc import ABC, abstractmethod +from onnxruntime.capi._pybind_state import GradientGraphBuilder +import onnx +import copy +from .onnxflow_pb2 import OnnxFlowParameter, OnnxFlowParameters + +def _build_gradient_model(model, requires_grad_params, frozen_params): + # Collect names of parameters that need gradients computed + trainable_parameters = set() + # Move all trainable and non trainable initializers to graph inputs. + # This allows training to pass in the parameters from outside the graph + # so as to share the parameters across multiple sessions. + graph_inputs = model.graph.input + initializers = [] + for initializer in model.graph.initializer: + if not initializer.name[0].isdigit(): + # Move onl those initializers as inputs that are not local + # to the onnx model. i.e. initializers that are model parameters. + # These are tpically those initializers without any number prefixed + # to their names. + graph_inputs.append( + onnx.helper.make_tensor_value_info(initializer.name, + initializer.data_type, + initializer.dims)) + if initializer.name not in frozen_params: + trainable_parameters.add(initializer.name) + else: + # All other initializers stay where they were. + initializers.append(initializer) + + # Graph and model with initializers as inputs. + graph_with_initializers_as_inputs = onnx.helper.make_graph(model.graph.node, + 'graph_with_initializers_as_inputs', + graph_inputs, model.graph.output, + initializer=initializers) + grad_model = onnx.helper.make_model(graph_with_initializers_as_inputs, + producer_name='onnxflow', + opset_imports=[ + onnx.helper.make_opsetid('com.microsoft', 1)] + \ + list(model.opset_import)) + + # Any parameter or input that requires gradient, should have been already added to + # requires_grad_params + for parameter_name in requires_grad_params: + trainable_parameters.add(parameter_name) + + # Assumption is that the graph has an output called `loss`. + builder = GradientGraphBuilder(grad_model.SerializeToString(), + {'loss'}, + trainable_parameters, + 'loss') + builder.build() + return onnx.load_from_string(builder.get_model()) + +def _build_gradient_accumulation_model(grad_model): + graph_inputs = grad_model.graph.input + graph_nodes = grad_model.graph.node + graph_outputs = grad_model.graph.output + for idx, graph_output in enumerate(grad_model.graph.output): + # if the graph output ends with _grad, + # assume that that output is a gradient output + if not graph_output.name.endswith('_grad'): + continue + + # gradient accumulation node inputs and output names + grad_name = graph_output.name + grad_accumulation_buffer_name = f'{grad_name}.accumulation.buffer' + grad_accumulation_output_name = f'{grad_name}.accumulation.out' + + # Gradient accumulation node + acc_node = onnx.helper.make_node("InPlaceAccumulator", + [grad_accumulation_buffer_name, grad_name], + [grad_accumulation_output_name], + name=f"GradientAccumulator{idx}", + domain='com.microsoft') + + graph_nodes.append(acc_node) + + # grad buffer is also a graph input + grad_accumulation_buffer_input = copy.deepcopy(graph_output) + grad_accumulation_buffer_input.name = grad_accumulation_buffer_name + graph_inputs.append(grad_accumulation_buffer_input) + + # accumulated gradient is also a graph output + grad_accumulation_output = copy.deepcopy(graph_output) + grad_accumulation_output.name = grad_accumulation_output_name + graph_outputs.append(grad_accumulation_output) + + graph = onnx.helper.make_graph(graph_nodes, 'GradientGraph', + graph_inputs, + graph_outputs, + grad_model.graph.initializer) + return onnx.helper.make_model(graph, producer_name='onnxflow', + opset_imports=list(grad_model.opset_import)) + + +def _get_model_parameters(model, requires_grad_params, frozen_params): + parameters = OnnxFlowParameters() + for initializer in model.graph.initializer: + if not initializer.name[0].isdigit(): + param = OnnxFlowParameter() + param.requires_grad = True + if initializer.name in frozen_params: + param.requires_grad = False + param.data.Pack(initializer) + param.is_parameter = True + parameters.parameters.append(param) + + requires_grad_params_set = set(requires_grad_params) + for graph_input in model.graph.input: + if graph_input.name in requires_grad_params_set: + param = OnnxFlowParameter() + param.requires_grad = True + param.data.Pack(graph_input) + param.is_parameter = False + parameters.parameters.append(param) + + return parameters + + +class Graph(ABC): + def __init__(self): + pass + + @abstractmethod + def build(self, *args, **kwargs): + ... + + def __call__(self, *args, **kwargs): + # build the user model + user_model = self.build(*args, **kwargs) + + # validate and check the model + onnx.checker.check_model(user_model, True) + + return user_model + +class TrainingGraph(Graph): + def __init__(self): + super(TrainingGraph, self).__init__() + self._frozen = set() + self._requires_grad = [] + self._parameters = None + + @abstractmethod + def build(self, *args, **kwargs): + ... + + def freeze_parameter(self, parameter_name): + self._frozen.add(parameter_name) + + def requires_grad(self, parameter_name): + self._requires_grad.append(parameter_name) + + def parameters(self): + # return parameters that can be serialized by the user + if self._parameters is None: + raise RuntimeError("Please build the training graph first before trying to retrieve the parameters.") + + return self._parameters + + def __call__(self, *args, **kwargs): + # build the user model + user_model = self.build(*args, **kwargs) + + # get all the model parameters for the user_model + self._parameters = _get_model_parameters(user_model, self._requires_grad, self._frozen) + + # build the gradient graph + grad_model = _build_gradient_model(user_model, self._requires_grad, self._frozen) + + # add gradient accumulation nodes + grad_model = _build_gradient_accumulation_model(grad_model) + + # validate and check the model + onnx.checker.check_model(grad_model, True) + + return grad_model diff --git a/orttraining/orttraining/onnxflow/onnxflow/loss.py b/orttraining/orttraining/onnxflow/onnxflow/loss.py new file mode 100644 index 0000000000..484cf28177 --- /dev/null +++ b/orttraining/orttraining/onnxflow/onnxflow/loss.py @@ -0,0 +1,136 @@ +from .graph import Graph +import onnx +import onnx +from onnx import helper +from onnx import TensorProto, OperatorSetIdProto +import copy + +class MSELoss(Graph): + def __init__(self): + super(MSELoss, self).__init__() + + def build(self, base_model, output, target='target', reduction='mean'): + # Ideally + # model = onnx_model.make_functional() + # loss_unreduced = onnx.Pow(onnx.Sub(model(), target), 2) + # if reduction == 'mean': + # loss = onnx.ReduceMean(loss_unreduced) + # elif reduction == 'sum': + # loss = onnx.ReduceSum(loss_unreduced) + # return loss + + # deepcopy the base model so we don't inadvertently change the original model + onnx_model = copy.deepcopy(base_model) + + # determine the reduction type + if reduction != 'mean' and reduction != 'sum': + raise RuntimeError('not supported reduction') + + graph_nodes = onnx_model.graph.node + graph_inputs = onnx_model.graph.input + + # create a new graph input. this is the target input needed to compare the + # graph output against to calculate loss. + target_input = copy.deepcopy(onnx_model.graph.output[0]) + target_input.name = target + graph_inputs.append(target_input) + + # create a new graph output for loss + graph_outputs = [helper.make_tensor_value_info('loss', TensorProto.FLOAT, [1, 1])] + + graph_initializers = onnx_model.graph.initializer + + # loss equation + # loss = reduce((output-target)^2) + + # create the sub node + sub_node_input_names = [output, target] + sub_node_output_names = ['loss_sub_output'] + sub_node = helper.make_node("Sub", + sub_node_input_names, + sub_node_output_names, + name=f"MSELossSub") + graph_nodes.append(sub_node) + + # create the square node + pow_node_input_names = sub_node_output_names + pow_node_input_names.append('0_pow_exponent') + pow_node_output_names = ['loss_pow_output'] + pow_node = helper.make_node("Pow", + pow_node_input_names, + pow_node_output_names, + name=f"MSELossPow") + graph_nodes.append(pow_node) + graph_initializers.append(helper.make_tensor('0_pow_exponent', TensorProto.FLOAT, [1], [2.0])) + + # create the reduce node + reduce_node_input_names = pow_node_output_names + reduce_node_output_names = ['loss'] + reduce_node = helper.make_node("ReduceMean" if reduction == 'mean' else "ReduceSum", + reduce_node_input_names, + reduce_node_output_names, + name=f"MSELossReduce") + graph_nodes.append(reduce_node) + + # generate the graph and model with above inputs, outputs, initializers and nodes + graph = helper.make_graph(graph_nodes, 'GraphWithLoss', graph_inputs, graph_outputs, graph_initializers) + model = helper.make_model(graph, producer_name='onnxflow', + opset_imports=[helper.make_opsetid('com.microsoft', 1)]+ list(base_model.opset_import)) + + return model + +class CrossEntropyLoss(Graph): + def __init__(self): + super(CrossEntropyLoss, self).__init__() + + def build(self, base_model, output, target='target', weights=None, reduction='mean', ignore_index=None, get_log_prob=False): + + # Ideally + # model = onnx_model.make_functional() + # loss = onnx.SoftmaxCrossEntropyLoss(output, target, weights, reduction) + # return loss + + # deepcopy the base model so we don't inadvertently change the original model + onnx_model = copy.deepcopy(base_model) + + # determine the reduction type + if reduction != 'mean' and reduction != 'sum': + raise RuntimeError('not supported reduction') + + graph_nodes = onnx_model.graph.node + graph_inputs = onnx_model.graph.input + + # create a new graph input. this is the target input needed to compare the + # graph output against to calculate loss. + target_input = copy.deepcopy(onnx_model.graph.output[0]) + target_input.name = target + target_input.type.tensor_type.elem_type = TensorProto.INT32 + graph_inputs.append(target_input) + + # create a new graph output for loss + graph_outputs = [helper.make_tensor_value_info('loss', TensorProto.FLOAT, [])] + graph_initializers = onnx_model.graph.initializer + + # create the loss node + loss_node_input_name = [output, target] + if weights: + loss_node_input_name.append('weights') + loss_node_output_name = ['loss', 'log_prob'] + loss_node = helper.make_node("SoftmaxCrossEntropyLoss", + loss_node_input_name, + loss_node_output_name, + reduction=reduction, + ignore_index=ignore_index, + name=f"SoftmaxCrossEntropyLoss") + graph_nodes.append(loss_node) + + # generate the graph and model with above inputs, outputs, initializers and nodes + # TODO: user model generated by opset 11 does not have SoftmaxCrossEntropyLoss. + # we need to probably enfore opset versions. + graph = helper.make_graph(graph_nodes, 'GraphWithLoss', graph_inputs, graph_outputs, graph_initializers) + model = helper.make_model(graph, producer_name='onnxflow', + opset_imports=[onnx.helper.make_opsetid("", 12)]) + + return model + +# TODO: BCEWithLogitsLoss diff --git a/orttraining/orttraining/onnxflow/onnxflow/onnxflow_pb2.py b/orttraining/orttraining/onnxflow/onnxflow/onnxflow_pb2.py new file mode 100644 index 0000000000..a09b66659b --- /dev/null +++ b/orttraining/orttraining/onnxflow/onnxflow/onnxflow_pb2.py @@ -0,0 +1,129 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: onnxflow.proto +"""Generated protocol buffer code.""" +from google.protobuf import descriptor as _descriptor +from google.protobuf import message as _message +from google.protobuf import reflection as _reflection +from google.protobuf import symbol_database as _symbol_database +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from google.protobuf import any_pb2 as google_dot_protobuf_dot_any__pb2 + + +DESCRIPTOR = _descriptor.FileDescriptor( + name='onnxflow.proto', + package='onnxflow', + syntax='proto2', + serialized_options=b'H\003', + create_key=_descriptor._internal_create_key, + serialized_pb=b'\n\x0eonnxflow.proto\x12\x08onnxflow\x1a\x19google/protobuf/any.proto\"d\n\x11OnnxFlowParameter\x12\"\n\x04\x64\x61ta\x18\x01 \x02(\x0b\x32\x14.google.protobuf.Any\x12\x15\n\rrequires_grad\x18\x02 \x02(\x08\x12\x14\n\x0cis_parameter\x18\x03 \x02(\x08\"E\n\x12OnnxFlowParameters\x12/\n\nparameters\x18\x01 \x03(\x0b\x32\x1b.onnxflow.OnnxFlowParameterB\x02H\x03' + , + dependencies=[google_dot_protobuf_dot_any__pb2.DESCRIPTOR,]) + + + + +_ONNXFLOWPARAMETER = _descriptor.Descriptor( + name='OnnxFlowParameter', + full_name='onnxflow.OnnxFlowParameter', + filename=None, + file=DESCRIPTOR, + containing_type=None, + create_key=_descriptor._internal_create_key, + fields=[ + _descriptor.FieldDescriptor( + name='data', full_name='onnxflow.OnnxFlowParameter.data', index=0, + number=1, type=11, cpp_type=10, label=2, + has_default_value=False, default_value=None, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + serialized_options=None, file=DESCRIPTOR, create_key=_descriptor._internal_create_key), + _descriptor.FieldDescriptor( + name='requires_grad', full_name='onnxflow.OnnxFlowParameter.requires_grad', index=1, + number=2, type=8, cpp_type=7, label=2, + has_default_value=False, default_value=False, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + serialized_options=None, file=DESCRIPTOR, create_key=_descriptor._internal_create_key), + _descriptor.FieldDescriptor( + name='is_parameter', full_name='onnxflow.OnnxFlowParameter.is_parameter', index=2, + number=3, type=8, cpp_type=7, label=2, + has_default_value=False, default_value=False, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + serialized_options=None, file=DESCRIPTOR, create_key=_descriptor._internal_create_key), + ], + extensions=[ + ], + nested_types=[], + enum_types=[ + ], + serialized_options=None, + is_extendable=False, + syntax='proto2', + extension_ranges=[], + oneofs=[ + ], + serialized_start=55, + serialized_end=155, +) + + +_ONNXFLOWPARAMETERS = _descriptor.Descriptor( + name='OnnxFlowParameters', + full_name='onnxflow.OnnxFlowParameters', + filename=None, + file=DESCRIPTOR, + containing_type=None, + create_key=_descriptor._internal_create_key, + fields=[ + _descriptor.FieldDescriptor( + name='parameters', full_name='onnxflow.OnnxFlowParameters.parameters', index=0, + number=1, type=11, cpp_type=10, label=3, + has_default_value=False, default_value=[], + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + serialized_options=None, file=DESCRIPTOR, create_key=_descriptor._internal_create_key), + ], + extensions=[ + ], + nested_types=[], + enum_types=[ + ], + serialized_options=None, + is_extendable=False, + syntax='proto2', + extension_ranges=[], + oneofs=[ + ], + serialized_start=157, + serialized_end=226, +) + +_ONNXFLOWPARAMETER.fields_by_name['data'].message_type = google_dot_protobuf_dot_any__pb2._ANY +_ONNXFLOWPARAMETERS.fields_by_name['parameters'].message_type = _ONNXFLOWPARAMETER +DESCRIPTOR.message_types_by_name['OnnxFlowParameter'] = _ONNXFLOWPARAMETER +DESCRIPTOR.message_types_by_name['OnnxFlowParameters'] = _ONNXFLOWPARAMETERS +_sym_db.RegisterFileDescriptor(DESCRIPTOR) + +OnnxFlowParameter = _reflection.GeneratedProtocolMessageType('OnnxFlowParameter', (_message.Message,), { + 'DESCRIPTOR' : _ONNXFLOWPARAMETER, + '__module__' : 'onnxflow_pb2' + # @@protoc_insertion_point(class_scope:onnxflow.OnnxFlowParameter) + }) +_sym_db.RegisterMessage(OnnxFlowParameter) + +OnnxFlowParameters = _reflection.GeneratedProtocolMessageType('OnnxFlowParameters', (_message.Message,), { + 'DESCRIPTOR' : _ONNXFLOWPARAMETERS, + '__module__' : 'onnxflow_pb2' + # @@protoc_insertion_point(class_scope:onnxflow.OnnxFlowParameters) + }) +_sym_db.RegisterMessage(OnnxFlowParameters) + + +DESCRIPTOR._options = None +# @@protoc_insertion_point(module_scope) diff --git a/orttraining/orttraining/onnxflow/onnxflow/optim.py b/orttraining/orttraining/onnxflow/onnxflow/optim.py new file mode 100644 index 0000000000..c872c96f04 --- /dev/null +++ b/orttraining/orttraining/onnxflow/onnxflow/optim.py @@ -0,0 +1,144 @@ +from .graph import Graph +import onnx +import onnx +from onnx import helper +from onnx import TensorProto, OperatorSetIdProto +import copy + + +class AdamW(Graph): + def __init__(self, bias_correction=True, betas=(0.9, 0.999), eps=1e-6, weight_decay=0.): + super(AdamW, self).__init__() + self.bias_correction = bias_correction + self.betas = betas + self.eps = eps + self.weight_decay = weight_decay + + # TODO: fix this to move outside of optimizer node in ORT backend + self.max_norm_clip = 1 + + def build(self, base_model): + # Ideally + # model = onnx_model.make_functional() + # loss_unreduced = onnx.Pow(onnx.Sub(model(), target), 2) + # if reduction == 'mean': + # loss = onnx.ReduceMean(loss_unreduced) + # elif reduction == 'sum': + # loss = onnx.ReduceSum(loss_unreduced) + # return loss + + learning_rate_name = 'learning_rate' + step_name = 'step' + + graph_nodes = [] + graph_inputs = [ + helper.make_tensor_value_info(learning_rate_name, TensorProto.FLOAT , [1]), + helper.make_tensor_value_info(step_name, TensorProto.INT64, [1])] + graph_outputs = [] + + for idx, graph_output in enumerate(base_model.graph.output): + if not graph_output.name.endswith('_grad'): + continue + + weight_name = graph_output.name[:-5] + grad_name = graph_output.name + first_order_moment_name = weight_name + '.exp_avg' + second_order_moment_name = weight_name + '.exp_avg_sq' + mixed_precision_name = weight_name + '.mixed_precision' + loss_scaler_name = weight_name + '.loss_scaler' + gradient_norm_name = weight_name + '.global_gradient_norm' + should_update_name = weight_name + '.should_update' + # prepare node (and graph) inputs and outputs + node_input_names = [learning_rate_name, # learning rate + step_name, # training step (used for beta correction) + weight_name, # weight to be updated + grad_name, # gradient of the weight to be used for update + first_order_moment_name, # first order moment for this weight + second_order_moment_name, # second order moment for this weight + mixed_precision_name, # mixed precision weight representation (required if computation to be done in mp) + loss_scaler_name, # used for gradient scaling + gradient_norm_name, # used for gradient scaling + should_update_name] # whether or not to skip updating the weights + + weight_tensor_value_info = copy.deepcopy(graph_output) + weight_tensor_value_info.name = weight_name + first_order_moment_tensor_value_info = copy.deepcopy(graph_output) + first_order_moment_tensor_value_info.name = first_order_moment_name + second_order_moment_tensor_value_info = copy.deepcopy(graph_output) + second_order_moment_tensor_value_info.name = second_order_moment_name + node_inputs = [ + weight_tensor_value_info, + copy.deepcopy(graph_output), + first_order_moment_tensor_value_info, + second_order_moment_tensor_value_info, + helper.make_tensor_value_info(mixed_precision_name, TensorProto.FLOAT16 , [0]), + helper.make_tensor_value_info(loss_scaler_name, TensorProto.FLOAT, []), + helper.make_tensor_value_info(gradient_norm_name, TensorProto.FLOAT, []), + helper.make_tensor_value_info(should_update_name, TensorProto.BOOL, [1]), + ] + graph_inputs.extend(node_inputs) + + step_output_name = f'{weight_name}.{step_name}.out' + first_order_moment_output_name = f'{first_order_moment_name}.out' + second_order_moment_output_name = f'{second_order_moment_name}.out' + weight_output_name = f'{weight_name}.out' + grad_output_name = f'{grad_name}.out' + mixed_precision_output_name = f'{mixed_precision_name}.out' + + first_order_moment_output_tensor_value_info = copy.deepcopy(graph_output) + first_order_moment_output_tensor_value_info.name = first_order_moment_output_name + second_order_moment_output_tensor_value_info = copy.deepcopy(graph_output) + second_order_moment_output_tensor_value_info.name = second_order_moment_output_name + weight_output_tensor_value_info = copy.deepcopy(graph_output) + weight_output_tensor_value_info.name = weight_output_name + grad_output_tensor_value_info = copy.deepcopy(graph_output) + grad_output_tensor_value_info.name = grad_output_name + + + node_output_names = [step_output_name, # step out + first_order_moment_output_name, # first order moment output + second_order_moment_output_name, # second order moment output + weight_output_name, # updated weights + grad_output_name, # gradients output + mixed_precision_output_name] # updated mixed precision weights + + node_outputs = [ + helper.make_tensor_value_info(step_output_name, TensorProto.INT64, [1]), + first_order_moment_output_tensor_value_info, + second_order_moment_output_tensor_value_info, + weight_output_tensor_value_info, + grad_output_tensor_value_info, + helper.make_tensor_value_info(mixed_precision_output_name, TensorProto.FLOAT16, [0]) + ] + graph_outputs.extend(node_outputs) + + # node attributes + node_attributes = { + 'alpha': self.betas[0], # beta1 + 'beta': self.betas[1], # beta2 + 'lambda': self.weight_decay, # weight decay + 'epsilon': self.eps, # epsilon + 'do_bias_correction': 1 if self.bias_correction else 0, # bias_correction + 'weight_decay_mode': 1, # weight decay mode 1 implies transformers adamw 0 implies pytorch adamw + 'max_norm_clip': self.max_norm_clip # used for gradient scaling + } + + # gradient scaling equation: + # if global_gradient_norm > loss_scaler*max_norm_clip: global_gradient_norm / max_norm_clip + # else: loss_scaler*max_norm_clip + + # make the node + optimizer_node = helper.make_node("AdamOptimizer", + node_input_names, + node_output_names, + name=f"AdamOptimizer{idx}", + domain='com.microsoft', + **node_attributes) + + graph_nodes.append(optimizer_node) + + # make the graph and the model + graph = helper.make_graph(graph_nodes, 'AdamOptimizerGraph', graph_inputs, graph_outputs) + model = helper.make_model(graph, producer_name='onnxflow', + opset_imports=[helper.make_opsetid('com.microsoft', 1)]) + return model diff --git a/orttraining/orttraining/onnxflow/sample.m.cpp b/orttraining/orttraining/onnxflow/sample.m.cpp new file mode 100644 index 0000000000..4830da74f9 --- /dev/null +++ b/orttraining/orttraining/onnxflow/sample.m.cpp @@ -0,0 +1,32 @@ +#include "orttraining/onnxflow/csrc/load_parameters.h" +#include +#include +#include + +int main() +{ + std::string path_to_parameters_proto; + std::cout << "Provide the absolute path to the parameters.of file\n"; + std::cin >> path_to_parameters_proto; + std::filesystem::path path{path_to_parameters_proto}; + + auto parameters = onnxflow::load_parameters(std::filesystem::absolute(path).string()); + + std::cout << "The parameters are:\n"; + for (const auto& param : parameters.parameters()) + { + if (param.is_parameter()) + { + onnx::TensorProto tensor; + param.data().UnpackTo(&tensor); + std::cout << "<" << tensor.name() << ", requires_grad=" << (param.requires_grad() ? "True" : "False") << ">" << std::endl; + } else + { + onnx::ValueInfoProto valueinfo; + param.data().UnpackTo(&valueinfo); + std::cout << "<" << valueinfo.name() << ", requires_grad=" << (param.requires_grad() ? "True" : "False") << ">" << std::endl; + } + } + + return 0; +} \ No newline at end of file diff --git a/orttraining/orttraining/onnxflow/sample.py b/orttraining/orttraining/onnxflow/sample.py new file mode 100644 index 0000000000..6d2cef6095 --- /dev/null +++ b/orttraining/orttraining/onnxflow/sample.py @@ -0,0 +1,36 @@ +from onnxflow import TrainingGraph, Graph +import onnxflow +import onnx + +class MyGraph(TrainingGraph): + def __init__(self, base_model): + super(MyGraph, self).__init__() + self.loss = onnxflow.loss.MSELoss() + self.base_model = base_model + + def build(self): + outputs = self.base_model.graph.output + lossful_graph = self.loss(self.base_model, outputs[0].name) + return lossful_graph + +onnxfile = 'models/simple_model.onnx' +model = onnx.load(onnxfile) + +graph = MyGraph(model) + +# remove in case of any model other than simple_model.onnx +graph.requires_grad('_original_module.fc1.weight') +graph.requires_grad('_original_module.fc1.bias') +graph.requires_grad('_original_module.fc2.weight') +graph.requires_grad('_original_module.fc2.bias') + +gradient_graph = graph() + +parameters = graph.parameters() +onnxflow.save(parameters, 'parameters.of') + +optimizer = onnxflow.optim.AdamW() +optimizer_graph = optimizer(gradient_graph) + +onnx.save(gradient_graph, "gradient_graph.onnx") +onnx.save(optimizer_graph, "optimizer_graph.onnx") diff --git a/tools/ci_build/build.py b/tools/ci_build/build.py index dcf8913925..17b71539fc 100644 --- a/tools/ci_build/build.py +++ b/tools/ci_build/build.py @@ -584,6 +584,11 @@ def parse_arguments(): parser.add_argument('--eager_customop_header', default=None, help='Header containing custom op definitions for eager mode.') + # on device training + parser.add_argument( + "--build_on_device_training", action='store_true', + help="Build on device training.") + parser.add_argument( "--enable_external_custom_op_schemas", action='store_true', help="Enable registering user defined custom operation schemas at shared library load time.\ @@ -860,6 +865,7 @@ def generate_build_tree(cmake_path, source_dir, build_dir, cuda_home, cudnn_home "-Donnxruntime_ENABLE_WEBASSEMBLY_DEBUG_INFO=" + ("ON" if args.enable_wasm_debug_info else "OFF"), "-Donnxruntime_ENABLE_WEBASSEMBLY_PROFILING=" + ("ON" if args.enable_wasm_profiling else "OFF"), "-Donnxruntime_ENABLE_EAGER_MODE=" + ("ON" if args.build_eager_mode else "OFF"), + "-Donnxruntime_ENABLE_ON_DEVICE_TRAINING=" + ("ON" if args.build_on_device_training else "OFF"), "-Donnxruntime_ENABLE_EXTERNAL_CUSTOM_OP_SCHEMAS=" + ("ON" if args.enable_external_custom_op_schemas else "OFF"), "-Donnxruntime_NVCC_THREADS=" + str(args.parallel),