Implement Custom STL Algorithm
This section guides you through implementing custom single-task learning (STL) algorithms for use in CLArena.
Single-task learning algorithms define the process of learning a single task.
Base Classes
In CLArena, single-task learning algorithms are implemented as subclasses of the base classes in clarena/stl_algorithms/base.py. The base classes are implemented inheriting Lightning module with additional features for single-task learning:
clarena.stl_algorithms.STLAlgorithm: the base class for all single-task learning algorithms.
Implement STL Algorithm
To implement STL algorithms:
- Inherit
STLAlgorithm. - Put your algorithm’s hyperparameters in
save_hyperparameters()in__init__(). This enables Lightning loggers to manage them automatically. - Implement
training_step(),validation_step(),test_step()to define the training, validation and test steps respectively. - Other hooks like
on_train_start(),on_validation_start(),on_test_start()are also free to customize if needed.
STLAlgorithm works the same as Lightning module. All the other hooks are free to customize. Please refer to the Lightning module documentation for details about the hooks.
Note that configure_optimizers() is already implemented in STLAlgorithm to manage the optimizer and learning rate scheduler for each task automatically. You don’t have to implement it unless you want to override the default behaviour.
For more details, please refer to the API Reference and source code. You may take implemented STL algorithms in CLArena as examples. Feel free to contribute by submitting pull requests in GitHub!