| 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", | |
| ] | |