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 54 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 36 of file train_with_dataset.py.
Referenced by train().
train_with_dataset.initParse | ( | ) |
Definition at line 5 of file train_with_dataset.py.
train_with_dataset.train | ( | args | ) |
Main function to train the model.
Definition at line 71 of file train_with_dataset.py.
References createLoss(), and createOptimizer().
train_with_dataset.args = initParse() |
Definition at line 123 of file train_with_dataset.py.