PatchMasking

class lightning_pose.callbacks.PatchMasking

Bases: Callback

Callback to apply curriculum patch masking during training.

Methods Summary

on_train_batch_start(trainer, pl_module, ...)	Apply patch masking to the batch before it goes to the model.
on_train_epoch_end(trainer, pl_module)	Log curriculum progress.

Methods Documentation

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

Apply patch masking to the batch before it goes to the model.

on_train_epoch_end(trainer: Trainer, pl_module: LightningModule) → None

Log curriculum progress.

__init__(patch_mask_config: dict | None = None, patch_seed: int = 0) → None

Initialize PatchMasking callback.

Parameters:

• patch_mask_config – dictionary configuring the masking curriculum, with optional keys init_step, final_step, init_ratio, and final_ratio.

• patch_seed – seed for reproducible patch selection.

__new__(**kwargs)