RichModelSummary

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:

pip install rich
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.