UnfreezeBackbone

class lightning_pose.callbacks.UnfreezeBackbone

Bases: Callback

Callback that ramps up the backbone learning rate from 0 to upsampling_lr on unfreeze_epoch or unfreeze_step.

Starts LR at initial_ratio * upsampling_lr. Grows lr by a factor of warm_up_ratio per epoch or step. Once LR reaches upsampling_lr, keeps it in sync with upsampling_lr.

Use instead of pl.callbacks.BackboneFinetuning in order to use multi-GPU (DDP). See lightning-ai/pytorch-lightning#20340 for context.

Methods Summary

on_train_batch_start(trainer, pl_module, ...)

Adjust the backbone learning rate at the start of each training batch.

Methods Documentation

on_train_batch_start(trainer: Trainer, pl_module: LightningModule, batch: Any, batch_idx: int) → None

Adjust the backbone learning rate at the start of each training batch.

__init__(unfreeze_epoch: int | None = None, unfreeze_step: int | None = None, initial_ratio: float = 0.1, warm_up_ratio: float = 1.5) → None

Initialize UnfreezeBackbone callback.

Exactly one of unfreeze_epoch or unfreeze_step must be provided.

Parameters:

• unfreeze_epoch – epoch at which to begin unfreezing the backbone.

• unfreeze_step – global step at which to begin unfreezing the backbone.

• initial_ratio – backbone LR starts at initial_ratio * upsampling_lr.

• warm_up_ratio – backbone LR is multiplied by this factor each epoch/step during warm-up.

__new__(**kwargs)