ModuleSummary

class stable_pretraining.callbacks.ModuleSummary

Bases: Callback

Logs detailed module parameter statistics in a formatted table.

This callback provides a comprehensive overview of all modules in the model, showing the number of trainable, non-trainable, uninitialized parameters, and buffers for each module. This helps understand model architecture and parameter distribution.

The summary is displayed during the setup phase and includes: - Module name and hierarchy - Trainable parameter count - Non-trainable (frozen) parameter count - Uninitialized parameter count (for lazy modules) - Buffer count (non-parameter persistent state)

setup(trainer, pl_module, stage)

Called when fit, validate, test, predict, or tune begins.