mirror of
https://github.com/NousResearch/atropos.git
synced 2026-04-25 17:10:42 +00:00
Convert FOB submodule to regular folder
This commit is contained in:
parent
94f046ad40
commit
94825011a0
74 changed files with 4563 additions and 0 deletions
37
environments/optimizer/FOB/pytorch_fob/optimizers/README.md
Normal file
37
environments/optimizer/FOB/pytorch_fob/optimizers/README.md
Normal file
|
|
@ -0,0 +1,37 @@
|
|||
# Optimizers
|
||||
We currently have the following optimizers:
|
||||
|
||||
| Name | Optimizer | LR Scheduler |
|
||||
| ---- | --------- | ------------ |
|
||||
| adamw_baseline | [AdamW](https://arxiv.org/abs/1711.05101) | [Cosine Annealing](https://arxiv.org/abs/1608.03983) with linear warmup |
|
||||
| adamcpr | [AdamCPR](https://arxiv.org/abs/2311.09058v2) | [Cosine Annealing](https://arxiv.org/abs/1608.03983) with linear warmup |
|
||||
| sgd_baseline | Stochastic Gradient Descent | [Cosine Annealing](https://arxiv.org/abs/1608.03983) |
|
||||
| sgd_stepwise | [Stochastic Gradient Descent](https://pytorch.org/docs/stable/generated/torch.optim.SGD.html#torch.optim.SGD) | [StepLR](https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.StepLR.html) |
|
||||
| adafactor | [Adafactor](https://arxiv.org/abs/1804.04235) | Constant |
|
||||
|
||||
|
||||
## Creating your own optimizer
|
||||
To add your own optimizer, you need to create a subfolder in the `optimizers` directory. The name of that folder will be the name used to invoke the optimizer. Within the folder you need to provide two files: `optimizer.py` and `default.yaml`. There is a `template` optimizer with useful comments, which can be used as a starting point.
|
||||
|
||||
### optimizer.py
|
||||
Here you need to implement a function `configure_optimizers` with the following signature:
|
||||
```python
|
||||
configure_optimizers(model: GroupedModel, config: OptimizerConfig) -> OptimizerLRScheduler
|
||||
```
|
||||
- The return type is the same as described [here](https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.core.LightningModule.html#lightning.pytorch.core.LightningModule.configure_optimizers).
|
||||
- The `GroupedModel` is a wrapper around a `torch.nn.Module`. It additionally provides a method `grouped_parameters`, which returns the model parameters grouped by their `weight_decay` and `learning_rate` settings. This is useful for some tasks that want to use e.g. lower learning rates for different parts of the model or to avoid applying weight decay to your norm layers. The underlying `torch.nn.Module` can be accessed with `model.model`.
|
||||
- The `OptimizerConfig` has the `lr_interval, max_steps, max_epochs` attributes. It also gains all attributes provided in the `optimizer` section of the `experiment.yaml`.
|
||||
|
||||
### default.yaml
|
||||
Here you can provide default values for all the hyperparameters your optimizer needs. These values will be added to the `OptimizerConfig` passed to the `configure_optimizers`. So if you have the following `default.yaml`:
|
||||
```yaml
|
||||
optimizer:
|
||||
name: my_awesome_optimizer
|
||||
output_dir_name: my_awesome_optimizer
|
||||
learning_rate: 1.e-3
|
||||
important:
|
||||
extra:
|
||||
parameter: 42
|
||||
```
|
||||
you could use `config.important.extra.parameter` in the `configure_optimizers` function.
|
||||
|
||||
|
|
@ -0,0 +1,4 @@
|
|||
from .optimizers import optimizer_names, optimizer_path, Optimizer
|
||||
from .lr_schedulers import lr_schedulers_path
|
||||
|
||||
__all__ = ["Optimizer", "optimizer_names", "optimizer_path", "lr_schedulers_path"]
|
||||
|
|
@ -0,0 +1,27 @@
|
|||
import importlib
|
||||
from pathlib import Path
|
||||
from lightning.pytorch.utilities.types import OptimizerLRScheduler
|
||||
from pytorch_fob.engine.parameter_groups import GroupedModel
|
||||
from pytorch_fob.engine.configs import OptimizerConfig
|
||||
|
||||
|
||||
def import_optimizer(name: str):
|
||||
return importlib.import_module(f"pytorch_fob.optimizers.{name}.optimizer")
|
||||
|
||||
|
||||
def optimizer_path(name: str) -> Path:
|
||||
return Path(__file__).resolve().parent / name
|
||||
|
||||
|
||||
def optimizer_names() -> list[str]:
|
||||
EXCLUDE = ["__pycache__", "lr_schedulers"]
|
||||
return [d.name for d in Path(__file__).parent.iterdir() if d.is_dir() and d.name not in EXCLUDE]
|
||||
|
||||
|
||||
class Optimizer():
|
||||
def __init__(self, config: OptimizerConfig) -> None:
|
||||
self.config = config
|
||||
|
||||
def configure_optimizers(self, model: GroupedModel) -> OptimizerLRScheduler:
|
||||
optimizer_module = import_optimizer(self.config.name)
|
||||
return optimizer_module.configure_optimizers(model, self.config)
|
||||
Loading…
Add table
Add a link
Reference in a new issue