RichModelSummary

class lightning.pytorch.callbacks.RichModelSummary(max_depth=1)

Bases: lightning.pytorch.callbacks.model_summary.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.

Raises

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