Source code for pytorch_lightning.callbacks.rich_model_summary

# Copyright The Lightning AI team.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import List, Tuple

from pytorch_lightning.callbacks import ModelSummary
from pytorch_lightning.utilities.imports import _RICH_AVAILABLE
from pytorch_lightning.utilities.model_summary import get_human_readable_count

    from rich import get_console
    from rich.table import Table

[docs]class RichModelSummary(ModelSummary): r""" Generates a summary of all layers in a :class:`~pytorch_lightning.core.module.LightningModule` with `rich text formatting <>`_. Install it with pip: .. code-block:: bash pip install rich .. code-block:: python from pytorch_lightning import Trainer from pytorch_lightning.callbacks import RichModelSummary trainer = Trainer(callbacks=RichModelSummary()) You could also enable ``RichModelSummary`` using the :class:`~pytorch_lightning.callbacks.RichProgressBar` .. code-block:: python from pytorch_lightning import Trainer from pytorch_lightning.callbacks import RichProgressBar trainer = Trainer(callbacks=RichProgressBar()) Args: max_depth: 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. """ def __init__(self, max_depth: int = 1) -> None: if not _RICH_AVAILABLE: raise ModuleNotFoundError( "`RichModelSummary` requires `rich` to be installed. Install it by running `pip install -U rich`." ) super().__init__(max_depth) @staticmethod def summarize( summary_data: List[Tuple[str, List[str]]], total_parameters: int, trainable_parameters: int, model_size: float, ) -> None: console = get_console() table = Table(header_style="bold magenta") table.add_column(" ", style="dim") table.add_column("Name", justify="left", no_wrap=True) table.add_column("Type") table.add_column("Params", justify="right") column_names = list(zip(*summary_data))[0] for column_name in ["In sizes", "Out sizes"]: if column_name in column_names: table.add_column(column_name, justify="right", style="white") rows = list(zip(*(arr[1] for arr in summary_data))) for row in rows: table.add_row(*row) console.print(table) parameters = [] for param in [trainable_parameters, total_parameters - trainable_parameters, total_parameters, model_size]: parameters.append("{:<{}}".format(get_human_readable_count(int(param)), 10)) grid = Table.grid(expand=True) grid.add_column() grid.add_column() grid.add_row(f"[bold]Trainable params[/]: {parameters[0]}") grid.add_row(f"[bold]Non-trainable params[/]: {parameters[1]}") grid.add_row(f"[bold]Total params[/]: {parameters[2]}") grid.add_row(f"[bold]Total estimated model params size (MB)[/]: {parameters[3]}") console.print(grid)

© Copyright Copyright (c) 2018-2023, Lightning AI et al...

Built with Sphinx using a theme provided by Read the Docs.