File size: 684 Bytes
b9fc164 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 | import torch
def load_optimizer(optimizer_name, model, cfg):
trainable_parameters = [param for param in model.parameters() if param.requires_grad]
if not trainable_parameters:
raise ValueError("No trainable parameters found when building optimizer.")
if optimizer_name == "adamw":
return torch.optim.AdamW(
trainable_parameters,
lr=cfg.learning_rate,
betas=(cfg.adam_beta1, cfg.adam_beta2),
weight_decay=cfg.adam_weight_decay,
eps=cfg.adam_epsilon,
)
raise ValueError(
f"Unknown optimizer name: {optimizer_name}, must be 'adamw'"
)
__all__ = [
"load_optimizer",
]
|