mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-21 19:18:55 +00:00
fb4707f76d
### Description Add cuda support to the on device training python bindings. ### Motivation and Context Now users can set the execution provider (cpu or cuda) when using python bindings for on device training apis.
36 lines
983 B
Python
36 lines
983 B
Python
# Copyright (c) Microsoft Corporation. All rights reserved.
|
|
# Licensed under the MIT License.
|
|
# optimizer.py
|
|
|
|
from onnxruntime.capi import _pybind_state as C
|
|
|
|
|
|
class Optimizer:
|
|
"""
|
|
Class for running Optimize Step in Training.
|
|
This class is a wrapper of Optimizer Class.
|
|
"""
|
|
|
|
def __init__(self, train_optimizer_uri, model) -> None:
|
|
"""
|
|
Initializes Optimizer with the optimizer onnx and the parameters from the model.
|
|
"""
|
|
self._optimizer = C.Optimizer(train_optimizer_uri, model._model, model._device)
|
|
|
|
def step(self):
|
|
"""
|
|
Run Optimizer Step.
|
|
"""
|
|
self._optimizer.optimizer_step()
|
|
|
|
def set_learning_rate(self, learning_rate: float) -> None:
|
|
"""
|
|
Set Learning Rate.
|
|
"""
|
|
self._optimizer.set_learning_rate(learning_rate)
|
|
|
|
def get_learning_rate(self) -> float:
|
|
"""
|
|
Get Learning Rate.
|
|
"""
|
|
return self._optimizer.get_learning_rate()
|