ONE - On-device Neural Engine
|
Functions | |
initParse () | |
createOptimizer (optimizer_type, learning_rate=0.001, **kwargs) | |
createLoss (loss_type, reduction="mean") | |
train_steps (args) | |
Variables | |
args = initParse() | |
train_step_with_dataset.createLoss | ( | loss_type, | |
reduction = "mean" |
|||
) |
Create a loss function based on the specified type and reduction. Args: loss_type (str): The type of loss function ('mse', 'cce'). reduction (str): Reduction type ('mean', 'sum'). Returns: object: An instance of the specified loss function.
Definition at line 49 of file train_step_with_dataset.py.
Referenced by train_steps().
train_step_with_dataset.createOptimizer | ( | optimizer_type, | |
learning_rate = 0.001 , |
|||
** | kwargs | ||
) |
Create an optimizer based on the specified type. Args: optimizer_type (str): The type of optimizer ('SGD' or 'Adam'). learning_rate (float): The learning rate for the optimizer. **kwargs: Additional parameters for the optimizer. Returns: Optimizer: The created optimizer instance.
Definition at line 31 of file train_step_with_dataset.py.
Referenced by train_steps().
train_step_with_dataset.initParse | ( | ) |
Definition at line 5 of file train_step_with_dataset.py.
train_step_with_dataset.train_steps | ( | args | ) |
Main function to train the model.
Definition at line 66 of file train_step_with_dataset.py.
References createLoss(), and createOptimizer().
train_step_with_dataset.args = initParse() |
Definition at line 135 of file train_step_with_dataset.py.