Shortcuts

Source code for pytorch_lightning.callbacks.rich_model_summary

# Copyright The PyTorch Lightning 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
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# 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

if _RICH_AVAILABLE:
    from rich.console import Console
    from rich.table import Table


[docs]class RichModelSummary(ModelSummary): r""" Generates a summary of all layers in a :class:`~pytorch_lightning.core.lightning.LightningModule` with `rich text formatting <https://github.com/willmcgugan/rich>`_. 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( "`RichProgressBar` 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 = 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-2021, William Falcon et al... Revision 7792b779.

Built with Sphinx using a theme provided by Read the Docs.
Read the Docs v: latest
Versions
latest
stable
1.5.4
1.5.3
1.5.2
1.5.1
1.5.0
1.4.9
1.4.8
1.4.7
1.4.6
1.4.5
1.4.4
1.4.3
1.4.2
1.4.1
1.4.0
1.3.8
1.3.7
1.3.6
1.3.5
1.3.4
1.3.3
1.3.2
1.3.1
1.3.0
1.2.10
1.2.8
1.2.7
1.2.6
1.2.5
1.2.4
1.2.3
1.2.2
1.2.1
1.2.0
1.1.8
1.1.7
1.1.6
1.1.5
1.1.4
1.1.3
1.1.2
1.1.1
1.1.0
1.0.8
1.0.7
1.0.6
1.0.5
1.0.4
1.0.3
1.0.2
1.0.1
1.0.0
0.10.0
0.9.0
0.8.5
0.8.4
0.8.3
0.8.2
0.8.1
0.8.0
0.7.6
0.7.5
0.7.4
0.7.3
0.7.2
0.7.1
0.7.0
0.6.0
0.5.3
0.4.9
ipynb-update
docs-search
Downloads
html
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.