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

On device training offline composition

Browse source
This commit is contained in:
Baiju Meswani 2022-04-06 00:35:00 +00:00 committed by Aishwarya Bhandare
parent 2cdc2be57e
commit ec3da68247
17 changed files with 1958 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
cmake/CMakeLists.txt
View file

@ -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()

52
cmake/onnxruntime_on_device_training.cmake Normal file
View file

@ -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 $<TARGET_FILE:protobuf::protoc>)
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")

28
orttraining/orttraining/onnxflow/README.md Normal file
View file

@ -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.

26
orttraining/orttraining/onnxflow/csrc/load_parameters.cpp Normal file
View file

@ -0,0 +1,26 @@
#include "load_parameters.h"
#include <fcntl.h>
#include <fstream>
#include <iostream>
#include <google/protobuf/text_format.h>
#include <google/protobuf/io/zero_copy_stream_impl.h>
#include <sstream>
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

9
orttraining/orttraining/onnxflow/csrc/load_parameters.h Normal file
View file

@ -0,0 +1,9 @@
#include "onnxflow.pb.h"
namespace onnxflow {
OnnxFlowParameters load_parameters(const std::string& path_to_file);
} // end namespace onnxflow

543
orttraining/orttraining/onnxflow/csrc/onnxflow.pb.cc Normal file
View file

@ -0,0 +1,543 @@
// Generated by the protocol buffer compiler. DO NOT EDIT!
// source: onnxflow.proto
#include "onnxflow.pb.h"
#include <algorithm>
#include <google/protobuf/io/coded_stream.h>
#include <google/protobuf/extension_set.h>
#include <google/protobuf/wire_format_lite.h>
#include <google/protobuf/io/zero_copy_stream_impl_lite.h>
// @@protoc_insertion_point(includes)
#include <google/protobuf/port_def.inc>
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<OnnxFlowParameter>()._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<std::string>(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<size_t>(reinterpret_cast<char*>(&is_parameter_) -
reinterpret_cast<char*>(&requires_grad_)) + sizeof(is_parameter_));
// @@protoc_insertion_point(copy_constructor:onnxflow.OnnxFlowParameter)
}
void OnnxFlowParameter::SharedCtor() {
::memset(reinterpret_cast<char*>(this) + static_cast<size_t>(
reinterpret_cast<char*>(&data_) - reinterpret_cast<char*>(this)),
0, static_cast<size_t>(reinterpret_cast<char*>(&is_parameter_) -
reinterpret_cast<char*>(&data_)) + sizeof(is_parameter_));
}
OnnxFlowParameter::~OnnxFlowParameter() {
// @@protoc_insertion_point(destructor:onnxflow.OnnxFlowParameter)
if (GetArenaForAllocation() != nullptr) return;
SharedDtor();
_internal_metadata_.Delete<std::string>();
}
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<size_t>(
reinterpret_cast<char*>(&is_parameter_) -
reinterpret_cast<char*>(&requires_grad_)) + sizeof(is_parameter_));
_has_bits_.Clear();
_internal_metadata_.Clear<std::string>();
}
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<std::string>(),
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<std::string>(::PROTOBUF_NAMESPACE_ID::internal::GetEmptyString).data(),
static_cast<int>(_internal_metadata_.unknown_fields<std::string>(::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<std::string>(::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<const OnnxFlowParameter*>(
&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<std::string>(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<char*>(&data_),
reinterpret_cast<char*>(&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<std::string>(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<std::string>();
}
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<std::string>();
}
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<std::string>(),
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<unsigned int>(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<std::string>(::PROTOBUF_NAMESPACE_ID::internal::GetEmptyString).data(),
static_cast<int>(_internal_metadata_.unknown_fields<std::string>(::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<std::string>(::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<const OnnxFlowParameters*>(
&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<std::string>(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 <google/protobuf/port_undef.inc>

602
orttraining/orttraining/onnxflow/csrc/onnxflow.pb.h Normal file
View file

@ -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 <limits>
#include <string>
#include <google/protobuf/port_def.inc>
#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 <google/protobuf/port_undef.inc>
#include <google/protobuf/io/coded_stream.h>
#include <google/protobuf/arena.h>
#include <google/protobuf/arenastring.h>
#include <google/protobuf/generated_message_table_driven.h>
#include <google/protobuf/generated_message_util.h>
#include <google/protobuf/metadata_lite.h>
#include <google/protobuf/message_lite.h>
#include <google/protobuf/repeated_field.h> // IWYU pragma: export
#include <google/protobuf/extension_set.h> // IWYU pragma: export
#include <google/protobuf/any.pb.h>
// @@protoc_insertion_point(includes)
#include <google/protobuf/port_def.inc>
#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<std::string>(::PROTOBUF_NAMESPACE_ID::internal::GetEmptyString);
}
inline std::string* mutable_unknown_fields() {
return _internal_metadata_.mutable_unknown_fields<std::string>();
}
static const OnnxFlowParameter& default_instance() {
return *internal_default_instance();
}
static inline const OnnxFlowParameter* internal_default_instance() {
return reinterpret_cast<const OnnxFlowParameter*>(
&_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<OnnxFlowParameter>(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 <typename T> 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<std::string>(::PROTOBUF_NAMESPACE_ID::internal::GetEmptyString);
}
inline std::string* mutable_unknown_fields() {
return _internal_metadata_.mutable_unknown_fields<std::string>();
}
static const OnnxFlowParameters& default_instance() {
return *internal_default_instance();
}
static inline const OnnxFlowParameters* internal_default_instance() {
return reinterpret_cast<const OnnxFlowParameters*>(
&_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<OnnxFlowParameters>(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 <typename T> 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<const ::PROTOBUF_NAMESPACE_ID::Any&>(
::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 &parameters_;
}
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 <google/protobuf/port_undef.inc>
#endif // GOOGLE_PROTOBUF_INCLUDED_GOOGLE_PROTOBUF_INCLUDED_onnxflow_2eproto

BIN
orttraining/orttraining/onnxflow/models/simple_model.onnx Normal file
View file

Binary file not shown.

17
orttraining/orttraining/onnxflow/onnxflow.proto Normal file
View file

@ -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;

7
orttraining/orttraining/onnxflow/onnxflow/__init__.py Normal file
View file

@ -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())

179
orttraining/orttraining/onnxflow/onnxflow/graph.py Normal file
View file

@ -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

136
orttraining/orttraining/onnxflow/onnxflow/loss.py Normal file
View file

@ -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

129
orttraining/orttraining/onnxflow/onnxflow/onnxflow_pb2.py Normal file
View file

@ -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)

144
orttraining/orttraining/onnxflow/onnxflow/optim.py Normal file
View file

@ -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

32
orttraining/orttraining/onnxflow/sample.m.cpp Normal file
View file

@ -0,0 +1,32 @@
#include "orttraining/onnxflow/csrc/load_parameters.h"
#include <filesystem>
#include <iostream>
#include <onnx/onnx_pb.h>
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;
}

36
orttraining/orttraining/onnxflow/sample.py Normal file
View file

@ -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")

6
tools/ci_build/build.py
View file

@ -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),