RichModelSummary — PyTorch Lightning 2.5.1.post0 documentation (original) (raw)

class lightning.pytorch.callbacks.RichModelSummary(max_depth=1, **summarize_kwargs)[source]¶

Bases: ModelSummary

Generates a summary of all layers in a LightningModule with rich text formatting.

Install it with pip:

from lightning.pytorch import Trainer from lightning.pytorch.callbacks import RichModelSummary

trainer = Trainer(callbacks=RichModelSummary())

You could also enable RichModelSummary using the RichProgressBar

from lightning.pytorch import Trainer from lightning.pytorch.callbacks import RichProgressBar

trainer = Trainer(callbacks=RichProgressBar())

Parameters:

• max_depth¶ (int) – The maximum depth of layer nesting that the summary will include. A value of 0 turns the layer summary off.
• **summarize_kwargs¶ (Any) – Additional arguments to pass to the summarize method.

Raises:

ModuleNotFoundError – If required rich package is not installed on the device.