|
ONE - On-device Neural Engine
|
Functions | |
| initParse () | |
| createOptimizer (optimizer_type, learning_rate=0.001, **kwargs) | |
| createLoss (loss_type, reduction="mean") | |
| train (args) | |
Variables | |
| args = initParse() | |
| train_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 56 of file train_with_dataset.py.
Referenced by train().
| train_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 38 of file train_with_dataset.py.
Referenced by train().
| train_with_dataset.initParse | ( | ) |
Definition at line 7 of file train_with_dataset.py.
| train_with_dataset.train | ( | args | ) |
Main function to train the model.
Definition at line 73 of file train_with_dataset.py.
References createLoss(), and createOptimizer().
| train_with_dataset.args = initParse() |
Definition at line 125 of file train_with_dataset.py.