Unified API to get any scheduler from its name. Args: name (`str` or `SchedulerType`): The name of the scheduler to use. optimizer (`torch.optim.Optimizer`): The optimizer that will be used during training. step_rules (`str`, *optional*):
(
name: Union[str, SchedulerType],
optimizer: Optimizer,
step_rules: Optional[str] = None,
num_warmup_steps: Optional[int] = None,
num_training_steps: Optional[int] = None,
num_cycles: int = 1,
power: float = 1.0,
last_epoch: int = -1,
)
| 287 | |
| 288 | |
| 289 | def get_scheduler( |
| 290 | name: Union[str, SchedulerType], |
| 291 | optimizer: Optimizer, |
| 292 | step_rules: Optional[str] = None, |
| 293 | num_warmup_steps: Optional[int] = None, |
| 294 | num_training_steps: Optional[int] = None, |
| 295 | num_cycles: int = 1, |
| 296 | power: float = 1.0, |
| 297 | last_epoch: int = -1, |
| 298 | ) -> LambdaLR: |
| 299 | """ |
| 300 | Unified API to get any scheduler from its name. |
| 301 | |
| 302 | Args: |
| 303 | name (`str` or `SchedulerType`): |
| 304 | The name of the scheduler to use. |
| 305 | optimizer (`torch.optim.Optimizer`): |
| 306 | The optimizer that will be used during training. |
| 307 | step_rules (`str`, *optional*): |
| 308 | A string representing the step rules to use. This is only used by the `PIECEWISE_CONSTANT` scheduler. |
| 309 | num_warmup_steps (`int`, *optional*): |
| 310 | The number of warmup steps to do. This is not required by all schedulers (hence the argument being |
| 311 | optional), the function will raise an error if it's unset and the scheduler type requires it. |
| 312 | num_training_steps (`int``, *optional*): |
| 313 | The number of training steps to do. This is not required by all schedulers (hence the argument being |
| 314 | optional), the function will raise an error if it's unset and the scheduler type requires it. |
| 315 | num_cycles (`int`, *optional*): |
| 316 | The number of hard restarts used in `COSINE_WITH_RESTARTS` scheduler. |
| 317 | power (`float`, *optional*, defaults to 1.0): |
| 318 | Power factor. See `POLYNOMIAL` scheduler |
| 319 | last_epoch (`int`, *optional*, defaults to -1): |
| 320 | The index of the last epoch when resuming training. |
| 321 | """ |
| 322 | name = SchedulerType(name) |
| 323 | schedule_func = TYPE_TO_SCHEDULER_FUNCTION[name] |
| 324 | if name == SchedulerType.CONSTANT: |
| 325 | return schedule_func(optimizer, last_epoch=last_epoch) |
| 326 | |
| 327 | if name == SchedulerType.PIECEWISE_CONSTANT: |
| 328 | return schedule_func(optimizer, step_rules=step_rules, last_epoch=last_epoch) |
| 329 | |
| 330 | # All other schedulers require `num_warmup_steps` |
| 331 | if num_warmup_steps is None: |
| 332 | raise ValueError(f"{name} requires `num_warmup_steps`, please provide that argument.") |
| 333 | |
| 334 | if name == SchedulerType.CONSTANT_WITH_WARMUP: |
| 335 | return schedule_func(optimizer, num_warmup_steps=num_warmup_steps, last_epoch=last_epoch) |
| 336 | |
| 337 | # All other schedulers require `num_training_steps` |
| 338 | if num_training_steps is None: |
| 339 | raise ValueError(f"{name} requires `num_training_steps`, please provide that argument.") |
| 340 | |
| 341 | if name == SchedulerType.COSINE_WITH_RESTARTS: |
| 342 | return schedule_func( |
| 343 | optimizer, |
| 344 | num_warmup_steps=num_warmup_steps, |
| 345 | num_training_steps=num_training_steps, |
| 346 | num_cycles=num_cycles, |
no test coverage detected