MannLabs / alphapeptdeep

Deep learning framework for proteomics
Apache License 2.0
101 stars 19 forks source link

Increase training paradigm flexibility #154

Closed mo-sameh closed 2 months ago

mo-sameh commented 2 months ago

In this pull request, I'm suggesting some minor changes that we used for Alphadia project that we've found beneficial and believe could also be useful for some alphapeptdeep users. Since the training process are inherited directly from the model interface rather than through composition. Consequently, making modifications to the training process using the existing interfaces posed a greater challenge. The focus of this PR is to provide more access to the training paradigm from the model interface without changing the architecture and more importantly ensuring compatibility with existing pipelines .

Here are the proposed changes:

1) Epoch and batch callbacks: These callbacks are triggered after each epoch and batch during training. They pass information such as the epoch/batch number and training loss, enabling monitoring and reporting of the loss between training epochs/batches. However, the key feature of these callbacks is the ability to return a "continue training" boolean value that could be used to stop the training. We utilized this feature to implement early stopping in the Alphadia fine-tuning module.

2) Learning rate scheduler: This change allows users to alter the learning rate scheduler as long as it conforms to the 'LR_SchedulerInterface'. One issue here is that the scheduler is currently initialized within the class. So to maintain backward compatibility the interface takes a class as an argument and initialize it in-ModelInterface, instead of using an object of the LR_scheduler.

Overall, we believe these proposed solutions should not disrupt any existing pipelines but rather enhance the flexibility of the training loops particularly for users who wish to extend the existing training framework. If you have any concerns about these changes or have alternative suggestions for achieving the same goal, please let me know.

mo-sameh commented 2 months ago

I've reviewed all the comments, updated the code, and based on that, I'll close this pull request. I'll then open two separate PRs: one for the Learning Rate Scheduler and another for the Training Callback Handlers.