Spaces:
Running on Zero
Running on Zero
File size: 1,248 Bytes
0122a25 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 | """Callback to configure learning rate during training."""
from __future__ import annotations
from collections.abc import Iterable
from typing import Any
import lightning.pytorch as pl
from mapdet3d.engine.optim.scheduler import LRSchedulerWrapper
from .base import Callback
class LRSchedulerCallback(Callback):
"""Callback to configure learning rate during training."""
def __init__(self) -> None:
"""Initialize the callback."""
super().__init__()
self.last_step = 0
def on_train_batch_end( # type: ignore
self,
trainer: pl.Trainer,
pl_module: pl.LightningModule,
outputs: Any,
batch: Any,
batch_idx: int,
) -> None:
"""Hook on training batch end."""
schedulers = pl_module.lr_schedulers()
if not isinstance(schedulers, Iterable):
schedulers = [schedulers] # type: ignore
if trainer.global_step != self.last_step:
for scheduler in schedulers:
if scheduler is None:
continue
assert isinstance(scheduler, LRSchedulerWrapper)
scheduler.step_on_batch(trainer.global_step)
self.last_step = trainer.global_step
|