diff --git a/freqtrade/freqai/base_models/PyTorchModelTrainer.py b/freqtrade/freqai/base_models/PyTorchModelTrainer.py index 41d26e31a..525d543af 100644 --- a/freqtrade/freqai/base_models/PyTorchModelTrainer.py +++ b/freqtrade/freqai/base_models/PyTorchModelTrainer.py @@ -27,6 +27,27 @@ class PyTorchModelTrainer: ): """ A class for training PyTorch models. + Implements the training loop logic, load/save methods. + + fit method - training loop logic: + - Calculates the predicted output for the batch using the PyTorch model. + - Calculates the loss between the predicted and actual output using a loss function. + - Computes the gradients of the loss with respect to the model's parameters using + backpropagation. + - Updates the model's parameters using an optimizer. + + save method: + called by DataDrawer + - Saving any nn.Module state_dict + - Saving model_meta_data, this dict should contain any additional data that the + user needs to store. e.g class_names for classification models. + + load method: + currently DataDrawer is responsible for the actual loading. + when using continual_learning the DataDrawer will load the dict + (saved by self.save(path)). and this class will populate the necessary + state_dict of the self.model & self.optimizer and self.model_meta_data. + :param model: The PyTorch model to be trained. :param optimizer: The optimizer to use for training. @@ -34,10 +55,11 @@ class PyTorchModelTrainer: :param device: The device to use for training (e.g. 'cpu', 'cuda'). :param batch_size: The size of the batches to use during training. :param max_iters: The number of training iterations to run. - iteration here refers to the number of times we call - self.optimizer.step() . used to calculate n_epochs. + iteration here refers to the number of times we call self.optimizer.step(). + used to calculate n_epochs. :param eval_iters: The number of iterations used to estimate the loss. - :param init_model: A dictionary containing the initial model parameters. + :param init_model: A dictionary containing the initial model/optimizer + state_dict and model_meta_data saved by self.save() method. :param model_meta_data: Additional metadata about the model (optional). """ self.model = model