Source code for

# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.

from __future__ import annotations

import os
from typing import TYPE_CHECKING

from onnxruntime.capi import _pybind_state as C

    from import Module

[docs]class Optimizer: """Class that provides methods to update the model parameters based on the computed gradients. Args: optimizer_uri: The path to the optimizer model. model: The module to be trained. """ def __init__(self, optimizer_uri: str | os.PathLike, module: Module): self._optimizer = C.Optimizer( os.fspath(optimizer_uri), module._state._state, module._device, module._session_options )
[docs] def step(self) -> None: """Updates the model parameters based on the computed gradients. This method updates the model parameters by taking a step in the direction of the computed gradients. The optimizer used depends on the optimizer model provided. """ self._optimizer.optimizer_step()
[docs] def set_learning_rate(self, learning_rate: float) -> None: """Sets the learning rate for the optimizer. Args: learning_rate: The learning rate to be set. """ self._optimizer.set_learning_rate(learning_rate)
[docs] def get_learning_rate(self) -> float: """Gets the current learning rate of the optimizer. Returns: The current learning rate. """ return self._optimizer.get_learning_rate()