mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-11 17:48:34 +00:00
Refine a bit (on device training) (#15803)
### Few minor refinements: - Simplify ParameterOptimizerState a bit - Use inlined containers - Remove GetStateDict APIs] - Re-enable cuda test for lr scheduler
This commit is contained in:
pengwa
GitHub
parent
346ec12377
commit
fed52053a7
9 changed files with 63 additions and 120 deletions
|
|
@ -191,8 +191,8 @@ TEST(CheckpointApiTest, LoadCheckpointToModel) {
|
|||
}
|
||||
|
||||
/**
|
||||
* Create Module with sets of parameters,
|
||||
* Create Optimizer passing in Module's parameters.
|
||||
* Create Module passing in checkpoint state,
|
||||
* Create Optimizer passing in checkpoint state.
|
||||
* Save Optimizer states into ORT checkpoint files,
|
||||
* Then load it into ORT, compare with the initial optimizer states values.
|
||||
*/
|
||||
|
|
@ -206,7 +206,8 @@ TEST(CheckpointApiTest, SaveOptimizerStateAsCheckpoint_ThenLoad_CUDA) {
|
|||
auto optim_uri = "testdata/training_api/adamw.onnx";
|
||||
|
||||
// Generate randomized weight values using synthetic data generator.
|
||||
constexpr int64_t fc2_weight_dim_in = 10, fc2_weight_dim_out = 500, fc1_weight_dim_in = 500, fc1_weight_dim_out = 784;
|
||||
constexpr int64_t fc2_weight_dim_in = 10, fc2_weight_dim_out = 500,
|
||||
fc1_weight_dim_in = 500, fc1_weight_dim_out = 784;
|
||||
const std::vector<int64_t> fc1_weight_shape{fc1_weight_dim_in, fc1_weight_dim_out};
|
||||
const std::vector<int64_t> fc1_bias_shape{fc1_weight_dim_in};
|
||||
const std::vector<int64_t> fc2_weight_shape{fc2_weight_dim_in, fc2_weight_dim_out};
|
||||
|
|
@ -249,12 +250,6 @@ TEST(CheckpointApiTest, SaveOptimizerStateAsCheckpoint_ThenLoad_CUDA) {
|
|||
auto optimizer = std::make_unique<Optimizer>(optim_uri, &state, session_option,
|
||||
*env, cuda_provider);
|
||||
|
||||
/// Phase 2 - Run Optimizer.GetStateDict and call save checkpoint APIs.
|
||||
/// And check the result checkpoint files.
|
||||
|
||||
CheckpointState checkpoint_state;
|
||||
ORT_ENFORCE(optimizer->GetStateDict(checkpoint_state.optimizer_checkpoint_state).IsOK());
|
||||
|
||||
// Remove the temporary directory if it already exists.
|
||||
auto ckpt_test_root_dir = ORT_TSTR("checkpointing_api_test_dir");
|
||||
if (Env::Default().FolderExists(ckpt_test_root_dir)) {
|
||||
|
|
@ -265,13 +260,14 @@ TEST(CheckpointApiTest, SaveOptimizerStateAsCheckpoint_ThenLoad_CUDA) {
|
|||
// Call Save APIs.
|
||||
PathString checkpoint_path{
|
||||
ConcatPathComponent<PathChar>(tmp_dir.Path(), ORT_TSTR("e2e_ckpt_save_cpu"))};
|
||||
ASSERT_STATUS_OK(SaveCheckpoint(checkpoint_state, checkpoint_path, true));
|
||||
ASSERT_STATUS_OK(SaveCheckpoint(state, checkpoint_path, true));
|
||||
|
||||
// Check the ckpt files in the directory.
|
||||
std::set<PathString> expected_file_names{
|
||||
ORT_TSTR("optim_group0_momentum0_tensors.pbseq"),
|
||||
ORT_TSTR("optim_group0_momentum1_tensors.pbseq"),
|
||||
ORT_TSTR("optim_group0_properties.pbseq"),
|
||||
ORT_TSTR("paramtrain_tensors.pbseq"),
|
||||
};
|
||||
|
||||
std::set<PathString> valid_file_names;
|
||||
|
|
@ -289,27 +285,27 @@ TEST(CheckpointApiTest, SaveOptimizerStateAsCheckpoint_ThenLoad_CUDA) {
|
|||
|
||||
ASSERT_EQ(expected_file_names, valid_file_names);
|
||||
|
||||
/// Phase 3 - Run load checkpoint APIs.
|
||||
/// Phase 2 - Run load checkpoint APIs.
|
||||
/// Validate the result matches with initial optimizer state values.
|
||||
|
||||
// Call Load APIs
|
||||
CheckpointState checkpoint_state_to_load;
|
||||
ASSERT_STATUS_OK(LoadCheckpoint(checkpoint_path, checkpoint_state_to_load));
|
||||
OptimizerCheckpointState optimizer_state = checkpoint_state_to_load.optimizer_checkpoint_state;
|
||||
std::unordered_map<std::string, std::shared_ptr<GroupOptimizerState>>&
|
||||
InlinedHashMap<std::string, std::shared_ptr<GroupOptimizerState>>&
|
||||
group_optimizer_states = optimizer_state.group_named_optimizer_states;
|
||||
|
||||
ASSERT_EQ(group_optimizer_states.size(), 1);
|
||||
ASSERT_EQ(group_optimizer_states.begin()->first, "group0");
|
||||
|
||||
std::unordered_map<std::string, ParameterOptimizerState>&
|
||||
InlinedHashMap<std::string, ParameterOptimizerState>&
|
||||
param_named_optimizer_states = group_optimizer_states["group0"]->param_named_optimizer_states;
|
||||
|
||||
ASSERT_EQ(param_named_optimizer_states.size(), named_parameters.size());
|
||||
|
||||
for (auto it = param_named_optimizer_states.begin(); it != param_named_optimizer_states.end(); ++it) {
|
||||
ASSERT_TRUE(named_parameters.find(it->first) != named_parameters.end());
|
||||
for (auto& [momentum_name, restored_ort_value] : it->second.momentum_named_states) {
|
||||
for (auto& [momentum_name, restored_ort_value] : it->second) {
|
||||
ASSERT_TRUE(momentum_name == "momentum0" || momentum_name == "momentum1");
|
||||
const OrtValue& param_ort_value = name_to_ort_value[it->first];
|
||||
ASSERT_TRUE(restored_ort_value.IsTensor() && param_ort_value.IsTensor());
|
||||
|
|
|
|||
|
|
@ -42,7 +42,7 @@ constexpr float INITIAL_LR = 1e-3f;
|
|||
*/
|
||||
Status CreateFakeOptimizerCheckpointStateOnCPU(
|
||||
const std::unordered_map<std::string, std::shared_ptr<Parameter>>& named_parameters,
|
||||
const std::vector<std::string>& momentum_keys,
|
||||
const InlinedVector<std::string>& momentum_keys,
|
||||
OptimizerCheckpointState& optimizer_checkpoint_state) {
|
||||
auto& grouped_optimizer_states = optimizer_checkpoint_state.group_named_optimizer_states;
|
||||
grouped_optimizer_states.insert({"group0", std::make_shared<GroupOptimizerState>()});
|
||||
|
|
@ -58,7 +58,7 @@ Status CreateFakeOptimizerCheckpointStateOnCPU(
|
|||
OrtValue param = pair.second->Data();
|
||||
const auto& param_tensor = param.template Get<Tensor>();
|
||||
GenerateRandomInput(param_tensor.Shape().GetDims(), param_moment_state);
|
||||
cur_param_optimizer_states.momentum_named_states.insert({state_name, std::move(param_moment_state)});
|
||||
cur_param_optimizer_states.insert({state_name, std::move(param_moment_state)});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -287,39 +287,7 @@ TEST(TrainingApiTest, OptimizerCreatedWithOptimizerCheckpointState) {
|
|||
std::shared_ptr<Optimizer> optim = std::make_shared<Optimizer>(
|
||||
ToUTF8String(optim_uri), &new_state, session_option, *env, providers);
|
||||
|
||||
// After loading state dict, check if optim state is updated to new states.
|
||||
OptimizerCheckpointState optimizer_states;
|
||||
ASSERT_STATUS_OK(optim->GetStateDict(optimizer_states));
|
||||
|
||||
for (auto& p : model->NamedParameters()) {
|
||||
auto param_name = p.first;
|
||||
ParameterOptimizerState& param_state =
|
||||
optimizer_states.group_named_optimizer_states["group0"]->param_named_optimizer_states.at(param_name);
|
||||
|
||||
ParameterOptimizerState& external_param_state =
|
||||
external_optimizer_checkpoint_state.group_named_optimizer_states["group0"]
|
||||
->param_named_optimizer_states.at(param_name);
|
||||
for (auto& param_p : param_state.momentum_named_states) {
|
||||
std::vector<float> moment_vec;
|
||||
if (run_cuda) {
|
||||
CudaOrtValueToCpuVec(param_state.momentum_named_states.at(param_p.first), moment_vec);
|
||||
} else {
|
||||
CpuOrtValueToVec(param_state.momentum_named_states.at(param_p.first), moment_vec);
|
||||
}
|
||||
std::vector<float> external_moment_vect;
|
||||
|
||||
if (run_cuda) {
|
||||
CudaOrtValueToCpuVec(external_param_state.momentum_named_states.at(param_p.first), external_moment_vect);
|
||||
} else {
|
||||
CpuOrtValueToVec(external_param_state.momentum_named_states.at(param_p.first), external_moment_vect);
|
||||
}
|
||||
|
||||
ASSERT_EQ(moment_vec.size(), external_moment_vect.size());
|
||||
for (size_t i = 0; i < moment_vec.size(); i++) {
|
||||
ASSERT_EQ(moment_vec[i], external_moment_vect[i]);
|
||||
}
|
||||
}
|
||||
}
|
||||
ASSERT_TRUE(optim.get() != nullptr);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -328,9 +296,9 @@ void TestLRSchduler(const std::basic_string<ORTCHAR_T>& test_file_name,
|
|||
int64_t total_step_count,
|
||||
int64_t warmup_step_count) {
|
||||
std::vector<bool> run_cuda_list{false};
|
||||
// #ifdef USE_CUDA
|
||||
// run_cuda_list.push_back(true);
|
||||
// #endif
|
||||
#ifdef USE_CUDA
|
||||
run_cuda_list.push_back(true);
|
||||
#endif
|
||||
|
||||
for (auto run_cuda : run_cuda_list) {
|
||||
std::vector<std::shared_ptr<IExecutionProvider>> providers;
|
||||
|
|
@ -392,8 +360,7 @@ void TestLRSchduler(const std::basic_string<ORTCHAR_T>& test_file_name,
|
|||
optim, warmup_step_count, total_step_count);
|
||||
|
||||
for (auto it = test_data.begin(); it != test_data.end(); ++it) {
|
||||
OptimizerCheckpointState optimizer_states;
|
||||
ASSERT_STATUS_OK(optim->GetStateDict(optimizer_states));
|
||||
OptimizerCheckpointState& optimizer_states = state.optimizer_checkpoint_state;
|
||||
auto group_optimizer_state = optimizer_states.group_named_optimizer_states["group0"];
|
||||
|
||||
constexpr const float rtol = 1e-4f, atol = 1e-5f;
|
||||
|
|
@ -493,11 +460,10 @@ TEST(TrainingApiTest, OptimStep) {
|
|||
std::string param_name = "fc2.weight";
|
||||
|
||||
// before training, check if optim state is initialized to 0
|
||||
onnxruntime::training::api::OptimizerCheckpointState optimizer_states;
|
||||
ASSERT_STATUS_OK(optim->GetStateDict(optimizer_states));
|
||||
onnxruntime::training::api::OptimizerCheckpointState& optimizer_states = state.optimizer_checkpoint_state;
|
||||
onnxruntime::training::api::ParameterOptimizerState& param_state =
|
||||
optimizer_states.group_named_optimizer_states["group0"]->param_named_optimizer_states.at(param_name);
|
||||
OrtValue& moment_1 = param_state.momentum_named_states.at("momentum0");
|
||||
OrtValue& moment_1 = param_state.at("momentum0");
|
||||
|
||||
std::vector<float> param_vec_before_optimizer_step;
|
||||
CudaOrtValueToCpuVec(model->NamedParameters().at(param_name)->Data(), param_vec_before_optimizer_step);
|
||||
|
|
|
|||
|
|
@ -260,7 +260,7 @@ Status OrtSaveOptimizerStatesInternal(OptimizerCheckpointState& optimizer_state,
|
|||
// Firstly indexed by momentum names; Secondly indexed by parameter names.
|
||||
InlinedHashMap<std::string, std::unordered_map<std::string, OrtValue>> optimizer_state_ort_values;
|
||||
for (const auto& [param_name, param_optimizer_state] : group_optimizer_state_ptr->param_named_optimizer_states) {
|
||||
for (const auto& [momentum_name, m_state_val] : param_optimizer_state.momentum_named_states) {
|
||||
for (const auto& [momentum_name, m_state_val] : param_optimizer_state) {
|
||||
if (optimizer_state_ort_values.find(momentum_name) == optimizer_state_ort_values.end()) {
|
||||
std::unordered_map<std::string, OrtValue> param_name_to_ortvalue{{param_name, m_state_val}};
|
||||
optimizer_state_ort_values.insert({momentum_name, param_name_to_ortvalue});
|
||||
|
|
@ -421,7 +421,7 @@ Status OrtLoadOptimizerStatesInternal(const PathString& optimizer_folder_path,
|
|||
}
|
||||
|
||||
auto& group_optimizer_state = grouped_optimizer_states[group_name];
|
||||
std::unordered_map<std::string, ParameterOptimizerState>&
|
||||
InlinedHashMap<std::string, ParameterOptimizerState>&
|
||||
param_optimizer_states = group_optimizer_state->param_named_optimizer_states;
|
||||
|
||||
const PathString& tensor_file_path = GetTensorProtoFilePath(optimizer_folder_path,
|
||||
|
|
@ -437,7 +437,7 @@ Status OrtLoadOptimizerStatesInternal(const PathString& optimizer_folder_path,
|
|||
ParameterOptimizerState param_state;
|
||||
param_optimizer_states.insert({param_name, param_state});
|
||||
}
|
||||
param_optimizer_states[param_name].momentum_named_states.insert({momentum_name, std::move(pair.second)});
|
||||
param_optimizer_states[param_name].insert({momentum_name, std::move(pair.second)});
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -178,12 +178,12 @@ Module::Module(const std::string& train_model_path_or_bytes,
|
|||
state_->module_checkpoint_state.train_session_data_transfer_mgr = &train_sess_->GetDataTransferManager();
|
||||
|
||||
// Extract model input and output names
|
||||
std::vector<std::string> train_input_names, train_output_names;
|
||||
InlinedVector<std::string> train_input_names, train_output_names;
|
||||
utils::GetGraphInputOutputNames(train_sess_, train_input_names, train_output_names);
|
||||
|
||||
// Reorder the extracted input names in the following order:
|
||||
// user inputs, weights, gradients, reset_grad
|
||||
std::vector<std::string> user_input_names, param_input_names, grad_input_names, reset_grad_name;
|
||||
InlinedVector<std::string> user_input_names, param_input_names, grad_input_names, reset_grad_name;
|
||||
|
||||
std::unordered_map<std::string, size_t> param_name_to_grad_input_index_map;
|
||||
for (const auto& input_name : train_input_names) {
|
||||
|
|
@ -283,7 +283,7 @@ Module::Module(const std::string& train_model_path_or_bytes,
|
|||
// We are making certain assumptions: Like the order in which parameters occur will be same between train and eval
|
||||
// graphs, and all the weights present in both graphs match.
|
||||
// TODO: Add the checks instead of making assumptions??
|
||||
std::vector<std::string> eval_user_input_names, eval_param_input_names;
|
||||
InlinedVector<std::string> eval_user_input_names, eval_param_input_names;
|
||||
for (const auto& input_name : eval_input_names_) {
|
||||
if (state_->module_checkpoint_state.named_parameters.find(input_name) !=
|
||||
state_->module_checkpoint_state.named_parameters.end()) {
|
||||
|
|
|
|||
|
|
@ -137,18 +137,22 @@ struct Module {
|
|||
private:
|
||||
std::unique_ptr<onnxruntime::InferenceSession> train_sess_{nullptr};
|
||||
std::unique_ptr<onnxruntime::InferenceSession> eval_sess_{nullptr};
|
||||
std::vector<std::string> train_input_names_;
|
||||
std::vector<std::string> train_output_names_;
|
||||
std::vector<std::string> eval_input_names_;
|
||||
std::vector<std::string> eval_output_names_;
|
||||
std::vector<std::string> weight_names_;
|
||||
std::vector<OrtValue> weights_;
|
||||
std::vector<OrtValue> gradients_;
|
||||
bool accumulate_gradient_ = false;
|
||||
|
||||
InlinedVector<std::string> train_input_names_;
|
||||
InlinedVector<std::string> train_output_names_;
|
||||
InlinedVector<std::string> eval_input_names_;
|
||||
InlinedVector<std::string> eval_output_names_;
|
||||
InlinedVector<std::string> weight_names_;
|
||||
|
||||
InlinedVector<OrtValue> weights_;
|
||||
InlinedVector<OrtValue> gradients_;
|
||||
|
||||
CheckpointState* state_; // Non owning pointer to the state.
|
||||
|
||||
bool accumulate_gradient_ = false;
|
||||
std::string eval_model_path_;
|
||||
size_t train_user_input_count_ = 0U;
|
||||
size_t eval_user_input_count_ = 0U;
|
||||
size_t train_user_input_count_{0U};
|
||||
size_t eval_user_input_count_{0U};
|
||||
};
|
||||
|
||||
} // namespace api
|
||||
|
|
|
|||
|
|
@ -113,7 +113,7 @@ Status Optimizer::GenerateMomentumNamedStates(OptimizerCheckpointState& optimize
|
|||
OrtValue param_state;
|
||||
ORT_ENFORCE(utils::CreateZeroValuedOrtValueLike(optim_sess_state, pair.second->Data(), param_state).IsOK(),
|
||||
"Error generating moment state for ", pair.first);
|
||||
cur_param_optimizer_states.momentum_named_states.insert({state_name, std::move(param_state)});
|
||||
cur_param_optimizer_states.insert({state_name, std::move(param_state)});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -127,8 +127,8 @@ Status Optimizer::ConstructInputs() {
|
|||
|
||||
auto& param_named_optimizer_states = optimizer_state_->param_named_optimizer_states;
|
||||
|
||||
std::vector<Tensor> params, grads;
|
||||
std::vector<std::vector<Tensor>> list_of_momentums;
|
||||
InlinedVector<Tensor> params, grads;
|
||||
InlinedVector<InlinedVector<Tensor>> list_of_momentums;
|
||||
list_of_momentums.resize(optimizer_algo_ptr_->momentum_keys.size());
|
||||
|
||||
// Collect all the non-user-defined inputs from the named_parameters_.
|
||||
|
|
@ -150,7 +150,7 @@ Status Optimizer::ConstructInputs() {
|
|||
for (size_t m_index = 0; m_index < optimizer_algo_ptr_->momentum_keys.size(); ++m_index) {
|
||||
auto* moment_tensor =
|
||||
param_named_optimizer_states.at(parameter_name)
|
||||
.momentum_named_states.at(optimizer_algo_ptr_->momentum_keys[m_index])
|
||||
.at(optimizer_algo_ptr_->momentum_keys[m_index])
|
||||
.GetMutable<Tensor>();
|
||||
list_of_momentums[m_index].emplace_back(
|
||||
Tensor(moment_tensor->DataType(), moment_tensor->Shape(),
|
||||
|
|
@ -258,20 +258,6 @@ Status Optimizer::Step() {
|
|||
return Status::OK();
|
||||
}
|
||||
|
||||
Status Optimizer::GetStateDict(OptimizerCheckpointState& optimizer_checkpoint_state) {
|
||||
auto& grouped_optimizer_states = optimizer_checkpoint_state.group_named_optimizer_states;
|
||||
|
||||
// To support multiple groups, the Optimizer constructor needs to accept information for grouping.
|
||||
grouped_optimizer_states.insert({GROUP_ZERO_NAME, std::make_shared<GroupOptimizerState>(*optimizer_state_)});
|
||||
|
||||
// Pass the optimizer session data transfer manager for data copying when saving.
|
||||
// An alternative is, we can do copy at this stage.
|
||||
ORT_RETURN_IF_NOT(optim_sess_, "optimizer session not initialized");
|
||||
const DataTransferManager& sess_data_transfer_manager = optim_sess_->GetDataTransferManager();
|
||||
optimizer_checkpoint_state.optimizer_session_data_transfer_mgr = &sess_data_transfer_manager;
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
Status Optimizer::LoadStateDict(OptimizerCheckpointState& optimizer_checkpoint_states) {
|
||||
auto group_optimizer_state_it =
|
||||
optimizer_checkpoint_states.group_named_optimizer_states.find(GROUP_ZERO_NAME);
|
||||
|
|
@ -293,8 +279,8 @@ Status Optimizer::LoadStateDict(OptimizerCheckpointState& optimizer_checkpoint_s
|
|||
ORT_ENFORCE(src_exist || !strict_match, "Parameter ", params_iter.first,
|
||||
" not found in the source optimizer checkpoint states.");
|
||||
|
||||
std::unordered_map<std::string, OrtValue>& momentum_named_states =
|
||||
param_named_optimizer_states.at(params_iter.first).momentum_named_states;
|
||||
InlinedHashMap<std::string, OrtValue>& momentum_named_states =
|
||||
param_named_optimizer_states.at(params_iter.first);
|
||||
|
||||
OrtValue& param_data = params_iter.second->Data();
|
||||
ORT_ENFORCE(param_data.IsTensor());
|
||||
|
|
|
|||
|
|
@ -19,9 +19,7 @@ namespace api {
|
|||
* For Adam optimizer, it looks like:
|
||||
* { "moment_0": OrtValue, "moment_1": OrtValue,}.
|
||||
*/
|
||||
struct ParameterOptimizerState {
|
||||
std::unordered_map<std::string, OrtValue> momentum_named_states;
|
||||
};
|
||||
typedef InlinedHashMap<std::string, OrtValue> ParameterOptimizerState;
|
||||
|
||||
/**
|
||||
* @brief States belong to one specific group of trainable Parameters.
|
||||
|
|
@ -33,7 +31,7 @@ struct GroupOptimizerState {
|
|||
// Adaptive learning rate as training proceeds. Be noted, learning_rate can be
|
||||
// restored by lr scheduler from given step and initial_lr, though, we still save/load this in checkpoint.
|
||||
float learning_rate{initial_lr};
|
||||
std::unordered_map<std::string, ParameterOptimizerState> param_named_optimizer_states;
|
||||
InlinedHashMap<std::string, ParameterOptimizerState> param_named_optimizer_states;
|
||||
};
|
||||
|
||||
/**
|
||||
|
|
@ -43,16 +41,16 @@ struct GroupOptimizerState {
|
|||
*/
|
||||
struct OptimizerCheckpointState {
|
||||
public:
|
||||
std::unordered_map<std::string, std::shared_ptr<GroupOptimizerState>> group_named_optimizer_states;
|
||||
InlinedHashMap<std::string, std::shared_ptr<GroupOptimizerState>> group_named_optimizer_states;
|
||||
const DataTransferManager* optimizer_session_data_transfer_mgr;
|
||||
};
|
||||
|
||||
struct OptimizerAlgorithmBase {
|
||||
OptimizerAlgorithmBase(const std::vector<std::string>& momentum_keys,
|
||||
const std::vector<std::string>& optimizer_states_inputs)
|
||||
OptimizerAlgorithmBase(const InlinedVector<std::string>& momentum_keys,
|
||||
const InlinedVector<std::string>& optimizer_states_inputs)
|
||||
: momentum_keys(momentum_keys), optimizer_states_inputs(optimizer_states_inputs) {}
|
||||
std::vector<std::string> momentum_keys;
|
||||
std::vector<std::string> optimizer_states_inputs;
|
||||
InlinedVector<std::string> momentum_keys;
|
||||
InlinedVector<std::string> optimizer_states_inputs;
|
||||
};
|
||||
|
||||
struct AdamWOptimizerAlgorithm : public OptimizerAlgorithmBase {
|
||||
|
|
@ -106,15 +104,6 @@ struct Optimizer {
|
|||
|
||||
Status Step();
|
||||
|
||||
/**
|
||||
* @brief Get the current optimizer state.
|
||||
*
|
||||
* Be noted the returned optimizer_checkpoint_states will hold new references to
|
||||
* original momentum states.
|
||||
* @return Status
|
||||
*/
|
||||
Status GetStateDict(OptimizerCheckpointState& optimizer_checkpoint_states);
|
||||
|
||||
Status SetLearningRate(float lr) {
|
||||
optimizer_state_->learning_rate = lr;
|
||||
return Status::OK();
|
||||
|
|
@ -159,11 +148,13 @@ struct Optimizer {
|
|||
|
||||
std::unique_ptr<OptimizerAlgorithmBase> optimizer_algo_ptr_;
|
||||
std::unique_ptr<onnxruntime::InferenceSession> optim_sess_;
|
||||
CheckpointState* state_; // Non owning pointer to the state.
|
||||
|
||||
CheckpointState* state_; // Non owning pointer to the state
|
||||
std::shared_ptr<GroupOptimizerState> optimizer_state_;
|
||||
std::vector<std::string> input_names_;
|
||||
std::vector<std::string> output_names_;
|
||||
std::vector<OrtValue> inputs_;
|
||||
|
||||
InlinedVector<std::string> input_names_;
|
||||
InlinedVector<std::string> output_names_;
|
||||
InlinedVector<OrtValue> inputs_;
|
||||
|
||||
int32_t group_count_{0};
|
||||
};
|
||||
|
|
|
|||
|
|
@ -19,9 +19,9 @@ namespace utils {
|
|||
const std::vector<std::string> GRAD_SUFFIX{"_grad.accumulation.buffer", "_grad", "_grad.accumulation.out"};
|
||||
|
||||
void GetGraphInputOutputNames(const std::unique_ptr<onnxruntime::InferenceSession>& session_object,
|
||||
std::vector<std::string>& input_names,
|
||||
std::vector<std::string>& output_names) {
|
||||
auto get_names = [](const std::vector<const NodeArg*>* node_args, std::vector<std::string>& names) {
|
||||
InlinedVector<std::string>& input_names,
|
||||
InlinedVector<std::string>& output_names) {
|
||||
auto get_names = [](const std::vector<const NodeArg*>* node_args, InlinedVector<std::string>& names) {
|
||||
ORT_ENFORCE(nullptr != node_args);
|
||||
for (const auto* arg : *node_args) {
|
||||
names.push_back(arg->Name());
|
||||
|
|
|
|||
|
|
@ -14,8 +14,8 @@ namespace utils {
|
|||
|
||||
// Get names of graph inputs and outputs
|
||||
void GetGraphInputOutputNames(const std::unique_ptr<onnxruntime::InferenceSession>& session_object,
|
||||
std::vector<std::string>& input_names,
|
||||
std::vector<std::string>& output_names);
|
||||
InlinedVector<std::string>& input_names,
|
||||
InlinedVector<std::string>& output_names);
|
||||
// Fetch the parameter name from suffix: name = param_name+suffix,
|
||||
// returns True if suffix is present in name else False
|
||||
bool GetParamNameFromSuffix(const std::string& name, const std::string& suffix, std::string& param_name);
|
||||
|
|
|
|||
Loading…
Reference in a new issue