onnxruntime_training_cxx_api.h

// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
 
#pragma once
#include "onnxruntime_training_c_api.h"
#include <optional>
#include <variant>
 
namespace Ort::detail {
 
#define ORT_DECLARE_TRAINING_RELEASE(NAME) \
  void OrtRelease(Ort##NAME* ptr);
 
// These release methods must be forward declared before including onnxruntime_cxx_api.h
// otherwise class Base won't be aware of them
ORT_DECLARE_TRAINING_RELEASE(CheckpointState);
ORT_DECLARE_TRAINING_RELEASE(TrainingSession);
 
}  // namespace Ort::detail
 
#include "onnxruntime_cxx_api.h"
 
namespace Ort {
 
inline const OrtTrainingApi& GetTrainingApi() { return *GetApi().GetTrainingApi(ORT_API_VERSION); }
 
namespace detail {
 
#define ORT_DEFINE_TRAINING_RELEASE(NAME) \
  inline void OrtRelease(Ort##NAME* ptr) { GetTrainingApi().Release##NAME(ptr); }
 
ORT_DEFINE_TRAINING_RELEASE(CheckpointState);
ORT_DEFINE_TRAINING_RELEASE(TrainingSession);
 
#undef ORT_DECLARE_TRAINING_RELEASE
#undef ORT_DEFINE_TRAINING_RELEASE
 
}  // namespace detail
 
using Property = std::variant<int64_t, float, std::string>;
 
class CheckpointState : public detail::Base<OrtCheckpointState> {
 private:
  CheckpointState(OrtCheckpointState* checkpoint_state) { p_ = checkpoint_state; }
 
 public:
  // Construct the checkpoint state by loading the checkpoint by calling LoadCheckpoint
  CheckpointState() = delete;
 
 
  static CheckpointState LoadCheckpoint(const std::basic_string<ORTCHAR_T>& path_to_checkpoint);
 
  static CheckpointState LoadCheckpointFromBuffer(const std::vector<uint8_t>& buffer);
 
  static void SaveCheckpoint(const CheckpointState& checkpoint_state,
                             const std::basic_string<ORTCHAR_T>& path_to_checkpoint,
                             const bool include_optimizer_state = false);
 
  void AddProperty(const std::string& property_name, const Property& property_value);
 
  Property GetProperty(const std::string& property_name);
 
  void UpdateParameter(const std::string& parameter_name, const Value& parameter);
 
  Value GetParameter(const std::string& parameter_name);
 
};
 
class TrainingSession : public detail::Base<OrtTrainingSession> {
 private:
  size_t training_model_output_count_, eval_model_output_count_;
 
 public:
 
  TrainingSession(const Env& env, const SessionOptions& session_options, CheckpointState& checkpoint_state,
                  const std::basic_string<ORTCHAR_T>& train_model_path,
                  const std::optional<std::basic_string<ORTCHAR_T>>& eval_model_path = std::nullopt,
                  const std::optional<std::basic_string<ORTCHAR_T>>& optimizer_model_path = std::nullopt);
 
  TrainingSession(const Env& env, const SessionOptions& session_options, CheckpointState& checkpoint_state,
                  const std::vector<uint8_t>& train_model_data, const std::vector<uint8_t>& eval_model_data = {},
                  const std::vector<uint8_t>& optim_model_data = {});
 
 
  std::vector<Value> TrainStep(const std::vector<Value>& input_values);
 
  void LazyResetGrad();
 
  std::vector<Value> EvalStep(const std::vector<Value>& input_values);
 
  void SetLearningRate(float learning_rate);
 
  float GetLearningRate() const;
 
  void RegisterLinearLRScheduler(int64_t warmup_step_count, int64_t total_step_count,
                                 float initial_lr);
 
  void SchedulerStep();
 
  void OptimizerStep();
 
 
 
  void ExportModelForInferencing(const std::basic_string<ORTCHAR_T>& inference_model_path,
                                 const std::vector<std::string>& graph_output_names);
 
 
 
  std::vector<std::string> InputNames(const bool training);
 
  std::vector<std::string> OutputNames(const bool training);
 
 
 
  Value ToBuffer(const bool only_trainable);
 
  void FromBuffer(Value& buffer);
 
};
 
 
void SetSeed(const int64_t seed);
 
 
}  // namespace Ort
 
#include "onnxruntime_training_cxx_inline.h"