RichModelSummary

class pytorch_lightning.callbacks.RichModelSummary(max_depth=1)

Bases: pytorch_lightning.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 pytorch_lightning import Trainer
from pytorch_lightning.callbacks import RichModelSummary

trainer = Trainer(callbacks=RichModelSummary())

You could also enable RichModelSummary using the RichProgressBar

from pytorch_lightning import Trainer
from pytorch_lightning.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.