├── .github ├── ISSUE_TEMPLATE │ ├── bug_report.md │ └── feature_request.md ├── codecov.yml └── workflows │ └── test.yml ├── .gitignore ├── .pre-commit-config.yaml ├── CITATION.cff ├── LICENSE ├── README.md ├── profiler.py ├── requirements-dev.txt ├── requirements.txt ├── ruff.toml ├── setup.cfg ├── setup.py ├── tests ├── __init__.py ├── colab_gpu_test.ipynb ├── conftest.py ├── exceptions_test.py ├── fixtures │ ├── __init__.py │ ├── genotype.py │ ├── models.py │ └── tmva_net.py ├── gpu_test.py ├── half_precision_test.py ├── jupyter_test.ipynb ├── test_output │ ├── ascii_only.out │ ├── autoencoder.out │ ├── basic_summary.out │ ├── batch_size_optimization.out │ ├── bert.out │ ├── cnn_module_list.out │ ├── compressai.out │ ├── container.out │ ├── custom_parameter.out │ ├── device.out │ ├── device_parallelism.out │ ├── dict_input.out │ ├── dict_parameters_1.out │ ├── dict_parameters_2.out │ ├── dict_parameters_3.out │ ├── edgecase_input_output_model.out │ ├── empty_module.out │ ├── empty_module_list.out │ ├── eval_order_doesnt_matter.out │ ├── fasterrcnn.out │ ├── flan_t5_small.out │ ├── formatting_options.out │ ├── forward_pass_exception.out │ ├── frozen_layers.out │ ├── genotype.out │ ├── google.out │ ├── groups.out │ ├── hide_recursive_layers.out │ ├── hide_recursive_layers_outside_loop.out │ ├── highly_nested_dict_model.out │ ├── input_size_possibilities.out │ ├── input_tensor.out │ ├── jit.out │ ├── linear.out │ ├── linear_model_half.out │ ├── lstm.out │ ├── lstm_custom_batch_size.out │ ├── lstm_half.out │ ├── mixed_trainable_parameters.out │ ├── model_with_args.out │ ├── module_dict.out │ ├── multiple_input_tensor_args.out │ ├── multiple_input_tensor_dict.out │ ├── multiple_input_tensor_list.out │ ├── multiple_input_types.out │ ├── namedtuple.out │ ├── nested_leftover_params.out │ ├── numpy_model.out │ ├── pack_padded.out │ ├── parameter_list.out │ ├── parameters_with_other_layers.out │ ├── partial_jit.out │ ├── pruned_adversary.out │ ├── pruning.out │ ├── recursive.out │ ├── recursive_with_missing_layers.out │ ├── register_parameter.out │ ├── resnet152.out │ ├── resnet18_depth_consistency.out │ ├── resnet50.out │ ├── return_dict.out │ ├── reusing_activation_layers.out │ ├── row_settings.out │ ├── siamese_net.out │ ├── single_input.out │ ├── single_input_all_cols.out │ ├── single_input_batch_dim.out │ ├── single_input_half.out │ ├── single_layer_network_on_gpu.out │ ├── single_layer_network_on_gpu_device.out │ ├── single_linear_layer.out │ ├── single_parameter_model.out │ ├── string_result.out │ ├── tmva_net_column_totals.out │ ├── too_many_linear.out │ ├── too_many_linear_plus_existing_hooks.out │ ├── too_many_relus.out │ ├── trainable_column.out │ ├── training_mode.out │ └── uninitialized_tensor.out ├── torchinfo_test.py └── torchinfo_xl_test.py └── torchinfo ├── __init__.py ├── enums.py ├── formatting.py ├── layer_info.py ├── model_statistics.py ├── py.typed └── torchinfo.py /.github/ISSUE_TEMPLATE/bug_report.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: Bug report 3 | about: Create a report to help us improve 4 | title: '' 5 | labels: '' 6 | assignees: '' 7 | 8 | --- 9 | 10 | **Describe the bug** 11 | A clear and concise description of what the bug is. 12 | 13 | **To Reproduce** 14 | Steps to reproduce the behavior: 15 | 1. Go to '...' 16 | 2. Click on '....' 17 | 3. Scroll down to '....' 18 | 4. See error 19 | 20 | **Expected behavior** 21 | A clear and concise description of what you expected to happen. 22 | 23 | **Screenshots** 24 | If applicable, add screenshots to help explain your problem. 25 | 26 | **Desktop (please complete the following information):** 27 | - OS: [e.g. iOS] 28 | - Browser [e.g. chrome, safari] 29 | - Version [e.g. 22] 30 | 31 | **Smartphone (please complete the following information):** 32 | - Device: [e.g. iPhone6] 33 | - OS: [e.g. iOS8.1] 34 | - Browser [e.g. stock browser, safari] 35 | - Version [e.g. 22] 36 | 37 | **Additional context** 38 | Add any other context about the problem here. 39 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/feature_request.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: Feature request 3 | about: Suggest an idea for this project 4 | title: '' 5 | labels: '' 6 | assignees: '' 7 | 8 | --- 9 | 10 | **Is your feature request related to a problem? Please describe.** 11 | A clear and concise description of what the problem is. Ex. I'm always frustrated when [...] 12 | 13 | **Describe the solution you'd like** 14 | A clear and concise description of what you want to happen. 15 | 16 | **Describe alternatives you've considered** 17 | A clear and concise description of any alternative solutions or features you've considered. 18 | 19 | **Additional context** 20 | Add any other context or screenshots about the feature request here. 21 | -------------------------------------------------------------------------------- /.github/codecov.yml: -------------------------------------------------------------------------------- 1 | coverage: 2 | status: 3 | project: 4 | default: 5 | informational: true 6 | patch: 7 | default: 8 | informational: true -------------------------------------------------------------------------------- /.github/workflows/test.yml: -------------------------------------------------------------------------------- 1 | name: Python package 2 | on: 3 | push: 4 | branches: [ main ] 5 | pull_request: 6 | branches: [ main ] 7 | 8 | jobs: 9 | build: 10 | runs-on: ubuntu-latest 11 | strategy: 12 | fail-fast: false 13 | matrix: 14 | python-version: ["3.8", "3.9", "3.10"] 15 | pytorch-version: ["1.4.0", "1.5.1", "1.6.0", "1.7.1", "1.8", "1.9", "1.10", "1.11", "1.12", "1.13", "2.0", "2.1", "2.2"] 16 | include: 17 | - python-version: 3.11 18 | pytorch-version: 2.0 19 | - python-version: 3.11 20 | pytorch-version: 2.1 21 | - python-version: 3.11 22 | pytorch-version: 2.2 23 | exclude: 24 | - python-version: 3.8 25 | pytorch-version: 1.7.1 26 | - python-version: 3.9 27 | pytorch-version: 1.4.0 28 | - python-version: 3.9 29 | pytorch-version: 1.5.1 30 | - python-version: 3.9 31 | pytorch-version: 1.6.0 32 | - python-version: 3.9 33 | pytorch-version: 1.7.1 34 | - python-version: 3.9 35 | pytorch-version: 1.8 36 | 37 | - python-version: 3.10 38 | pytorch-version: 1.4.0 39 | - python-version: 3.10 40 | pytorch-version: 1.5.1 41 | - python-version: 3.10 42 | pytorch-version: 1.6.0 43 | - python-version: 3.10 44 | pytorch-version: 1.7.1 45 | - python-version: 3.10 46 | pytorch-version: 1.8 47 | - python-version: 3.10 48 | pytorch-version: 1.6.0 49 | - python-version: 3.10 50 | pytorch-version: 1.9 51 | - python-version: 3.10 52 | pytorch-version: 1.10 53 | 54 | steps: 55 | - uses: actions/checkout@v2 56 | - name: Set up Python ${{ matrix.python-version }} 57 | uses: actions/setup-python@v2 58 | with: 59 | python-version: ${{ matrix.python-version }} 60 | - name: Install dependencies 61 | run: | 62 | python -m pip install --upgrade uv 63 | uv pip install --system pytest pytest-cov 64 | uv pip install --system torch==${{ matrix.pytorch-version }} torchvision transformers 65 | - name: mypy 66 | if: ${{ matrix.pytorch-version == '2.2' }} 67 | run: | 68 | uv pip install --system mypy==1.9.0 69 | mypy --install-types --non-interactive . 70 | - name: pytest 71 | if: ${{ matrix.pytorch-version == '2.2' }} 72 | run: | 73 | pytest --cov=torchinfo --cov-report= --durations=0 74 | - name: pytest 75 | if: ${{ matrix.pytorch-version != '2.2' }} 76 | run: | 77 | pytest --no-output -k "not test_eval_order_doesnt_matter and not test_google and not test_uninitialized_tensor and not test_input_size_half_precision and not test_recursive_with_missing_layers and not test_flan_t5_small" 78 | - name: codecov 79 | uses: codecov/codecov-action@v1 80 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | .vscode/ 6 | .DS_Store 7 | profile.txt 8 | Archive.zip 9 | 10 | # C extensions 11 | *.so 12 | 13 | # Distribution / packaging 14 | .Python 15 | build/ 16 | develop-eggs/ 17 | dist/ 18 | downloads/ 19 | eggs/ 20 | .eggs/ 21 | lib/ 22 | lib64/ 23 | parts/ 24 | sdist/ 25 | var/ 26 | wheels/ 27 | pip-wheel-metadata/ 28 | share/python-wheels/ 29 | *.egg-info/ 30 | .installed.cfg 31 | *.egg 32 | MANIFEST 33 | 34 | # PyInstaller 35 | # Usually these files are written by a python script from a template 36 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 37 | *.manifest 38 | *.spec 39 | 40 | # Installer logs 41 | pip-log.txt 42 | pip-delete-this-directory.txt 43 | 44 | # Unit test / coverage reports 45 | htmlcov/ 46 | .tox/ 47 | .nox/ 48 | .coverage 49 | .coverage.* 50 | .cache 51 | nosetests.xml 52 | coverage.xml 53 | *.cover 54 | *.py,cover 55 | .hypothesis/ 56 | .pytest_cache/ 57 | 58 | # Translations 59 | *.mo 60 | *.pot 61 | 62 | # Django stuff: 63 | *.log 64 | local_settings.py 65 | db.sqlite3 66 | db.sqlite3-journal 67 | 68 | # Flask stuff: 69 | instance/ 70 | .webassets-cache 71 | 72 | # Scrapy stuff: 73 | .scrapy 74 | 75 | # Sphinx documentation 76 | docs/_build/ 77 | 78 | # PyBuilder 79 | target/ 80 | 81 | # Jupyter Notebook 82 | .ipynb_checkpoints 83 | 84 | # IPython 85 | profile_default/ 86 | ipython_config.py 87 | 88 | # pyenv 89 | .python-version 90 | 91 | # pipenv 92 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 93 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 94 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 95 | # install all needed dependencies. 96 | #Pipfile.lock 97 | 98 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 99 | __pypackages__/ 100 | 101 | # Celery stuff 102 | celerybeat-schedule 103 | celerybeat.pid 104 | 105 | # SageMath parsed files 106 | *.sage.py 107 | 108 | # Environments 109 | .env 110 | .venv 111 | env/ 112 | venv/ 113 | ENV/ 114 | env.bak/ 115 | venv.bak/ 116 | 117 | # Spyder project settings 118 | .spyderproject 119 | .spyproject 120 | 121 | # Rope project settings 122 | .ropeproject 123 | 124 | # mkdocs documentation 125 | /site 126 | 127 | # mypy 128 | .mypy_cache/ 129 | .dmypy.json 130 | dmypy.json 131 | 132 | # Pyre type checker 133 | .pyre/ 134 | -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | ci: 2 | skip: [mypy, pytest] 3 | repos: 4 | - repo: https://github.com/astral-sh/ruff-pre-commit 5 | rev: v0.9.10 6 | hooks: 7 | - id: ruff 8 | args: [--fix] 9 | - id: ruff-format 10 | 11 | - repo: local 12 | hooks: 13 | - id: mypy 14 | name: mypy 15 | entry: mypy 16 | language: python 17 | types: [python] 18 | require_serial: true 19 | 20 | - id: pytest 21 | name: pytest 22 | entry: pytest --cov=torchinfo --cov-report=html --durations=0 23 | language: python 24 | types: [python] 25 | always_run: true 26 | pass_filenames: false 27 | verbose: true 28 | -------------------------------------------------------------------------------- /CITATION.cff: -------------------------------------------------------------------------------- 1 | cff-version: 1.2.0 2 | title: torchinfo 3 | message: If you use this software, please cite it as below. 4 | type: software 5 | authors: 6 | - given-names: Tyler 7 | family-names: Yep 8 | email: tyep@cs.stanford.edu 9 | identifiers: 10 | - type: url 11 | value: 'https://github.com/TylerYep/torchinfo' 12 | description: View model summaries in PyTorch! 13 | repository-code: 'https://github.com/TylerYep/torchinfo' 14 | abstract: >- 15 | Torchinfo provides information complementary to 16 | what is provided by print(your_model) in PyTorch. 17 | keywords: 18 | - torch 19 | - pytorch 20 | - torchinfo 21 | - torchsummary 22 | license: MIT 23 | date-released: '2020-03-16' 24 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 Tyler Yep 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /profiler.py: -------------------------------------------------------------------------------- 1 | import cProfile 2 | import pstats 3 | import random 4 | 5 | import torchvision # type: ignore[import-untyped] # noqa: F401 6 | from tqdm import trange # noqa: F401 7 | 8 | from torchinfo import summary # noqa: F401 9 | 10 | 11 | def profile() -> None: 12 | """ 13 | Prints top N methods, sorted by time. 14 | Equivalent to: 15 | python -m cProfile -o data/profile.txt main.py -n 100 16 | Options: 17 | time, cumulative, line, name, nfl, calls 18 | ----------- 19 | ncalls - for the number of calls. 20 | 21 | time/tottime - for the total time spent in the given function 22 | (and excluding time made in calls to sub-functions) 23 | 24 | cumulative/cumtime - is the cumulative time spent in this and all subfunctions 25 | (from invocation till exit). This figure is accurate even for recursive functions. 26 | """ 27 | random.seed(0) 28 | command = ( 29 | "for _ in trange(10): " 30 | "summary(torchvision.models.resnet152(), (1, 3, 224, 224), verbose=0)" 31 | ) 32 | profile_file = "profile.txt" 33 | sort = "time" 34 | 35 | cProfile.run(command, filename=profile_file, sort=sort) 36 | stats = pstats.Stats(profile_file) 37 | stats.sort_stats(sort).print_stats(50) 38 | 39 | 40 | if __name__ == "__main__": 41 | profile() 42 | -------------------------------------------------------------------------------- /requirements-dev.txt: -------------------------------------------------------------------------------- 1 | codecov 2 | mypy 3 | pytest 4 | pytest-cov 5 | pre-commit 6 | ruff 7 | transformers 8 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | torch 2 | torchvision 3 | numpy<2 4 | -------------------------------------------------------------------------------- /ruff.toml: -------------------------------------------------------------------------------- 1 | target-version = "py38" 2 | lint.select = ["ALL"] 3 | lint.ignore = [ 4 | "ANN401", # Dynamically typed expressions (typing.Any) are disallowed 5 | "C901", # function is too complex (12 > 10) 6 | "COM812", # Trailing comma missing 7 | "D", # Docstring rules 8 | "EM101", # Exception must not use a string literal, assign to variable first 9 | "EM102", # Exception must not use an f-string literal, assign to variable first 10 | "ERA001", # Found commented-out code 11 | "FBT001", # Boolean positional arg in function definition 12 | "FBT002", # Boolean default value in function definition 13 | "FBT003", # Boolean positional value in function call 14 | "FIX002", # Line contains TODO 15 | "ISC001", # Isort 16 | "PLR0911", # Too many return statements (11 > 6) 17 | "PLR2004", # Magic value used in comparison, consider replacing 2 with a constant variable 18 | "PLR0912", # Too many branches 19 | "PLR0913", # Too many arguments to function call 20 | "PLR0915", # Too many statements 21 | "S101", # Use of `assert` detected 22 | "S311", # Standard pseudo-random generators are not suitable for cryptographic purposes 23 | "T201", # print() found 24 | "T203", # pprint() found 25 | "TD002", # Missing author in TODO; try: `# TODO(): ...` 26 | "TD003", # Missing issue link on the line following this TODO 27 | "TD005", # Missing issue description after `TODO` 28 | "TRY003", # Avoid specifying long messages outside the exception class 29 | 30 | # torchinfo-specific ignores 31 | "N803", # Argument name `A_i` should be lowercase 32 | "N806", # Variable `G` in function should be lowercase 33 | "BLE001", # Do not catch blind exception: `Exception` 34 | "PLW0602", # Using global for `_cached_forward_pass` but no assignment is done 35 | "PLW0603", # Using the global statement to update `_cached_forward_pass` is discouraged 36 | "PLW2901", # `for` loop variable `name` overwritten by assignment target 37 | "SIM108", # [*] Use ternary operator `model_mode = Mode.EVAL if mode is None else Mode(mode)` instead of `if`-`else`-block 38 | "SLF001", # Private member accessed: `_modules` 39 | "TCH002", # Move third-party import into a type-checking block 40 | "TRY004", # Prefer `TypeError` exception for invalid type 41 | "TRY301", # Abstract `raise` to an inner function 42 | ] 43 | exclude = ["tests"] # TODO: check tests too 44 | 45 | [lint.flake8-pytest-style] 46 | fixture-parentheses = false 47 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | [metadata] 2 | name = torchinfo 3 | version = attr: torchinfo.__version__ 4 | description = Model summary in PyTorch, based off of the original torchsummary. 5 | long_description = file: README.md 6 | long_description_content_type = text/markdown 7 | url = https://github.com/tyleryep/torchinfo 8 | author = Tyler Yep @tyleryep 9 | author_email = tyep@cs.stanford.edu 10 | license = MIT 11 | license_file = LICENSE 12 | classifiers = 13 | License :: OSI Approved :: MIT License 14 | Operating System :: OS Independent 15 | Programming Language :: Python :: 3 16 | Programming Language :: Python :: 3 :: Only 17 | Programming Language :: Python :: 3.6 18 | Programming Language :: Python :: 3.7 19 | Programming Language :: Python :: 3.8 20 | Programming Language :: Python :: 3.9 21 | Programming Language :: Python :: 3.10 22 | Programming Language :: Python :: 3.11 23 | Programming Language :: Python :: 3.12 24 | keywords = torch pytorch torchsummary torch-summary summary keras deep-learning ml torchinfo torch-info visualize model statistics layer stats 25 | 26 | [options] 27 | packages = torchinfo 28 | python_requires = >=3.8 29 | include_package_data = True 30 | 31 | [options.package_data] 32 | torchinfo = py.typed 33 | 34 | [mypy] 35 | strict = True 36 | warn_unreachable = True 37 | disallow_any_unimported = True 38 | extra_checks = True 39 | enable_error_code = ignore-without-code 40 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | 3 | setup() 4 | -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TylerYep/torchinfo/e67e74885ce103d7508b71c20e279aee1d9eb4ce/tests/__init__.py -------------------------------------------------------------------------------- /tests/conftest.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import warnings 3 | from pathlib import Path 4 | from typing import Iterator, Tuple 5 | 6 | import pytest 7 | 8 | from torchinfo import ModelStatistics 9 | from torchinfo.enums import ColumnSettings 10 | from torchinfo.formatting import HEADER_TITLES 11 | from torchinfo.torchinfo import clear_cached_forward_pass 12 | 13 | 14 | def pytest_addoption(parser: pytest.Parser) -> None: 15 | """This allows us to check for these params in sys.argv.""" 16 | parser.addoption("--overwrite", action="store_true", default=False) 17 | parser.addoption("--no-output", action="store_true", default=False) 18 | 19 | 20 | @pytest.fixture(autouse=True) 21 | def verify_capsys( 22 | capsys: pytest.CaptureFixture[str], request: pytest.FixtureRequest 23 | ) -> Iterator[None]: 24 | yield 25 | clear_cached_forward_pass() 26 | if "--no-output" in sys.argv: 27 | return 28 | 29 | test_name = request.node.name.replace("test_", "") 30 | if sys.version_info < (3, 8) and test_name == "tmva_net_column_totals": 31 | warnings.warn( 32 | "sys.getsizeof can return different results on earlier Python versions.", 33 | stacklevel=2, 34 | ) 35 | return 36 | 37 | if test_name == "input_size_half_precision": 38 | return 39 | 40 | verify_output(capsys, f"tests/test_output/{test_name}.out") 41 | 42 | 43 | def verify_output(capsys: pytest.CaptureFixture[str], filename: str) -> None: 44 | """ 45 | Utility function to ensure output matches file. 46 | If you are writing new tests, set overwrite_file=True to generate the 47 | new test_output file. 48 | """ 49 | captured, _ = capsys.readouterr() 50 | filepath = Path(filename) 51 | if not captured and not filepath.exists(): 52 | return 53 | if "--overwrite" in sys.argv: 54 | filepath.parent.mkdir(exist_ok=True) 55 | filepath.touch(exist_ok=True) 56 | filepath.write_text(captured, encoding="utf-8") 57 | 58 | verify_output_str(captured, filename) 59 | 60 | 61 | def verify_output_str(output: str, filename: str) -> None: 62 | expected = Path(filename).read_text(encoding="utf-8") 63 | # Verify the input size has the same unit 64 | output_input_size, output_input_unit = get_input_size_and_unit(output) 65 | expected_input_size, expected_input_unit = get_input_size_and_unit(expected) 66 | assert output_input_unit == expected_input_unit 67 | 68 | # Sometime it does not have the same exact value, depending on torch version. 69 | # We assume the variation cannot be too large. 70 | if output_input_size != 0: 71 | assert abs(output_input_size - expected_input_size)/output_input_size < 1e-2 72 | 73 | if output_input_size != expected_input_size: 74 | # In case of a difference, replace the expected input size. 75 | expected = replace_input_size(expected, expected_input_unit, expected_input_size, output_input_size) 76 | assert output == expected 77 | for category in (ColumnSettings.NUM_PARAMS, ColumnSettings.MULT_ADDS): 78 | assert_sum_column_totals_match(output, category) 79 | 80 | def replace_input_size(output: str, unit: str, old_value: float, new_value: float) -> str: 81 | return output.replace(f"Input size {unit}: {old_value:.2f}", f"Input size {unit}: {new_value:.2f}") 82 | 83 | def get_input_size_and_unit(output_str: str) -> Tuple[float, str]: 84 | input_size = float(output_str.split('Input size')[1].split(':')[1].split('\n')[0].strip()) 85 | input_unit = output_str.split('Input size')[1].split(':')[0].strip() 86 | return input_size, input_unit 87 | 88 | def get_column_value_for_row(line: str, offset: int) -> int: 89 | """Helper function for getting the column totals.""" 90 | col_value = line[offset:] 91 | end = col_value.find(" ") 92 | if end != -1: 93 | col_value = col_value[:end] 94 | if ( 95 | not col_value 96 | or col_value in ("--", "(recursive)") 97 | or col_value.startswith(("└─", "├─")) 98 | ): 99 | return 0 100 | return int(col_value.replace(",", "").replace("(", "").replace(")", "")) 101 | 102 | 103 | def assert_sum_column_totals_match(output: str, category: ColumnSettings) -> None: 104 | """Asserts that column totals match the total from the table summary.""" 105 | lines = output.replace("=", "").split("\n\n") 106 | header_row = lines[0].strip() 107 | offset = header_row.find(HEADER_TITLES[category]) 108 | if offset == -1: 109 | return 110 | layers = lines[1].split("\n") 111 | calculated_total = float(sum(get_column_value_for_row(line, offset) for line in layers)) 112 | results = lines[2].split("\n") 113 | 114 | if category == ColumnSettings.NUM_PARAMS: 115 | total_params = results[0].split(":")[1].replace(",", "") 116 | splitted_results = results[0].split('(') 117 | if len(splitted_results) > 1: 118 | units = splitted_results[1][0] 119 | if units == 'T': 120 | calculated_total /= 1e12 121 | elif units == 'G': 122 | calculated_total /= 1e9 123 | elif units == 'M': 124 | calculated_total /= 1e6 125 | elif units == 'k': 126 | calculated_total /= 1e3 127 | assert calculated_total == float(total_params) 128 | elif category == ColumnSettings.MULT_ADDS: 129 | total_mult_adds = results[-1].split(":")[1].replace(",", "") 130 | assert float( 131 | f"{ModelStatistics.to_readable(calculated_total)[1]:0.2f}" 132 | ) == float(total_mult_adds) 133 | -------------------------------------------------------------------------------- /tests/exceptions_test.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import torch 3 | 4 | from tests.fixtures.models import CustomParameter, EdgeCaseModel, IdentityModel 5 | from torchinfo import summary 6 | 7 | 8 | def test_invalid_user_params() -> None: 9 | test = IdentityModel() 10 | 11 | with pytest.raises(ValueError): 12 | summary(test, verbose=4) 13 | with pytest.raises(ValueError): 14 | summary(test, input_size=(1, 28, 28), col_names=("invalid_name",)) 15 | with pytest.raises(ValueError): 16 | summary(test, col_width=0) 17 | with pytest.raises(ValueError): 18 | summary(test, row_settings=("invalid_name",)) 19 | with pytest.raises(ValueError): 20 | summary(test, col_names=("output_size",)) 21 | with pytest.raises(RuntimeError): 22 | summary(test, (1, 28, 28), torch.randn(1, 28, 28)) 23 | with pytest.raises(RuntimeError): 24 | summary(test, device="lol") 25 | with pytest.raises(RuntimeError): 26 | summary(test, device=torch.device("lol")) 27 | 28 | 29 | def test_incorrect_model_forward() -> None: 30 | # Warning: these tests always raise RuntimeError. 31 | with pytest.raises(RuntimeError): 32 | summary(EdgeCaseModel(throw_error=True), input_size=(5, 1, 28, 28)) 33 | with pytest.raises(RuntimeError): 34 | summary(EdgeCaseModel(return_str=True), input_size=(5, 1, 28, 28)) 35 | with pytest.raises(RuntimeError): 36 | summary(EdgeCaseModel(return_class=True), input_size=(5, 1, 28, 28)) 37 | with pytest.raises(RuntimeError): 38 | summary( 39 | EdgeCaseModel(throw_error=True), input_data=[[[torch.randn(1, 28, 28)]]] 40 | ) 41 | 42 | 43 | def test_input_size_possible_exceptions() -> None: 44 | test = CustomParameter(2, 3) 45 | 46 | with pytest.raises(ValueError): 47 | summary(test, input_size=[(3, 0)]) 48 | with pytest.raises(TypeError): 49 | summary(test, input_size={0: 1}) # type: ignore[arg-type] 50 | with pytest.raises(TypeError): 51 | summary(test, input_size="hello") 52 | 53 | 54 | def test_forward_pass_exception() -> None: 55 | input_size = (1, 1, 28, 28) 56 | summary(EdgeCaseModel(throw_error=False), input_size=input_size) 57 | with pytest.raises(RuntimeError): 58 | summary(EdgeCaseModel(throw_error=True), input_size=input_size) 59 | -------------------------------------------------------------------------------- /tests/fixtures/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TylerYep/torchinfo/e67e74885ce103d7508b71c20e279aee1d9eb4ce/tests/fixtures/__init__.py -------------------------------------------------------------------------------- /tests/fixtures/genotype.py: -------------------------------------------------------------------------------- 1 | # type: ignore 2 | from collections import namedtuple 3 | 4 | import torch 5 | from torch import nn 6 | 7 | from tests.fixtures.models import IdentityModel 8 | 9 | 10 | class ReLUConvBN(nn.Module): 11 | def __init__(self, C_in, C_out, kernel_size=1, stride=1, padding=0, affine=True): 12 | super().__init__() 13 | self.op = nn.Sequential( 14 | nn.ReLU(inplace=False), 15 | nn.Conv2d( 16 | C_in, C_out, kernel_size, stride=stride, padding=padding, bias=False 17 | ), 18 | nn.BatchNorm2d(C_out, affine=affine), 19 | ) 20 | 21 | def forward(self, x): 22 | return self.op(x) 23 | 24 | 25 | class SepConv(nn.Module): 26 | def __init__(self, C_in, C_out, stride, kernel_size=3, padding=1, affine=True): 27 | super().__init__() 28 | self.op = nn.Sequential( 29 | nn.ReLU(inplace=False), 30 | nn.Conv2d( 31 | C_in, 32 | C_in, 33 | kernel_size=kernel_size, 34 | stride=stride, 35 | padding=padding, 36 | groups=C_in, 37 | bias=False, 38 | ), 39 | nn.Conv2d(C_in, C_in, kernel_size=1, padding=0, bias=False), 40 | nn.BatchNorm2d(C_in, affine=affine), 41 | nn.ReLU(inplace=False), 42 | nn.Conv2d( 43 | C_in, 44 | C_in, 45 | kernel_size=kernel_size, 46 | stride=1, 47 | padding=padding, 48 | groups=C_in, 49 | bias=False, 50 | ), 51 | nn.Conv2d(C_in, C_out, kernel_size=1, padding=0, bias=False), 52 | nn.BatchNorm2d(C_out, affine=affine), 53 | ) 54 | 55 | def forward(self, x): 56 | return self.op(x) 57 | 58 | 59 | class FactorizedReduce(nn.Module): 60 | def __init__(self, C_in, C_out, affine=True): 61 | super().__init__() 62 | assert C_out % 2 == 0 63 | self.relu = nn.ReLU(inplace=False) 64 | self.conv_1 = nn.Conv2d(C_in, C_out // 2, 1, stride=2, padding=0, bias=False) 65 | self.conv_2 = nn.Conv2d(C_in, C_out // 2, 1, stride=2, padding=0, bias=False) 66 | self.bn = nn.BatchNorm2d(C_out, affine=affine) 67 | 68 | def forward(self, x): 69 | x = self.relu(x) 70 | out = torch.cat([self.conv_1(x), self.conv_2(x[:, :, 1:, 1:])], dim=1) 71 | out = self.bn(out) 72 | return out 73 | 74 | 75 | class Cell(nn.Module): 76 | def __init__(self, C_prev_prev, C_prev, C, reduction=False, reduction_prev=False): 77 | super().__init__() 78 | Genotype = namedtuple("Genotype", "normal normal_concat reduce reduce_concat") 79 | genotype = Genotype( 80 | normal=[("skip_connect", 1), ("skip_connect", 0)], 81 | normal_concat=range(1), 82 | reduce=[("sep_conv_3x3", 0), ("sep_conv_3x3", 1)], 83 | reduce_concat=range(1), 84 | ) 85 | if reduction_prev: 86 | self.preprocess0 = FactorizedReduce(C_prev_prev, C) 87 | else: 88 | self.preprocess0 = ReLUConvBN(C_prev_prev, C) 89 | self.preprocess1 = ReLUConvBN(C_prev, C) 90 | 91 | if reduction: 92 | op_names, indices = zip(*genotype.reduce) 93 | concat = genotype.reduce_concat 94 | else: 95 | op_names, indices = zip(*genotype.normal) 96 | concat = genotype.normal_concat 97 | 98 | self._steps = len(op_names) // 2 99 | self._concat = concat 100 | self.multiplier = len(concat) 101 | 102 | self._ops = nn.ModuleList() 103 | for name, index in zip(op_names, indices): 104 | stride = 2 if reduction and index < 2 else 1 105 | if name == "skip_connect": 106 | op = IdentityModel() if stride == 1 else FactorizedReduce(C, C) 107 | elif name == "sep_conv_3x3": 108 | op = SepConv(C, C, stride) 109 | self._ops += [op] 110 | self._indices = indices 111 | 112 | def forward(self, s0, s1, drop_prob): 113 | s0 = self.preprocess0(s0) 114 | s1 = self.preprocess1(s1) 115 | 116 | states = [s0, s1] 117 | for i in range(self._steps): 118 | h1 = states[self._indices[2 * i]] 119 | h2 = states[self._indices[2 * i + 1]] 120 | op1 = self._ops[2 * i] 121 | op2 = self._ops[2 * i + 1] 122 | h1 = op1(h1) 123 | h2 = op2(h2) 124 | s = h1 + h2 125 | states += [s] 126 | return torch.cat([states[i] for i in self._concat], dim=1) 127 | 128 | 129 | class GenotypeNetwork(nn.Module): 130 | def __init__(self, C=16, num_classes=10, layers=1, auxiliary=False): 131 | super().__init__() 132 | self._layers = layers 133 | self._auxiliary = auxiliary 134 | self.drop_path_prob = 0.0 135 | 136 | stem_multiplier = 3 137 | C_curr = stem_multiplier * C 138 | self.stem = nn.Sequential( 139 | nn.Conv2d(3, C_curr, 3, padding=1, bias=False), nn.BatchNorm2d(C_curr) 140 | ) 141 | 142 | C_prev_prev, C_prev, C_curr = C_curr, C_curr, C 143 | self.cells = nn.ModuleList() 144 | for i in range(layers): 145 | if i in (layers // 3, 2 * layers // 3): 146 | C_curr *= 2 147 | cell = Cell(C_prev_prev, C_prev, C_curr) 148 | self.cells += [cell] 149 | C_prev_prev, C_prev = C_prev, cell.multiplier * C_curr 150 | 151 | def forward(self, input_): 152 | s0 = s1 = self.stem(input_) 153 | for cell in self.cells: 154 | s0, s1 = s1, cell(s0, s1, self.drop_path_prob) 155 | -------------------------------------------------------------------------------- /tests/gpu_test.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import torch 3 | 4 | from tests.fixtures.models import MultiDeviceModel, SingleInputNet 5 | from torchinfo import summary 6 | 7 | 8 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="GPU must be enabled.") 9 | class TestGPU: 10 | """GPU-only tests.""" 11 | 12 | @staticmethod 13 | def test_single_layer_network_on_gpu() -> None: 14 | model = torch.nn.Linear(2, 5) 15 | model.cuda() 16 | 17 | results = summary(model, input_size=(1, 2)) 18 | 19 | assert results.total_params == 15 20 | assert results.trainable_params == 15 21 | 22 | @staticmethod 23 | def test_single_layer_network_on_gpu_device() -> None: 24 | model = torch.nn.Linear(2, 5) 25 | 26 | results = summary(model, input_size=(1, 2), device="cuda") 27 | 28 | assert results.total_params == 15 29 | assert results.trainable_params == 15 30 | 31 | @staticmethod 32 | def test_input_size_half_precision() -> None: 33 | # run this test case in gpu since 34 | # half precision is not supported in pytorch v1.12 35 | test = torch.nn.Linear(2, 5).half().to(torch.device("cuda")) 36 | with pytest.warns( 37 | UserWarning, 38 | match=( 39 | "Half precision is not supported with input_size parameter, and " 40 | "may output incorrect results. Try passing input_data directly." 41 | ), 42 | ): 43 | summary(test, dtypes=[torch.float16], input_size=(10, 2), device="cuda") 44 | 45 | @staticmethod 46 | def test_device() -> None: 47 | model = SingleInputNet() 48 | # input_size 49 | summary(model, input_size=(5, 1, 28, 28), device="cuda") 50 | 51 | # input_data 52 | model = SingleInputNet() 53 | input_data = torch.randn(5, 1, 28, 28) 54 | summary(model, input_data=input_data) 55 | summary(model, input_data=input_data, device="cuda") 56 | summary(model, input_data=input_data.to("cuda")) 57 | summary(model, input_data=input_data.to("cuda"), device=torch.device("cpu")) 58 | 59 | 60 | @pytest.mark.skipif( 61 | not torch.cuda.device_count() >= 2, reason="Only relevant for multi-GPU" 62 | ) 63 | class TestMultiGPU: 64 | """multi-GPU-only tests""" 65 | 66 | @staticmethod 67 | def test_model_stays_on_device_if_gpu() -> None: 68 | model = torch.nn.Linear(10, 10).to("cuda:1") 69 | summary(model) 70 | model_parameter = next(model.parameters()) 71 | assert model_parameter.device == torch.device("cuda:1") 72 | 73 | @staticmethod 74 | def test_different_model_parts_on_different_devices() -> None: 75 | model = torch.nn.Sequential( 76 | torch.nn.Linear(10, 10).to(1), torch.nn.Linear(10, 10).to(0) 77 | ) 78 | summary(model) 79 | 80 | 81 | @pytest.mark.skipif( 82 | not torch.cuda.is_available(), reason="Need CUDA to test parallelism." 83 | ) 84 | def test_device_parallelism() -> None: 85 | model = MultiDeviceModel("cpu", "cuda") 86 | input_data = torch.randn(10) 87 | summary(model, input_data=input_data) 88 | assert not next(model.net1.parameters()).is_cuda 89 | assert next(model.net2.parameters()).is_cuda 90 | -------------------------------------------------------------------------------- /tests/half_precision_test.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import torch 3 | 4 | from tests.fixtures.models import LinearModel, LSTMNet, SingleInputNet 5 | from torchinfo import summary 6 | from torchinfo.model_statistics import ModelStatistics 7 | 8 | 9 | @pytest.mark.skipif( 10 | not torch.cuda.is_available(), reason="Cuda must be available to test half models." 11 | ) 12 | class TestHalfPrecision: 13 | """ 14 | Iterate on these tests by using the following Google Colab link: 15 | https://colab.research.google.com/drive/1e_86DcAL6Q0r1OkjFcOuYKQRJaKsIro8 16 | """ 17 | 18 | @staticmethod 19 | def test_single_input_half() -> None: 20 | model = SingleInputNet() 21 | model.half().cuda() 22 | 23 | input_data = torch.randn((2, 1, 28, 28), dtype=torch.float16, device="cuda") 24 | results = summary(model, input_data=input_data) 25 | 26 | assert ModelStatistics.to_megabytes(results.total_param_bytes) == 0.04368 27 | assert ModelStatistics.to_megabytes(results.total_output_bytes) == 0.0568 28 | 29 | @staticmethod 30 | def test_linear_model_half() -> None: 31 | x = torch.randn((64, 128)) 32 | 33 | model = LinearModel() 34 | results = summary(model, input_data=x) 35 | 36 | model.half().cuda() 37 | x = x.type(torch.float16).cuda() 38 | results_half = summary(model, input_data=x) 39 | 40 | assert ModelStatistics.to_megabytes( 41 | results_half.total_param_bytes 42 | ) == pytest.approx(ModelStatistics.to_megabytes(results.total_param_bytes) / 2) 43 | assert ModelStatistics.to_megabytes( 44 | results_half.total_output_bytes 45 | ) == pytest.approx(ModelStatistics.to_megabytes(results.total_output_bytes) / 2) 46 | 47 | @staticmethod 48 | def test_lstm_half() -> None: 49 | model = LSTMNet() 50 | model.half() 51 | results = summary( 52 | model, 53 | input_size=(1, 100), 54 | dtypes=[torch.long], 55 | col_width=20, 56 | col_names=("kernel_size", "output_size", "num_params", "mult_adds"), 57 | row_settings=("var_names",), 58 | ) 59 | 60 | assert ModelStatistics.to_megabytes(results.total_param_bytes) == 7.56916 61 | assert ModelStatistics.to_megabytes(results.total_output_bytes) == 0.3328 62 | -------------------------------------------------------------------------------- /tests/jupyter_test.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": 1, 6 | "metadata": {}, 7 | "outputs": [], 8 | "source": [ 9 | "import torch\n", 10 | "import torch.nn as nn\n", 11 | "from torch.nn import functional as F\n", 12 | "from torchinfo import summary\n", 13 | "\n", 14 | "class SingleInputNet(nn.Module):\n", 15 | " \"\"\" Simple CNN model. \"\"\"\n", 16 | " def __init__(self) -> None:\n", 17 | " super().__init__()\n", 18 | " self.conv1 = nn.Conv2d(1, 10, kernel_size=5)\n", 19 | " self.conv2 = nn.Conv2d(10, 20, kernel_size=5)\n", 20 | " self.conv2_drop = nn.Dropout2d(0.3)\n", 21 | " self.fc1 = nn.Linear(320, 50)\n", 22 | " self.fc2 = nn.Linear(50, 10)\n", 23 | " def forward(self, x: torch.Tensor) -> torch.Tensor:\n", 24 | " x = F.relu(F.max_pool2d(self.conv1(x), 2))\n", 25 | " x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))\n", 26 | " x = x.view(-1, 320)\n", 27 | " x = F.relu(self.fc1(x))\n", 28 | " x = self.fc2(x)\n", 29 | " return F.log_softmax(x, dim=1)\n", 30 | "\n", 31 | "model = SingleInputNet()\n", 32 | "input_size = (2, 1, 28, 28)" 33 | ] 34 | }, 35 | { 36 | "cell_type": "code", 37 | "execution_count": 2, 38 | "metadata": {}, 39 | "outputs": [ 40 | { 41 | "data": { 42 | "text/plain": [ 43 | "==========================================================================================\n", 44 | "Layer (type:depth-idx) Output Shape Param #\n", 45 | "==========================================================================================\n", 46 | "├─Conv2d: 1-1 [2, 10, 24, 24] 260\n", 47 | "├─Conv2d: 1-2 [2, 20, 8, 8] 5,020\n", 48 | "├─Dropout2d: 1-3 [2, 20, 8, 8] --\n", 49 | "├─Linear: 1-4 [2, 50] 16,050\n", 50 | "├─Linear: 1-5 [2, 10] 510\n", 51 | "==========================================================================================\n", 52 | "Total params: 21,840\n", 53 | "Trainable params: 21,840\n", 54 | "Non-trainable params: 0\n", 55 | "Total mult-adds (M): 0.96\n", 56 | "==========================================================================================\n", 57 | "Input size (MB): 0.01\n", 58 | "Forward/backward pass size (MB): 0.11\n", 59 | "Params size (MB): 0.09\n", 60 | "Estimated Total Size (MB): 0.21\n", 61 | "==========================================================================================" 62 | ] 63 | }, 64 | "execution_count": 2, 65 | "metadata": {}, 66 | "output_type": "execute_result" 67 | } 68 | ], 69 | "source": [ 70 | "summary(model, input_size=input_size, verbose=0)" 71 | ] 72 | }, 73 | { 74 | "cell_type": "code", 75 | "execution_count": 3, 76 | "metadata": {}, 77 | "outputs": [ 78 | { 79 | "data": { 80 | "text/plain": [ 81 | "==========================================================================================\n", 82 | "Layer (type:depth-idx) Output Shape Param #\n", 83 | "==========================================================================================\n", 84 | "├─Conv2d: 1-1 [2, 10, 24, 24] 260\n", 85 | "├─Conv2d: 1-2 [2, 20, 8, 8] 5,020\n", 86 | "├─Dropout2d: 1-3 [2, 20, 8, 8] --\n", 87 | "├─Linear: 1-4 [2, 50] 16,050\n", 88 | "├─Linear: 1-5 [2, 10] 510\n", 89 | "==========================================================================================\n", 90 | "Total params: 21,840\n", 91 | "Trainable params: 21,840\n", 92 | "Non-trainable params: 0\n", 93 | "Total mult-adds (M): 0.96\n", 94 | "==========================================================================================\n", 95 | "Input size (MB): 0.01\n", 96 | "Forward/backward pass size (MB): 0.11\n", 97 | "Params size (MB): 0.09\n", 98 | "Estimated Total Size (MB): 0.21\n", 99 | "==========================================================================================" 100 | ] 101 | }, 102 | "execution_count": 3, 103 | "metadata": {}, 104 | "output_type": "execute_result" 105 | } 106 | ], 107 | "source": [ 108 | "summary(model, input_size=input_size)" 109 | ] 110 | }, 111 | { 112 | "cell_type": "code", 113 | "execution_count": 4, 114 | "metadata": {}, 115 | "outputs": [ 116 | { 117 | "data": { 118 | "text/plain": [ 119 | "4" 120 | ] 121 | }, 122 | "execution_count": 4, 123 | "metadata": {}, 124 | "output_type": "execute_result" 125 | } 126 | ], 127 | "source": [ 128 | "summary(model, input_size=input_size)\n", 129 | "2+2" 130 | ] 131 | }, 132 | { 133 | "cell_type": "code", 134 | "execution_count": 5, 135 | "metadata": {}, 136 | "outputs": [ 137 | { 138 | "name": "stdout", 139 | "output_type": "stream", 140 | "text": [ 141 | "==========================================================================================\n", 142 | "Layer (type:depth-idx) Output Shape Param #\n", 143 | "==========================================================================================\n", 144 | "├─Conv2d: 1-1 [2, 10, 24, 24] 260\n", 145 | "├─Conv2d: 1-2 [2, 20, 8, 8] 5,020\n", 146 | "├─Dropout2d: 1-3 [2, 20, 8, 8] --\n", 147 | "├─Linear: 1-4 [2, 50] 16,050\n", 148 | "├─Linear: 1-5 [2, 10] 510\n", 149 | "==========================================================================================\n", 150 | "Total params: 21,840\n", 151 | "Trainable params: 21,840\n", 152 | "Non-trainable params: 0\n", 153 | "Total mult-adds (M): 0.96\n", 154 | "==========================================================================================\n", 155 | "Input size (MB): 0.01\n", 156 | "Forward/backward pass size (MB): 0.11\n", 157 | "Params size (MB): 0.09\n", 158 | "Estimated Total Size (MB): 0.21\n", 159 | "==========================================================================================\n" 160 | ] 161 | }, 162 | { 163 | "data": { 164 | "text/plain": [ 165 | "4" 166 | ] 167 | }, 168 | "execution_count": 5, 169 | "metadata": {}, 170 | "output_type": "execute_result" 171 | } 172 | ], 173 | "source": [ 174 | "print(summary(model, input_size=input_size))\n", 175 | "2+2" 176 | ] 177 | }, 178 | { 179 | "cell_type": "code", 180 | "execution_count": null, 181 | "metadata": {}, 182 | "outputs": [], 183 | "source": [] 184 | } 185 | ], 186 | "metadata": { 187 | "kernelspec": { 188 | "display_name": "Python 3", 189 | "language": "python", 190 | "name": "python3" 191 | }, 192 | "language_info": { 193 | "codemirror_mode": { 194 | "name": "ipython", 195 | "version": 3 196 | }, 197 | "file_extension": ".py", 198 | "mimetype": "text/x-python", 199 | "name": "python", 200 | "nbconvert_exporter": "python", 201 | "pygments_lexer": "ipython3", 202 | "version": "3.8.3" 203 | } 204 | }, 205 | "nbformat": 4, 206 | "nbformat_minor": 2 207 | } 208 | -------------------------------------------------------------------------------- /tests/test_output/ascii_only.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type) Output Shape Param # 3 | ========================================================================================== 4 | ResNet [1, 1000] -- 5 | + Conv2d [1, 64, 32, 32] 9,408 6 | + BatchNorm2d [1, 64, 32, 32] 128 7 | + ReLU [1, 64, 32, 32] -- 8 | + MaxPool2d [1, 64, 16, 16] -- 9 | + Sequential [1, 64, 16, 16] -- 10 | | + BasicBlock [1, 64, 16, 16] -- 11 | | | + Conv2d [1, 64, 16, 16] 36,864 12 | | | + BatchNorm2d [1, 64, 16, 16] 128 13 | | | + ReLU [1, 64, 16, 16] -- 14 | | | + Conv2d [1, 64, 16, 16] 36,864 15 | | | + BatchNorm2d [1, 64, 16, 16] 128 16 | | | + ReLU [1, 64, 16, 16] -- 17 | | + BasicBlock [1, 64, 16, 16] -- 18 | | | + Conv2d [1, 64, 16, 16] 36,864 19 | | | + BatchNorm2d [1, 64, 16, 16] 128 20 | | | + ReLU [1, 64, 16, 16] -- 21 | | | + Conv2d [1, 64, 16, 16] 36,864 22 | | | + BatchNorm2d [1, 64, 16, 16] 128 23 | | | + ReLU [1, 64, 16, 16] -- 24 | + Sequential [1, 128, 8, 8] -- 25 | | + BasicBlock [1, 128, 8, 8] -- 26 | | | + Conv2d [1, 128, 8, 8] 73,728 27 | | | + BatchNorm2d [1, 128, 8, 8] 256 28 | | | + ReLU [1, 128, 8, 8] -- 29 | | | + Conv2d [1, 128, 8, 8] 147,456 30 | | | + BatchNorm2d [1, 128, 8, 8] 256 31 | | | + Sequential [1, 128, 8, 8] 8,448 32 | | | + ReLU [1, 128, 8, 8] -- 33 | | + BasicBlock [1, 128, 8, 8] -- 34 | | | + Conv2d [1, 128, 8, 8] 147,456 35 | | | + BatchNorm2d [1, 128, 8, 8] 256 36 | | | + ReLU [1, 128, 8, 8] -- 37 | | | + Conv2d [1, 128, 8, 8] 147,456 38 | | | + BatchNorm2d [1, 128, 8, 8] 256 39 | | | + ReLU [1, 128, 8, 8] -- 40 | + Sequential [1, 256, 4, 4] -- 41 | | + BasicBlock [1, 256, 4, 4] -- 42 | | | + Conv2d [1, 256, 4, 4] 294,912 43 | | | + BatchNorm2d [1, 256, 4, 4] 512 44 | | | + ReLU [1, 256, 4, 4] -- 45 | | | + Conv2d [1, 256, 4, 4] 589,824 46 | | | + BatchNorm2d [1, 256, 4, 4] 512 47 | | | + Sequential [1, 256, 4, 4] 33,280 48 | | | + ReLU [1, 256, 4, 4] -- 49 | | + BasicBlock [1, 256, 4, 4] -- 50 | | | + Conv2d [1, 256, 4, 4] 589,824 51 | | | + BatchNorm2d [1, 256, 4, 4] 512 52 | | | + ReLU [1, 256, 4, 4] -- 53 | | | + Conv2d [1, 256, 4, 4] 589,824 54 | | | + BatchNorm2d [1, 256, 4, 4] 512 55 | | | + ReLU [1, 256, 4, 4] -- 56 | + Sequential [1, 512, 2, 2] -- 57 | | + BasicBlock [1, 512, 2, 2] -- 58 | | | + Conv2d [1, 512, 2, 2] 1,179,648 59 | | | + BatchNorm2d [1, 512, 2, 2] 1,024 60 | | | + ReLU [1, 512, 2, 2] -- 61 | | | + Conv2d [1, 512, 2, 2] 2,359,296 62 | | | + BatchNorm2d [1, 512, 2, 2] 1,024 63 | | | + Sequential [1, 512, 2, 2] 132,096 64 | | | + ReLU [1, 512, 2, 2] -- 65 | | + BasicBlock [1, 512, 2, 2] -- 66 | | | + Conv2d [1, 512, 2, 2] 2,359,296 67 | | | + BatchNorm2d [1, 512, 2, 2] 1,024 68 | | | + ReLU [1, 512, 2, 2] -- 69 | | | + Conv2d [1, 512, 2, 2] 2,359,296 70 | | | + BatchNorm2d [1, 512, 2, 2] 1,024 71 | | | + ReLU [1, 512, 2, 2] -- 72 | + AdaptiveAvgPool2d [1, 512, 1, 1] -- 73 | + Linear [1, 1000] 513,000 74 | ========================================================================================== 75 | Total params: 11,689,512 76 | Trainable params: 11,689,512 77 | Non-trainable params: 0 78 | Total mult-adds (M): 148.57 79 | ========================================================================================== 80 | Input size (kB): 49.22 81 | Forward/backward pass size (MB): 3.25 82 | Params size (MB): 46.76 83 | Estimated Total Size (MB): 50.06 84 | ========================================================================================== 85 | -------------------------------------------------------------------------------- /tests/test_output/autoencoder.out: -------------------------------------------------------------------------------- 1 | =================================================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # Kernel Shape 3 | =================================================================================================================== 4 | AutoEncoder [1, 3, 64, 64] -- -- 5 | ├─Sequential: 1-1 [1, 16, 64, 64] -- -- 6 | │ └─Conv2d: 2-1 [1, 16, 64, 64] 448 [3, 3] 7 | │ └─ReLU: 2-2 [1, 16, 64, 64] -- -- 8 | ├─MaxPool2d: 1-2 [1, 16, 32, 32] -- 2 9 | ├─MaxUnpool2d: 1-3 [1, 16, 64, 64] -- [2, 2] 10 | ├─Sequential: 1-4 [1, 3, 64, 64] -- -- 11 | │ └─Conv2d: 2-3 [1, 3, 64, 64] 435 [3, 3] 12 | │ └─ReLU: 2-4 [1, 3, 64, 64] -- -- 13 | =================================================================================================================== 14 | Total params: 883 15 | Trainable params: 883 16 | Non-trainable params: 0 17 | Total mult-adds (M): 3.62 18 | =================================================================================================================== 19 | Input size (kB): 49.22 20 | Forward/backward pass size (kB): 622.59 21 | Params size (kB): 3.53 22 | Estimated Total Size (kB): 675.35 23 | =================================================================================================================== 24 | -------------------------------------------------------------------------------- /tests/test_output/basic_summary.out: -------------------------------------------------------------------------------- 1 | ================================================================= 2 | Layer (type:depth-idx) Param # 3 | ================================================================= 4 | SingleInputNet -- 5 | ├─Conv2d: 1-1 260 6 | ├─Conv2d: 1-2 5,020 7 | ├─Dropout2d: 1-3 -- 8 | ├─Linear: 1-4 16,050 9 | ├─Linear: 1-5 510 10 | ================================================================= 11 | Total params: 21,840 12 | Trainable params: 21,840 13 | Non-trainable params: 0 14 | ================================================================= 15 | -------------------------------------------------------------------------------- /tests/test_output/batch_size_optimization.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | SingleInputNet [1, 10] -- 5 | ├─Conv2d: 1-1 [1, 10, 24, 24] 260 6 | ├─Conv2d: 1-2 [1, 20, 8, 8] 5,020 7 | ├─Dropout2d: 1-3 [1, 20, 8, 8] -- 8 | ├─Linear: 1-4 [1, 50] 16,050 9 | ├─Linear: 1-5 [1, 10] 510 10 | ========================================================================================== 11 | Total params: 21,840 12 | Trainable params: 21,840 13 | Non-trainable params: 0 14 | Total mult-adds (k): 487.60 15 | ========================================================================================== 16 | Input size (kB): 3.21 17 | Forward/backward pass size (kB): 56.80 18 | Params size (kB): 87.36 19 | Estimated Total Size (kB): 147.37 20 | ========================================================================================== 21 | -------------------------------------------------------------------------------- /tests/test_output/bert.out: -------------------------------------------------------------------------------- 1 | ==================================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ==================================================================================================== 4 | BertModel [2, 768] -- 5 | ├─BertEmbeddings: 1-1 [2, 512, 768] -- 6 | │ └─Embedding: 2-1 [2, 512, 768] 23,440,896 7 | │ └─Embedding: 2-2 [2, 512, 768] 1,536 8 | │ └─Embedding: 2-3 [1, 512, 768] 393,216 9 | │ └─LayerNorm: 2-4 [2, 512, 768] 1,536 10 | │ └─Dropout: 2-5 [2, 512, 768] -- 11 | ├─BertEncoder: 1-2 [2, 512, 768] -- 12 | │ └─ModuleList: 2-6 -- -- 13 | │ │ └─BertLayer: 3-1 [2, 512, 768] 7,087,872 14 | │ │ └─BertLayer: 3-2 [2, 512, 768] 7,087,872 15 | │ │ └─BertLayer: 3-3 [2, 512, 768] 7,087,872 16 | │ │ └─BertLayer: 3-4 [2, 512, 768] 7,087,872 17 | │ │ └─BertLayer: 3-5 [2, 512, 768] 7,087,872 18 | │ │ └─BertLayer: 3-6 [2, 512, 768] 7,087,872 19 | │ │ └─BertLayer: 3-7 [2, 512, 768] 7,087,872 20 | │ │ └─BertLayer: 3-8 [2, 512, 768] 7,087,872 21 | │ │ └─BertLayer: 3-9 [2, 512, 768] 7,087,872 22 | │ │ └─BertLayer: 3-10 [2, 512, 768] 7,087,872 23 | │ │ └─BertLayer: 3-11 [2, 512, 768] 7,087,872 24 | │ │ └─BertLayer: 3-12 [2, 512, 768] 7,087,872 25 | ├─BertPooler: 1-3 [2, 768] -- 26 | │ └─Linear: 2-7 [2, 768] 590,592 27 | │ └─Tanh: 2-8 [2, 768] -- 28 | ==================================================================================================== 29 | Total params: 109,482,240 30 | Trainable params: 109,482,240 31 | Non-trainable params: 0 32 | Total mult-adds (M): 218.57 33 | ==================================================================================================== 34 | Input size (MB): 0.01 35 | Forward/backward pass size (MB): 852.50 36 | Params size (MB): 437.93 37 | Estimated Total Size (MB): 1290.45 38 | ==================================================================================================== 39 | -------------------------------------------------------------------------------- /tests/test_output/cnn_module_list.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | CNNModuleList [1, 1, 10] -- 5 | ├─ModuleList: 1-1 -- -- 6 | │ └─ConvLayerA: 2-1 [1, 1, 10] -- 7 | │ │ └─Conv1d: 3-1 [1, 1, 10] 2 8 | │ │ └─ReLU: 3-2 [1, 1, 10] -- 9 | │ │ └─MaxPool1d: 3-3 [1, 1, 10] -- 10 | │ └─ConvLayerA: 2-2 [1, 1, 10] -- 11 | │ │ └─Conv1d: 3-4 [1, 1, 10] 2 12 | │ │ └─ReLU: 3-5 [1, 1, 10] -- 13 | │ │ └─MaxPool1d: 3-6 [1, 1, 10] -- 14 | │ └─ConvLayerA: 2-3 [1, 1, 10] -- 15 | │ │ └─Conv1d: 3-7 [1, 1, 10] 2 16 | │ │ └─ReLU: 3-8 [1, 1, 10] -- 17 | │ │ └─MaxPool1d: 3-9 [1, 1, 10] -- 18 | │ └─ConvLayerA: 2-4 [1, 1, 10] -- 19 | │ │ └─Conv1d: 3-10 [1, 1, 10] 2 20 | │ │ └─ReLU: 3-11 [1, 1, 10] -- 21 | │ │ └─MaxPool1d: 3-12 [1, 1, 10] -- 22 | │ └─ConvLayerA: 2-5 [1, 1, 10] -- 23 | │ │ └─Conv1d: 3-13 [1, 1, 10] 2 24 | │ │ └─ReLU: 3-14 [1, 1, 10] -- 25 | │ │ └─MaxPool1d: 3-15 [1, 1, 10] -- 26 | ========================================================================================== 27 | Total params: 10 28 | Trainable params: 10 29 | Non-trainable params: 0 30 | Total mult-adds: 100 31 | ========================================================================================== 32 | Input size: 112 33 | Forward/backward pass size: 400 34 | Params size: 40 35 | Estimated Total Size: 552 36 | ========================================================================================== 37 | ========================================================================================== 38 | Layer (type:depth-idx) Output Shape Param # 39 | ========================================================================================== 40 | CNNModuleList [1, 1, 10] -- 41 | ├─ModuleList: 1-1 -- -- 42 | │ └─ConvLayerB: 2-1 [1, 1, 10] -- 43 | │ │ └─Conv1d: 3-1 [1, 1, 10] 2 44 | │ │ └─ReLU: 3-2 [1, 1, 10] -- 45 | │ │ └─MaxPool1d: 3-3 [1, 1, 10] -- 46 | │ └─ConvLayerB: 2-2 [1, 1, 10] -- 47 | │ │ └─Conv1d: 3-4 [1, 1, 10] 2 48 | │ │ └─ReLU: 3-5 [1, 1, 10] -- 49 | │ │ └─MaxPool1d: 3-6 [1, 1, 10] -- 50 | │ └─ConvLayerB: 2-3 [1, 1, 10] -- 51 | │ │ └─Conv1d: 3-7 [1, 1, 10] 2 52 | │ │ └─ReLU: 3-8 [1, 1, 10] -- 53 | │ │ └─MaxPool1d: 3-9 [1, 1, 10] -- 54 | │ └─ConvLayerB: 2-4 [1, 1, 10] -- 55 | │ │ └─Conv1d: 3-10 [1, 1, 10] 2 56 | │ │ └─ReLU: 3-11 [1, 1, 10] -- 57 | │ │ └─MaxPool1d: 3-12 [1, 1, 10] -- 58 | │ └─ConvLayerB: 2-5 [1, 1, 10] -- 59 | │ │ └─Conv1d: 3-13 [1, 1, 10] 2 60 | │ │ └─ReLU: 3-14 [1, 1, 10] -- 61 | │ │ └─MaxPool1d: 3-15 [1, 1, 10] -- 62 | ========================================================================================== 63 | Total params: 10 64 | Trainable params: 10 65 | Non-trainable params: 0 66 | Total mult-adds: 100 67 | ========================================================================================== 68 | Input size: 112 69 | Forward/backward pass size: 400 70 | Params size: 40 71 | Estimated Total Size: 552 72 | ========================================================================================== 73 | -------------------------------------------------------------------------------- /tests/test_output/compressai.out: -------------------------------------------------------------------------------- 1 | =============================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | =============================================================================================== 4 | FactorizedPrior [1, 192, 16, 16] -- 5 | ├─Sequential: 1-1 [1, 192, 16, 16] -- 6 | │ └─Conv2d: 2-1 [1, 128, 128, 128] 9,728 7 | │ └─GDN: 2-2 [1, 128, 128, 128] 16,512 8 | │ │ └─NonNegativeParametrizer: 3-1 [128] -- 9 | │ │ └─NonNegativeParametrizer: 3-2 [128, 128] -- 10 | │ └─Conv2d: 2-3 [1, 128, 64, 64] 409,728 11 | │ └─GDN: 2-4 [1, 128, 64, 64] 16,512 12 | │ │ └─NonNegativeParametrizer: 3-3 [128] -- 13 | │ │ └─NonNegativeParametrizer: 3-4 [128, 128] -- 14 | │ └─Conv2d: 2-5 [1, 128, 32, 32] 409,728 15 | │ └─GDN: 2-6 [1, 128, 32, 32] 16,512 16 | │ │ └─NonNegativeParametrizer: 3-5 [128] -- 17 | │ │ └─NonNegativeParametrizer: 3-6 [128, 128] -- 18 | │ └─Conv2d: 2-7 [1, 192, 16, 16] 614,592 19 | ├─EntropyBottleneck: 1-2 [1, 192, 16, 16] 11,712 20 | │ └─LowerBound: 2-8 [192, 1, 256] -- 21 | ├─Sequential: 1-3 [1, 3, 256, 256] -- 22 | │ └─ConvTranspose2d: 2-9 [1, 128, 32, 32] 614,528 23 | │ └─GDN: 2-10 [1, 128, 32, 32] 16,512 24 | │ │ └─NonNegativeParametrizer: 3-7 [128] -- 25 | │ │ └─NonNegativeParametrizer: 3-8 [128, 128] -- 26 | │ └─ConvTranspose2d: 2-11 [1, 128, 64, 64] 409,728 27 | │ └─GDN: 2-12 [1, 128, 64, 64] 16,512 28 | │ │ └─NonNegativeParametrizer: 3-9 [128] -- 29 | │ │ └─NonNegativeParametrizer: 3-10 [128, 128] -- 30 | │ └─ConvTranspose2d: 2-13 [1, 128, 128, 128] 409,728 31 | │ └─GDN: 2-14 [1, 128, 128, 128] 16,512 32 | │ │ └─NonNegativeParametrizer: 3-11 [128] -- 33 | │ │ └─NonNegativeParametrizer: 3-12 [128, 128] -- 34 | │ └─ConvTranspose2d: 2-15 [1, 3, 256, 256] 9,603 35 | =============================================================================================== 36 | Total params: 2,998,147 37 | Trainable params: 2,998,147 38 | Non-trainable params: 0 39 | Total mult-adds (G): 12.06 40 | =============================================================================================== 41 | Input size (MB): 0.79 42 | Forward/backward pass size (MB): 46.01 43 | Params size (MB): 11.55 44 | Estimated Total Size (MB): 58.34 45 | =============================================================================================== 46 | -------------------------------------------------------------------------------- /tests/test_output/container.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | ContainerModule [1, 5] -- 5 | ├─ModuleList: 1-1 -- -- 6 | │ └─Linear: 2-1 [1, 5] 30 7 | │ └─ContainerChildModule: 2-2 [1, 5] -- 8 | │ │ └─Sequential: 3-1 [1, 5] -- 9 | │ │ │ └─Linear: 4-1 [1, 5] 30 10 | │ │ │ └─Linear: 4-2 [1, 5] 30 11 | │ │ └─Linear: 3-2 [1, 5] 30 12 | │ │ └─Sequential: 3-3 -- (recursive) 13 | │ │ │ └─Linear: 4-3 [1, 5] (recursive) 14 | │ │ │ └─Linear: 4-4 [1, 5] (recursive) 15 | │ │ └─Sequential: 3-4 [1, 5] (recursive) 16 | │ │ │ └─Linear: 4-5 [1, 5] (recursive) 17 | │ │ │ └─Linear: 4-6 [1, 5] (recursive) 18 | │ │ │ └─Linear: 4-7 [1, 5] (recursive) 19 | │ │ │ └─Linear: 4-8 [1, 5] (recursive) 20 | │ └─Linear: 2-3 [1, 5] 30 21 | ========================================================================================== 22 | Total params: 150 23 | Trainable params: 150 24 | Non-trainable params: 0 25 | Total mult-adds: 330 26 | ========================================================================================== 27 | Input size: 92 28 | Forward/backward pass size: 440 29 | Params size: 600 30 | Estimated Total Size (kB): 1.13 31 | ========================================================================================== 32 | ========================================================================================== 33 | Layer (type:depth-idx) Output Shape Param # 34 | ========================================================================================== 35 | ContainerModule [5] -- 36 | ├─ModuleList: 1-1 -- -- 37 | │ └─Linear: 2-1 [5] 30 38 | │ └─ContainerChildModule: 2-2 [5] -- 39 | │ │ └─Sequential: 3-1 [5] 60 40 | │ │ └─Linear: 3-2 [5] 30 41 | │ │ └─Sequential: 3-3 -- (recursive) 42 | │ │ └─Sequential: 3-4 [5] (recursive) 43 | │ └─Linear: 2-3 [5] 30 44 | ========================================================================================== 45 | Total params: 150 46 | Trainable params: 150 47 | Non-trainable params: 0 48 | Total mult-adds: 330 49 | ========================================================================================== 50 | Input size: 92 51 | Forward/backward pass size: 440 52 | Params size: 600 53 | Estimated Total Size (kB): 1.13 54 | ========================================================================================== 55 | -------------------------------------------------------------------------------- /tests/test_output/custom_parameter.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | CustomParameter [4, 8] 32 5 | ========================================================================================== 6 | Total params: 32 7 | Trainable params: 32 8 | Non-trainable params: 0 9 | Total mult-adds: 128 10 | ========================================================================================== 11 | Input size: 76 12 | Forward/backward pass size: 256 13 | Params size: 128 14 | Estimated Total Size: 460 15 | ========================================================================================== 16 | -------------------------------------------------------------------------------- /tests/test_output/device.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | SingleInputNet [5, 10] -- 5 | ├─Conv2d: 1-1 [5, 10, 24, 24] 260 6 | ├─Conv2d: 1-2 [5, 20, 8, 8] 5,020 7 | ├─Dropout2d: 1-3 [5, 20, 8, 8] -- 8 | ├─Linear: 1-4 [5, 50] 16,050 9 | ├─Linear: 1-5 [5, 10] 510 10 | ========================================================================================== 11 | Total params: 21,840 12 | Trainable params: 21,840 13 | Non-trainable params: 0 14 | Total mult-adds (M): 2.44 15 | ========================================================================================== 16 | Input size (kB): 15.75 17 | Forward/backward pass size (kB): 284 18 | Params size (kB): 87.36 19 | Estimated Total Size (kB): 387.11 20 | ========================================================================================== 21 | ========================================================================================== 22 | Layer (type:depth-idx) Output Shape Param # 23 | ========================================================================================== 24 | SingleInputNet [5, 10] -- 25 | ├─Conv2d: 1-1 [5, 10, 24, 24] 260 26 | ├─Conv2d: 1-2 [5, 20, 8, 8] 5,020 27 | ├─Dropout2d: 1-3 [5, 20, 8, 8] -- 28 | ├─Linear: 1-4 [5, 50] 16,050 29 | ├─Linear: 1-5 [5, 10] 510 30 | ========================================================================================== 31 | Total params: 21,840 32 | Trainable params: 21,840 33 | Non-trainable params: 0 34 | Total mult-adds (M): 2.44 35 | ========================================================================================== 36 | Input size (kB): 15.75 37 | Forward/backward pass size (kB): 284 38 | Params size (kB): 87.36 39 | Estimated Total Size (kB): 387.11 40 | ========================================================================================== 41 | ========================================================================================== 42 | Layer (type:depth-idx) Output Shape Param # 43 | ========================================================================================== 44 | SingleInputNet [5, 10] -- 45 | ├─Conv2d: 1-1 [5, 10, 24, 24] 260 46 | ├─Conv2d: 1-2 [5, 20, 8, 8] 5,020 47 | ├─Dropout2d: 1-3 [5, 20, 8, 8] -- 48 | ├─Linear: 1-4 [5, 50] 16,050 49 | ├─Linear: 1-5 [5, 10] 510 50 | ========================================================================================== 51 | Total params: 21,840 52 | Trainable params: 21,840 53 | Non-trainable params: 0 54 | Total mult-adds (M): 2.44 55 | ========================================================================================== 56 | Input size (kB): 15.75 57 | Forward/backward pass size (kB): 284 58 | Params size (kB): 87.36 59 | Estimated Total Size (kB): 387.11 60 | ========================================================================================== 61 | ========================================================================================== 62 | Layer (type:depth-idx) Output Shape Param # 63 | ========================================================================================== 64 | SingleInputNet [5, 10] -- 65 | ├─Conv2d: 1-1 [5, 10, 24, 24] 260 66 | ├─Conv2d: 1-2 [5, 20, 8, 8] 5,020 67 | ├─Dropout2d: 1-3 [5, 20, 8, 8] -- 68 | ├─Linear: 1-4 [5, 50] 16,050 69 | ├─Linear: 1-5 [5, 10] 510 70 | ========================================================================================== 71 | Total params: 21,840 72 | Trainable params: 21,840 73 | Non-trainable params: 0 74 | Total mult-adds (M): 2.44 75 | ========================================================================================== 76 | Input size (kB): 15.75 77 | Forward/backward pass size (kB): 284 78 | Params size (kB): 87.36 79 | Estimated Total Size (kB): 387.11 80 | ========================================================================================== 81 | ========================================================================================== 82 | Layer (type:depth-idx) Output Shape Param # 83 | ========================================================================================== 84 | SingleInputNet [5, 10] -- 85 | ├─Conv2d: 1-1 [5, 10, 24, 24] 260 86 | ├─Conv2d: 1-2 [5, 20, 8, 8] 5,020 87 | ├─Dropout2d: 1-3 [5, 20, 8, 8] -- 88 | ├─Linear: 1-4 [5, 50] 16,050 89 | ├─Linear: 1-5 [5, 10] 510 90 | ========================================================================================== 91 | Total params: 21,840 92 | Trainable params: 21,840 93 | Non-trainable params: 0 94 | Total mult-adds (M): 2.44 95 | ========================================================================================== 96 | Input size (kB): 15.75 97 | Forward/backward pass size (kB): 284 98 | Params size (kB): 87.36 99 | Estimated Total Size (kB): 387.11 100 | ========================================================================================== 101 | -------------------------------------------------------------------------------- /tests/test_output/device_parallelism.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | MultiDeviceModel [5] -- 5 | ├─Linear: 1-1 [10] 110 6 | ├─ReLU: 1-2 [10] -- 7 | ├─Linear: 1-3 [5] 55 8 | ========================================================================================== 9 | Total params: 165 10 | Trainable params: 165 11 | Non-trainable params: 0 12 | Total mult-adds (M): 0.00 13 | ========================================================================================== 14 | Input size (MB): 0.00 15 | Forward/backward pass size (MB): 0.00 16 | Params size (MB): 0.00 17 | Estimated Total Size (MB): 0.00 18 | ========================================================================================== 19 | -------------------------------------------------------------------------------- /tests/test_output/dict_input.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | MultipleInputNetDifferentDtypes [2, 10] -- 5 | ├─Linear: 1-1 [1, 50] 15,050 6 | ├─Linear: 1-2 [1, 10] 510 7 | ├─Linear: 1-3 [1, 50] 15,050 8 | ├─Linear: 1-4 [1, 10] 510 9 | ========================================================================================== 10 | Total params: 31,120 11 | Trainable params: 31,120 12 | Non-trainable params: 0 13 | Total mult-adds (k): 31.12 14 | ========================================================================================== 15 | Input size (kB): 3.74 16 | Forward/backward pass size: 960 17 | Params size (kB): 124.48 18 | Estimated Total Size (kB): 129.18 19 | ========================================================================================== 20 | -------------------------------------------------------------------------------- /tests/test_output/dict_parameters_1.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | DictParameter [10, 1] -- 5 | ========================================================================================== 6 | Total params: 0 7 | Trainable params: 0 8 | Non-trainable params: 0 9 | Total mult-adds: 0 10 | ========================================================================================== 11 | Input size: 229 12 | Forward/backward pass size: 0 13 | Params size: 0 14 | Estimated Total Size: 229 15 | ========================================================================================== 16 | -------------------------------------------------------------------------------- /tests/test_output/dict_parameters_2.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | DictParameter [10, 1] -- 5 | ========================================================================================== 6 | Total params: 0 7 | Trainable params: 0 8 | Non-trainable params: 0 9 | Total mult-adds: 0 10 | ========================================================================================== 11 | Input size: 224 12 | Forward/backward pass size: 0 13 | Params size: 0 14 | Estimated Total Size: 224 15 | ========================================================================================== 16 | -------------------------------------------------------------------------------- /tests/test_output/dict_parameters_3.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | DictParameter [10, 1] -- 5 | ========================================================================================== 6 | Total params: 0 7 | Trainable params: 0 8 | Non-trainable params: 0 9 | Total mult-adds: 0 10 | ========================================================================================== 11 | Input size: 224 12 | Forward/backward pass size: 0 13 | Params size: 0 14 | Estimated Total Size: 224 15 | ========================================================================================== 16 | -------------------------------------------------------------------------------- /tests/test_output/edgecase_input_output_model.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | EdgecaseInputOutputModel -- -- 5 | ├─Linear: 1-1 [1] 4 6 | ========================================================================================== 7 | Total params: 4 8 | Trainable params: 4 9 | Non-trainable params: 0 10 | Total mult-adds: 4 11 | ========================================================================================== 12 | Input size: 0 13 | Forward/backward pass size: 8 14 | Params size: 16 15 | Estimated Total Size: 24 16 | ========================================================================================== 17 | -------------------------------------------------------------------------------- /tests/test_output/empty_module.out: -------------------------------------------------------------------------------- 1 | ================================================================= 2 | Layer (type:depth-idx) Param # 3 | ================================================================= 4 | EmptyModule -- 5 | ================================================================= 6 | Total params: 0 7 | Trainable params: 0 8 | Non-trainable params: 0 9 | ================================================================= 10 | -------------------------------------------------------------------------------- /tests/test_output/empty_module_list.out: -------------------------------------------------------------------------------- 1 | ================================================================= 2 | Layer (type:depth-idx) Param # 3 | ================================================================= 4 | ModuleList -- 5 | ================================================================= 6 | Total params: 0 7 | Trainable params: 0 8 | Non-trainable params: 0 9 | ================================================================= 10 | -------------------------------------------------------------------------------- /tests/test_output/fasterrcnn.out: -------------------------------------------------------------------------------- 1 | ========================================================================================================= 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================================= 4 | FasterRCNN -- -- 5 | ├─BackboneWithFPN: 1 -- -- 6 | │ └─FeaturePyramidNetwork: 2 -- -- 7 | │ │ └─ModuleList: 3-1 -- 984,064 8 | │ │ └─ModuleList: 3-2 -- 2,360,320 9 | ├─GeneralizedRCNNTransform: 1-1 -- -- 10 | ├─BackboneWithFPN: 1-2 [1, 256, 13, 13] -- 11 | │ └─IntermediateLayerGetter: 2-1 [1, 2048, 25, 25] -- 12 | │ │ └─Conv2d: 3-3 [1, 64, 400, 400] 9,408 13 | │ │ └─FrozenBatchNorm2d: 3-4 [1, 64, 400, 400] -- 14 | │ │ └─ReLU: 3-5 [1, 64, 400, 400] -- 15 | │ │ └─MaxPool2d: 3-6 [1, 64, 200, 200] -- 16 | │ │ └─Sequential: 3-7 [1, 256, 200, 200] 212,992 17 | │ │ └─Sequential: 3-8 [1, 512, 100, 100] 1,212,416 18 | │ │ └─Sequential: 3-9 [1, 1024, 50, 50] 7,077,888 19 | │ │ └─Sequential: 3-10 [1, 2048, 25, 25] 14,942,208 20 | │ └─FeaturePyramidNetwork: 2-2 [1, 256, 13, 13] -- 21 | │ │ └─LastLevelMaxPool: 3-11 [1, 256, 200, 200] -- 22 | ├─RegionProposalNetwork: 1-3 [1000, 4] -- 23 | │ └─RPNHead: 2-3 [1, 3, 200, 200] -- 24 | │ │ └─Conv2d: 3-12 [1, 256, 200, 200] 590,080 25 | │ │ └─Conv2d: 3-13 [1, 3, 200, 200] 771 26 | │ │ └─Conv2d: 3-14 [1, 12, 200, 200] 3,084 27 | │ │ └─Conv2d: 3-15 [1, 256, 100, 100] (recursive) 28 | │ │ └─Conv2d: 3-16 [1, 3, 100, 100] (recursive) 29 | │ │ └─Conv2d: 3-17 [1, 12, 100, 100] (recursive) 30 | │ │ └─Conv2d: 3-18 [1, 256, 50, 50] (recursive) 31 | │ │ └─Conv2d: 3-19 [1, 3, 50, 50] (recursive) 32 | │ │ └─Conv2d: 3-20 [1, 12, 50, 50] (recursive) 33 | │ │ └─Conv2d: 3-21 [1, 256, 25, 25] (recursive) 34 | │ │ └─Conv2d: 3-22 [1, 3, 25, 25] (recursive) 35 | │ │ └─Conv2d: 3-23 [1, 12, 25, 25] (recursive) 36 | │ │ └─Conv2d: 3-24 [1, 256, 13, 13] (recursive) 37 | │ │ └─Conv2d: 3-25 [1, 3, 13, 13] (recursive) 38 | │ │ └─Conv2d: 3-26 [1, 12, 13, 13] (recursive) 39 | │ └─AnchorGenerator: 2-4 [159882, 4] -- 40 | ├─RoIHeads: 1-4 -- -- 41 | │ └─MultiScaleRoIAlign: 2-5 [1000, 256, 7, 7] -- 42 | │ └─TwoMLPHead: 2-6 [1000, 1024] -- 43 | │ │ └─Linear: 3-27 [1000, 1024] 12,846,080 44 | │ │ └─Linear: 3-28 [1000, 1024] 1,049,600 45 | │ └─FastRCNNPredictor: 2-7 [1000, 91] -- 46 | │ │ └─Linear: 3-29 [1000, 91] 93,275 47 | │ │ └─Linear: 3-30 [1000, 364] 373,100 48 | ========================================================================================================= 49 | Total params: 41,755,286 50 | Trainable params: 41,755,286 51 | Non-trainable params: 0 52 | Total mult-adds (G): 134.42 53 | ========================================================================================================= 54 | Input size (MB): 0.15 55 | Forward/backward pass size (MB): 1458.42 56 | Params size (MB): 167.02 57 | Estimated Total Size (MB): 1625.60 58 | ========================================================================================================= 59 | -------------------------------------------------------------------------------- /tests/test_output/flan_t5_small.out: -------------------------------------------------------------------------------- 1 | ============================================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ============================================================================================================== 4 | T5ForConditionalGeneration [3, 100, 512] -- 5 | ├─T5Stack: 1-1 [3, 100, 512] 35,332,800 6 | ├─T5Stack: 1-2 -- (recursive) 7 | │ └─Embedding: 2-1 [3, 100, 512] 16,449,536 8 | ├─T5Stack: 1-3 -- (recursive) 9 | │ └─Dropout: 2-2 [3, 100, 512] -- 10 | │ └─ModuleList: 2-3 -- -- 11 | │ │ └─T5Block: 3-1 [3, 100, 512] 2,360,512 12 | │ │ └─T5Block: 3-2 [3, 100, 512] 2,360,320 13 | │ │ └─T5Block: 3-3 [3, 100, 512] 2,360,320 14 | │ │ └─T5Block: 3-4 [3, 100, 512] 2,360,320 15 | │ │ └─T5Block: 3-5 [3, 100, 512] 2,360,320 16 | │ │ └─T5Block: 3-6 [3, 100, 512] 2,360,320 17 | │ │ └─T5Block: 3-7 [3, 100, 512] 2,360,320 18 | │ │ └─T5Block: 3-8 [3, 100, 512] 2,360,320 19 | │ └─T5LayerNorm: 2-4 [3, 100, 512] 512 20 | │ └─Dropout: 2-5 [3, 100, 512] -- 21 | ├─T5Stack: 1-4 [3, 6, 100, 64] 16,449,536 22 | │ └─Embedding: 2-6 [3, 100, 512] (recursive) 23 | │ └─Dropout: 2-7 [3, 100, 512] -- 24 | │ └─ModuleList: 2-8 -- -- 25 | │ │ └─T5Block: 3-9 [3, 100, 512] 3,147,456 26 | │ │ └─T5Block: 3-10 [3, 100, 512] 3,147,264 27 | │ │ └─T5Block: 3-11 [3, 100, 512] 3,147,264 28 | │ │ └─T5Block: 3-12 [3, 100, 512] 3,147,264 29 | │ │ └─T5Block: 3-13 [3, 100, 512] 3,147,264 30 | │ │ └─T5Block: 3-14 [3, 100, 512] 3,147,264 31 | │ │ └─T5Block: 3-15 [3, 100, 512] 3,147,264 32 | │ │ └─T5Block: 3-16 [3, 100, 512] 3,147,264 33 | │ └─T5LayerNorm: 2-9 [3, 100, 512] 512 34 | │ └─Dropout: 2-10 [3, 100, 512] -- 35 | ├─Linear: 1-5 [3, 100, 32128] 16,449,536 36 | ============================================================================================================== 37 | Total params: 128,743,488 38 | Trainable params: 128,743,488 39 | Non-trainable params: 0 40 | Total mult-adds (G): 18.25 41 | ============================================================================================================== 42 | Input size (kB): 7.42 43 | Forward/backward pass size (MB): 326.28 44 | Params size (MB): 307.84 45 | Estimated Total Size (MB): 634.14 46 | ============================================================================================================== 47 | -------------------------------------------------------------------------------- /tests/test_output/formatting_options.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | SingleInputNet [16, 10] -- 5 | ├─Conv2d: 1-1 [16, 10, 24, 24] 260 6 | ├─Conv2d: 1-2 [16, 20, 8, 8] 5,020 7 | ├─Dropout2d: 1-3 [16, 20, 8, 8] -- 8 | ├─Linear: 1-4 [16, 50] 16,050 9 | ├─Linear: 1-5 [16, 10] 510 10 | ========================================================================================== 11 | Total params: 21,840 12 | Trainable params: 21,840 13 | Non-trainable params: 0 14 | Total mult-adds: 7,801,600 15 | ========================================================================================== 16 | Input size (kB): 50.25 17 | Forward/backward pass size (kB): 908.80 18 | Params size (kB): 87.36 19 | Estimated Total Size (MB): 1.05 20 | ========================================================================================== 21 | ========================================================================================== 22 | Layer (type:depth-idx) Output Shape Param # 23 | ========================================================================================== 24 | SingleInputNet [16, 10] -- 25 | ├─Conv2d: 1-1 [16, 10, 24, 24] 260 26 | ├─Conv2d: 1-2 [16, 20, 8, 8] 5,020 27 | ├─Dropout2d: 1-3 [16, 20, 8, 8] -- 28 | ├─Linear: 1-4 [16, 50] 16,050 29 | ├─Linear: 1-5 [16, 10] 510 30 | ========================================================================================== 31 | Total params: 21,840 32 | Trainable params: 21,840 33 | Non-trainable params: 0 34 | Total mult-adds (T): 0.00 35 | ========================================================================================== 36 | Input size (kB): 50.25 37 | Forward/backward pass size (kB): 908.80 38 | Params size (kB): 87.36 39 | Estimated Total Size (MB): 1.05 40 | ========================================================================================== 41 | -------------------------------------------------------------------------------- /tests/test_output/forward_pass_exception.out: -------------------------------------------------------------------------------- 1 | ============================================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ============================================================================================================== 4 | EdgeCaseModel [1, 10, 24, 24] -- 5 | ├─Conv2d: 1-1 [1, 10, 24, 24] 260 6 | ├─LayerWithRidiculouslyLongNameAndDoesntDoAnything: 1-2 [1, 10, 24, 24] -- 7 | │ └─Identity: 2-1 [1, 10, 24, 24] -- 8 | ============================================================================================================== 9 | Total params: 260 10 | Trainable params: 260 11 | Non-trainable params: 0 12 | Total mult-adds (k): 149.76 13 | ============================================================================================================== 14 | Input size (kB): 3.21 15 | Forward/backward pass size (kB): 46.08 16 | Params size (kB): 1.04 17 | Estimated Total Size (kB): 50.33 18 | ============================================================================================================== 19 | -------------------------------------------------------------------------------- /tests/test_output/frozen_layers.out: -------------------------------------------------------------------------------- 1 | ============================================================================================================================================ 2 | Layer (type:depth-idx) Output Shape Param # Kernel Shape Mult-Adds 3 | ============================================================================================================================================ 4 | ResNet [1, 1000] -- -- -- 5 | ├─Conv2d: 1-1 [1, 64, 32, 32] (9,408) [7, 7] 9,633,792 6 | ├─BatchNorm2d: 1-2 [1, 64, 32, 32] (128) -- 128 7 | ├─ReLU: 1-3 [1, 64, 32, 32] -- -- -- 8 | ├─MaxPool2d: 1-4 [1, 64, 16, 16] -- 3 -- 9 | ├─Sequential: 1-5 [1, 64, 16, 16] -- -- -- 10 | │ └─BasicBlock: 2-1 [1, 64, 16, 16] -- -- -- 11 | │ │ └─Conv2d: 3-1 [1, 64, 16, 16] (36,864) [3, 3] 9,437,184 12 | │ │ └─BatchNorm2d: 3-2 [1, 64, 16, 16] (128) -- 128 13 | │ │ └─ReLU: 3-3 [1, 64, 16, 16] -- -- -- 14 | │ │ └─Conv2d: 3-4 [1, 64, 16, 16] (36,864) [3, 3] 9,437,184 15 | │ │ └─BatchNorm2d: 3-5 [1, 64, 16, 16] (128) -- 128 16 | │ │ └─ReLU: 3-6 [1, 64, 16, 16] -- -- -- 17 | │ └─BasicBlock: 2-2 [1, 64, 16, 16] -- -- -- 18 | │ │ └─Conv2d: 3-7 [1, 64, 16, 16] (36,864) [3, 3] 9,437,184 19 | │ │ └─BatchNorm2d: 3-8 [1, 64, 16, 16] (128) -- 128 20 | │ │ └─ReLU: 3-9 [1, 64, 16, 16] -- -- -- 21 | │ │ └─Conv2d: 3-10 [1, 64, 16, 16] (36,864) [3, 3] 9,437,184 22 | │ │ └─BatchNorm2d: 3-11 [1, 64, 16, 16] (128) -- 128 23 | │ │ └─ReLU: 3-12 [1, 64, 16, 16] -- -- -- 24 | ├─Sequential: 1-6 [1, 128, 8, 8] -- -- -- 25 | │ └─BasicBlock: 2-3 [1, 128, 8, 8] -- -- -- 26 | │ │ └─Conv2d: 3-13 [1, 128, 8, 8] (73,728) [3, 3] 4,718,592 27 | │ │ └─BatchNorm2d: 3-14 [1, 128, 8, 8] (256) -- 256 28 | │ │ └─ReLU: 3-15 [1, 128, 8, 8] -- -- -- 29 | │ │ └─Conv2d: 3-16 [1, 128, 8, 8] (147,456) [3, 3] 9,437,184 30 | │ │ └─BatchNorm2d: 3-17 [1, 128, 8, 8] (256) -- 256 31 | │ │ └─Sequential: 3-18 [1, 128, 8, 8] (8,448) -- 524,544 32 | │ │ └─ReLU: 3-19 [1, 128, 8, 8] -- -- -- 33 | │ └─BasicBlock: 2-4 [1, 128, 8, 8] -- -- -- 34 | │ │ └─Conv2d: 3-20 [1, 128, 8, 8] (147,456) [3, 3] 9,437,184 35 | │ │ └─BatchNorm2d: 3-21 [1, 128, 8, 8] (256) -- 256 36 | │ │ └─ReLU: 3-22 [1, 128, 8, 8] -- -- -- 37 | │ │ └─Conv2d: 3-23 [1, 128, 8, 8] (147,456) [3, 3] 9,437,184 38 | │ │ └─BatchNorm2d: 3-24 [1, 128, 8, 8] (256) -- 256 39 | │ │ └─ReLU: 3-25 [1, 128, 8, 8] -- -- -- 40 | ├─Sequential: 1-7 [1, 256, 4, 4] -- -- -- 41 | │ └─BasicBlock: 2-5 [1, 256, 4, 4] -- -- -- 42 | │ │ └─Conv2d: 3-26 [1, 256, 4, 4] 294,912 [3, 3] 4,718,592 43 | │ │ └─BatchNorm2d: 3-27 [1, 256, 4, 4] 512 -- 512 44 | │ │ └─ReLU: 3-28 [1, 256, 4, 4] -- -- -- 45 | │ │ └─Conv2d: 3-29 [1, 256, 4, 4] 589,824 [3, 3] 9,437,184 46 | │ │ └─BatchNorm2d: 3-30 [1, 256, 4, 4] 512 -- 512 47 | │ │ └─Sequential: 3-31 [1, 256, 4, 4] 33,280 -- 524,800 48 | │ │ └─ReLU: 3-32 [1, 256, 4, 4] -- -- -- 49 | │ └─BasicBlock: 2-6 [1, 256, 4, 4] -- -- -- 50 | │ │ └─Conv2d: 3-33 [1, 256, 4, 4] 589,824 [3, 3] 9,437,184 51 | │ │ └─BatchNorm2d: 3-34 [1, 256, 4, 4] 512 -- 512 52 | │ │ └─ReLU: 3-35 [1, 256, 4, 4] -- -- -- 53 | │ │ └─Conv2d: 3-36 [1, 256, 4, 4] 589,824 [3, 3] 9,437,184 54 | │ │ └─BatchNorm2d: 3-37 [1, 256, 4, 4] 512 -- 512 55 | │ │ └─ReLU: 3-38 [1, 256, 4, 4] -- -- -- 56 | ├─Sequential: 1-8 [1, 512, 2, 2] -- -- -- 57 | │ └─BasicBlock: 2-7 [1, 512, 2, 2] -- -- -- 58 | │ │ └─Conv2d: 3-39 [1, 512, 2, 2] 1,179,648 [3, 3] 4,718,592 59 | │ │ └─BatchNorm2d: 3-40 [1, 512, 2, 2] 1,024 -- 1,024 60 | │ │ └─ReLU: 3-41 [1, 512, 2, 2] -- -- -- 61 | │ │ └─Conv2d: 3-42 [1, 512, 2, 2] 2,359,296 [3, 3] 9,437,184 62 | │ │ └─BatchNorm2d: 3-43 [1, 512, 2, 2] 1,024 -- 1,024 63 | │ │ └─Sequential: 3-44 [1, 512, 2, 2] 132,096 -- 525,312 64 | │ │ └─ReLU: 3-45 [1, 512, 2, 2] -- -- -- 65 | │ └─BasicBlock: 2-8 [1, 512, 2, 2] -- -- -- 66 | │ │ └─Conv2d: 3-46 [1, 512, 2, 2] 2,359,296 [3, 3] 9,437,184 67 | │ │ └─BatchNorm2d: 3-47 [1, 512, 2, 2] 1,024 -- 1,024 68 | │ │ └─ReLU: 3-48 [1, 512, 2, 2] -- -- -- 69 | │ │ └─Conv2d: 3-49 [1, 512, 2, 2] 2,359,296 [3, 3] 9,437,184 70 | │ │ └─BatchNorm2d: 3-50 [1, 512, 2, 2] 1,024 -- 1,024 71 | │ │ └─ReLU: 3-51 [1, 512, 2, 2] -- -- -- 72 | ├─AdaptiveAvgPool2d: 1-9 [1, 512, 1, 1] -- -- -- 73 | ├─Linear: 1-10 [1, 1000] 513,000 -- 513,000 74 | ============================================================================================================================================ 75 | Total params: 11,689,512 76 | Trainable params: 11,006,440 77 | Non-trainable params: 683,072 78 | Total mult-adds (M): 148.57 79 | ============================================================================================================================================ 80 | Input size (kB): 49.22 81 | Forward/backward pass size (MB): 3.25 82 | Params size (MB): 46.76 83 | Estimated Total Size (MB): 50.06 84 | ============================================================================================================================================ 85 | -------------------------------------------------------------------------------- /tests/test_output/genotype.out: -------------------------------------------------------------------------------- 1 | =============================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | =============================================================================================== 4 | GenotypeNetwork -- -- 5 | ├─Sequential: 1-1 [2, 48, 32, 32] -- 6 | │ └─Conv2d: 2-1 [2, 48, 32, 32] 1,296 7 | │ └─BatchNorm2d: 2-2 [2, 48, 32, 32] 96 8 | ├─ModuleList: 1-2 -- -- 9 | │ └─Cell: 2-3 [2, 32, 32, 32] -- 10 | │ │ └─ReLUConvBN: 3-1 [2, 32, 32, 32] 1,600 11 | │ │ └─ReLUConvBN: 3-2 [2, 32, 32, 32] 1,600 12 | │ │ └─ModuleList: 3-3 -- -- 13 | =============================================================================================== 14 | Total params: 4,592 15 | Trainable params: 4,592 16 | Non-trainable params: 0 17 | Total mult-adds (M): 8.95 18 | =============================================================================================== 19 | Input size (kB): 24.65 20 | Forward/backward pass size (MB): 3.67 21 | Params size (kB): 18.37 22 | Estimated Total Size (MB): 3.71 23 | =============================================================================================== 24 | =============================================================================================== 25 | Layer (type:depth-idx) Output Shape Param # 26 | =============================================================================================== 27 | GenotypeNetwork -- -- 28 | ├─Sequential: 1-1 [2, 48, 32, 32] -- 29 | │ └─Conv2d: 2-1 [2, 48, 32, 32] 1,296 30 | │ └─BatchNorm2d: 2-2 [2, 48, 32, 32] 96 31 | ├─ModuleList: 1-2 -- -- 32 | │ └─Cell: 2-3 [2, 32, 32, 32] -- 33 | │ │ └─ReLUConvBN: 3-1 [2, 32, 32, 32] -- 34 | │ │ │ └─Sequential: 4-1 [2, 32, 32, 32] -- 35 | │ │ │ │ └─ReLU: 5-1 [2, 48, 32, 32] -- 36 | │ │ │ │ └─Conv2d: 5-2 [2, 32, 32, 32] 1,536 37 | │ │ │ │ └─BatchNorm2d: 5-3 [2, 32, 32, 32] 64 38 | │ │ └─ReLUConvBN: 3-2 [2, 32, 32, 32] -- 39 | │ │ │ └─Sequential: 4-2 [2, 32, 32, 32] -- 40 | │ │ │ │ └─ReLU: 5-4 [2, 48, 32, 32] -- 41 | │ │ │ │ └─Conv2d: 5-5 [2, 32, 32, 32] 1,536 42 | │ │ │ │ └─BatchNorm2d: 5-6 [2, 32, 32, 32] 64 43 | │ │ └─ModuleList: 3-3 -- -- 44 | │ │ │ └─IdentityModel: 4-3 [2, 32, 32, 32] -- 45 | │ │ │ │ └─Identity: 5-7 [2, 32, 32, 32] -- 46 | │ │ │ └─IdentityModel: 4-4 [2, 32, 32, 32] -- 47 | │ │ │ │ └─Identity: 5-8 [2, 32, 32, 32] -- 48 | =============================================================================================== 49 | Total params: 4,592 50 | Trainable params: 4,592 51 | Non-trainable params: 0 52 | Total mult-adds (M): 8.95 53 | =============================================================================================== 54 | Input size (kB): 24.65 55 | Forward/backward pass size (MB): 3.67 56 | Params size (kB): 18.37 57 | Estimated Total Size (MB): 3.71 58 | =============================================================================================== 59 | -------------------------------------------------------------------------------- /tests/test_output/groups.out: -------------------------------------------------------------------------------- 1 | ================================================================================================================================================================ 2 | Layer (type:depth-idx) Kernel Shape Groups Input Shape Output Shape Param # Mult-Adds 3 | ================================================================================================================================================================ 4 | Conv2d [3, 3] 4 [7, 16, 28, 28] [7, 32, 26, 26] 1,184 5,602,688 5 | ================================================================================================================================================================ 6 | Total params: 1,184 7 | Trainable params: 1,184 8 | Non-trainable params: 0 9 | Total mult-adds (M): 5.60 10 | ================================================================================================================================================================ 11 | Input size (kB): 351.30 12 | Forward/backward pass size (MB): 1.21 13 | Params size (kB): 4.74 14 | Estimated Total Size (MB): 1.57 15 | ================================================================================================================================================================ 16 | -------------------------------------------------------------------------------- /tests/test_output/hide_recursive_layers.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | SimpleRNN [2, 2] 9 5 | ├─LSTMCell: 1-1 [2, 2] 56 6 | ├─Tanh: 1-2 [2, 2] -- 7 | ├─LSTMCell: 1-3 [2, 2] (recursive) 8 | ├─Tanh: 1-4 [2, 2] -- 9 | ├─LSTMCell: 1-5 [2, 2] (recursive) 10 | ├─Tanh: 1-6 [2, 2] -- 11 | ├─LSTMCell: 1-7 [2, 2] (recursive) 12 | ├─Tanh: 1-8 [2, 2] -- 13 | ========================================================================================== 14 | Total params: 65 15 | Trainable params: 65 16 | Non-trainable params: 0 17 | Total mult-adds: 896 18 | ========================================================================================== 19 | Input size: 96 20 | Forward/backward pass size: 128 21 | Params size: 224 22 | Estimated Total Size: 448 23 | ========================================================================================== 24 | ========================================================================================== 25 | Layer (type:depth-idx) Output Shape Param # 26 | ========================================================================================== 27 | SimpleRNN [2, 2] 9 28 | ├─LSTMCell: 1-1 [2, 2] 56 29 | ├─Tanh: 1-2 [2, 2] -- 30 | ========================================================================================== 31 | Total params: 65 32 | Trainable params: 65 33 | Non-trainable params: 0 34 | Total mult-adds: 896 35 | ========================================================================================== 36 | Input size: 96 37 | Forward/backward pass size: 128 38 | Params size: 224 39 | Estimated Total Size: 448 40 | ========================================================================================== 41 | -------------------------------------------------------------------------------- /tests/test_output/hide_recursive_layers_outside_loop.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | SimpleRNN [2, 3] -- 5 | ├─LSTMCell: 1-1 [2, 2] 56 6 | ├─Tanh: 1-2 [2, 2] -- 7 | ├─LSTMCell: 1-3 [2, 2] (recursive) 8 | ├─Tanh: 1-4 [2, 2] -- 9 | ├─LSTMCell: 1-5 [2, 2] (recursive) 10 | ├─Tanh: 1-6 [2, 2] -- 11 | ├─LSTMCell: 1-7 [2, 2] (recursive) 12 | ├─Tanh: 1-8 [2, 2] -- 13 | ├─Linear: 1-9 [2, 3] 9 14 | ├─Tanh: 1-10 [2, 3] -- 15 | ========================================================================================== 16 | Total params: 65 17 | Trainable params: 65 18 | Non-trainable params: 0 19 | Total mult-adds: 914 20 | ========================================================================================== 21 | Input size: 96 22 | Forward/backward pass size: 176 23 | Params size: 260 24 | Estimated Total Size: 532 25 | ========================================================================================== 26 | ========================================================================================== 27 | Layer (type:depth-idx) Output Shape Param # 28 | ========================================================================================== 29 | SimpleRNN [2, 3] -- 30 | ├─LSTMCell: 1-1 [2, 2] 56 31 | ├─Tanh: 1-2 [2, 2] -- 32 | ├─Linear: 1-9 [2, 3] 9 33 | ========================================================================================== 34 | Total params: 65 35 | Trainable params: 65 36 | Non-trainable params: 0 37 | Total mult-adds: 914 38 | ========================================================================================== 39 | Input size: 96 40 | Forward/backward pass size: 176 41 | Params size: 260 42 | Estimated Total Size: 532 43 | ========================================================================================== 44 | -------------------------------------------------------------------------------- /tests/test_output/highly_nested_dict_model.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | HighlyNestedDictModel [10] -- 5 | ├─Linear: 1-1 [10] 110 6 | ├─Linear: 1-2 [10] 110 7 | ========================================================================================== 8 | Total params: 220 9 | Trainable params: 220 10 | Non-trainable params: 0 11 | Total mult-adds: 220 12 | ========================================================================================== 13 | Input size: 112 14 | Forward/backward pass size: 160 15 | Params size: 880 16 | Estimated Total Size (kB): 1.15 17 | ========================================================================================== 18 | -------------------------------------------------------------------------------- /tests/test_output/input_size_possibilities.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | CustomParameter [3, 2] 6 5 | ========================================================================================== 6 | Total params: 6 7 | Trainable params: 6 8 | Non-trainable params: 0 9 | Total mult-adds: 18 10 | ========================================================================================== 11 | Input size: 80 12 | Forward/backward pass size: 48 13 | Params size: 24 14 | Estimated Total Size: 152 15 | ========================================================================================== 16 | ========================================================================================== 17 | Layer (type:depth-idx) Output Shape Param # 18 | ========================================================================================== 19 | CustomParameter [3, 2] 6 20 | ========================================================================================== 21 | Total params: 6 22 | Trainable params: 6 23 | Non-trainable params: 0 24 | Total mult-adds: 18 25 | ========================================================================================== 26 | Input size: 80 27 | Forward/backward pass size: 48 28 | Params size: 24 29 | Estimated Total Size: 152 30 | ========================================================================================== 31 | ========================================================================================== 32 | Layer (type:depth-idx) Output Shape Param # 33 | ========================================================================================== 34 | CustomParameter [3, 2] 6 35 | ========================================================================================== 36 | Total params: 6 37 | Trainable params: 6 38 | Non-trainable params: 0 39 | Total mult-adds: 18 40 | ========================================================================================== 41 | Input size: 80 42 | Forward/backward pass size: 48 43 | Params size: 24 44 | Estimated Total Size: 152 45 | ========================================================================================== 46 | ========================================================================================== 47 | Layer (type:depth-idx) Output Shape Param # 48 | ========================================================================================== 49 | CustomParameter [3, 2] 6 50 | ========================================================================================== 51 | Total params: 6 52 | Trainable params: 6 53 | Non-trainable params: 0 54 | Total mult-adds: 18 55 | ========================================================================================== 56 | Input size: 80 57 | Forward/backward pass size: 48 58 | Params size: 24 59 | Estimated Total Size: 152 60 | ========================================================================================== 61 | -------------------------------------------------------------------------------- /tests/test_output/input_tensor.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | SingleInputNet [5, 10] -- 5 | ├─Conv2d: 1-1 [5, 10, 24, 24] 260 6 | ├─Conv2d: 1-2 [5, 20, 8, 8] 5,020 7 | ├─Dropout2d: 1-3 [5, 20, 8, 8] -- 8 | ├─Linear: 1-4 [5, 50] 16,050 9 | ├─Linear: 1-5 [5, 10] 510 10 | ========================================================================================== 11 | Total params: 21,840 12 | Trainable params: 21,840 13 | Non-trainable params: 0 14 | Total mult-adds (M): 2.44 15 | ========================================================================================== 16 | Input size (kB): 15.75 17 | Forward/backward pass size (kB): 284 18 | Params size (kB): 87.36 19 | Estimated Total Size (kB): 387.11 20 | ========================================================================================== 21 | -------------------------------------------------------------------------------- /tests/test_output/jit.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | LinearModel [64, 1] -- 5 | ├─Sequential: 1-1 [64, 1] -- 6 | │ └─Linear: 2-1 [64, 128] 16,512 7 | │ └─ReLU: 2-2 [64, 128] -- 8 | │ └─Linear: 2-3 [64, 128] 16,512 9 | │ └─ReLU: 2-4 [64, 128] -- 10 | │ └─Linear: 2-5 [64, 1] 129 11 | ========================================================================================== 12 | Total params: 33,153 13 | Trainable params: 33,153 14 | Non-trainable params: 0 15 | Total mult-adds (M): 2.12 16 | ========================================================================================== 17 | Input size (kB): 32.84 18 | Forward/backward pass size (kB): 131.58 19 | Params size (kB): 132.61 20 | Estimated Total Size (kB): 297.04 21 | ========================================================================================== 22 | ========================================================================================== 23 | Layer (type:depth-idx) Output Shape Param # 24 | ========================================================================================== 25 | LinearModel -- -- 26 | ├─Sequential: 1-1 -- -- 27 | │ └─Linear: 2-1 -- 16,512 28 | │ └─ReLU: 2-2 -- -- 29 | │ └─Linear: 2-3 -- 16,512 30 | │ └─ReLU: 2-4 -- -- 31 | │ └─Linear: 2-5 -- 129 32 | ========================================================================================== 33 | Total params: 33,153 34 | Trainable params: 33,153 35 | Non-trainable params: 0 36 | Total mult-adds: 0 37 | ========================================================================================== 38 | Input size (kB): 32.84 39 | Forward/backward pass size: 0 40 | Params size (kB): 132.61 41 | Estimated Total Size (kB): 165.45 42 | ========================================================================================== 43 | -------------------------------------------------------------------------------- /tests/test_output/linear.out: -------------------------------------------------------------------------------- 1 | ======================================================================================================================== 2 | Layer (type:depth-idx) Input Shape Output Shape Param # Mult-Adds 3 | ======================================================================================================================== 4 | Linear [32, 16, 8] [32, 16, 64] 576 294,912 5 | ======================================================================================================================== 6 | Total params: 576 7 | Trainable params: 576 8 | Non-trainable params: 0 9 | Total mult-adds (k): 294.91 10 | ======================================================================================================================== 11 | Input size (kB): 16.46 12 | Forward/backward pass size (kB): 262.14 13 | Params size (kB): 2.30 14 | Estimated Total Size (kB): 280.90 15 | ======================================================================================================================== 16 | -------------------------------------------------------------------------------- /tests/test_output/linear_model_half.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | LinearModel [64, 1] -- 5 | ├─Sequential: 1-1 [64, 1] -- 6 | │ └─Linear: 2-1 [64, 128] 16,512 7 | │ └─ReLU: 2-2 [64, 128] -- 8 | │ └─Linear: 2-3 [64, 128] 16,512 9 | │ └─ReLU: 2-4 [64, 128] -- 10 | │ └─Linear: 2-5 [64, 1] 129 11 | ========================================================================================== 12 | Total params: 33,153 13 | Trainable params: 33,153 14 | Non-trainable params: 0 15 | Total mult-adds (M): 2.12 16 | ========================================================================================== 17 | Input size (MB): 0.03 18 | Forward/backward pass size (MB): 0.13 19 | Params size (MB): 0.13 20 | Estimated Total Size (MB): 0.30 21 | ========================================================================================== 22 | ========================================================================================== 23 | Layer (type:depth-idx) Output Shape Param # 24 | ========================================================================================== 25 | LinearModel [64, 1] -- 26 | ├─Sequential: 1-1 [64, 1] -- 27 | │ └─Linear: 2-1 [64, 128] 16,512 28 | │ └─ReLU: 2-2 [64, 128] -- 29 | │ └─Linear: 2-3 [64, 128] 16,512 30 | │ └─ReLU: 2-4 [64, 128] -- 31 | │ └─Linear: 2-5 [64, 1] 129 32 | ========================================================================================== 33 | Total params: 33,153 34 | Trainable params: 33,153 35 | Non-trainable params: 0 36 | Total mult-adds (M): 2.12 37 | ========================================================================================== 38 | Input size (MB): 0.02 39 | Forward/backward pass size (MB): 0.07 40 | Params size (MB): 0.07 41 | Estimated Total Size (MB): 0.15 42 | ========================================================================================== 43 | -------------------------------------------------------------------------------- /tests/test_output/lstm.out: -------------------------------------------------------------------------------- 1 | ======================================================================================================================== 2 | Layer (type (var_name)) Kernel Shape Output Shape Param # Mult-Adds 3 | ======================================================================================================================== 4 | LSTMNet (LSTMNet) -- [100, 20] -- -- 5 | ├─Embedding (embedding) -- [1, 100, 300] 6,000 6,000 6 | │ └─weight [300, 20] └─6,000 7 | ├─LSTM (encoder) -- [1, 100, 512] 3,768,320 376,832,000 8 | │ └─weight_ih_l0 [2048, 300] ├─614,400 9 | │ └─weight_hh_l0 [2048, 512] ├─1,048,576 10 | │ └─bias_ih_l0 [2048] ├─2,048 11 | │ └─bias_hh_l0 [2048] ├─2,048 12 | │ └─weight_ih_l1 [2048, 512] ├─1,048,576 13 | │ └─weight_hh_l1 [2048, 512] ├─1,048,576 14 | │ └─bias_ih_l1 [2048] ├─2,048 15 | │ └─bias_hh_l1 [2048] └─2,048 16 | ├─Linear (decoder) -- [1, 100, 20] 10,260 1,026,000 17 | │ └─weight [512, 20] ├─10,240 18 | │ └─bias [20] └─20 19 | ======================================================================================================================== 20 | Total params: 3,784,580 21 | Trainable params: 3,784,580 22 | Non-trainable params: 0 23 | Total mult-adds (M): 377.86 24 | ======================================================================================================================== 25 | Input size: 872 26 | Forward/backward pass size (kB): 665.60 27 | Params size (MB): 15.14 28 | Estimated Total Size (MB): 15.80 29 | ======================================================================================================================== 30 | -------------------------------------------------------------------------------- /tests/test_output/lstm_custom_batch_size.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | LSTMNet [100, 1] -- 5 | ├─Embedding: 1-1 [100, 1, 300] 6,000 6 | ├─LSTM: 1-2 [100, 1, 512] 3,768,320 7 | ├─Linear: 1-3 [100, 1, 20] 10,260 8 | ========================================================================================== 9 | Total params: 3,784,580 10 | Trainable params: 3,784,580 11 | Non-trainable params: 0 12 | Total mult-adds (M): 378.46 13 | ========================================================================================== 14 | Input size: 872 15 | Forward/backward pass size (kB): 665.60 16 | Params size (MB): 15.14 17 | Estimated Total Size (MB): 15.80 18 | ========================================================================================== 19 | -------------------------------------------------------------------------------- /tests/test_output/lstm_half.out: -------------------------------------------------------------------------------- 1 | ======================================================================================================================== 2 | Layer (type (var_name)) Kernel Shape Output Shape Param # Mult-Adds 3 | ======================================================================================================================== 4 | LSTMNet (LSTMNet) -- [100, 20] -- -- 5 | ├─Embedding (embedding) -- [1, 100, 300] 6,000 6,000 6 | ├─LSTM (encoder) -- [1, 100, 512] 3,768,320 376,832,000 7 | ├─Linear (decoder) -- [1, 100, 20] 10,260 1,026,000 8 | ======================================================================================================================== 9 | Total params: 3,784,580 10 | Trainable params: 3,784,580 11 | Non-trainable params: 0 12 | Total mult-adds (M): 377.86 13 | ======================================================================================================================== 14 | Input size (MB): 0.00 15 | Forward/backward pass size (MB): 0.33 16 | Params size (MB): 7.57 17 | Estimated Total Size (MB): 7.90 18 | ======================================================================================================================== 19 | -------------------------------------------------------------------------------- /tests/test_output/mixed_trainable_parameters.out: -------------------------------------------------------------------------------- 1 | ================================================================= 2 | Layer (type:depth-idx) Param # 3 | ================================================================= 4 | MixedTrainableParameters 20 5 | ├─w ├─10 6 | ├─b └─10 7 | ================================================================= 8 | Total params: 20 9 | Trainable params: 10 10 | Non-trainable params: 10 11 | ================================================================= 12 | -------------------------------------------------------------------------------- /tests/test_output/model_with_args.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | RecursiveNet [1, 64, 28, 28] -- 5 | ├─Conv2d: 1-1 [1, 64, 28, 28] 36,928 6 | ├─Conv2d: 1-2 [1, 64, 28, 28] (recursive) 7 | ├─Conv2d: 1-3 [1, 64, 28, 28] (recursive) 8 | ├─Conv2d: 1-4 [1, 64, 28, 28] (recursive) 9 | ├─Conv2d: 1-5 [1, 64, 28, 28] (recursive) 10 | ├─Conv2d: 1-6 [1, 64, 28, 28] (recursive) 11 | ========================================================================================== 12 | Total params: 36,928 13 | Trainable params: 36,928 14 | Non-trainable params: 0 15 | Total mult-adds (M): 173.71 16 | ========================================================================================== 17 | Input size (kB): 200.78 18 | Forward/backward pass size (MB): 2.41 19 | Params size (kB): 147.71 20 | Estimated Total Size (MB): 2.76 21 | ========================================================================================== 22 | -------------------------------------------------------------------------------- /tests/test_output/module_dict.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | ModuleDictModel [1, 10, 1, 1] -- 5 | ├─ModuleDict: 1-1 -- -- 6 | │ └─Conv2d: 2-1 [1, 10, 1, 1] 910 7 | ├─ModuleDict: 1-2 -- 1 8 | │ └─LeakyReLU: 2-2 [1, 10, 1, 1] -- 9 | ========================================================================================== 10 | Total params: 911 11 | Trainable params: 911 12 | Non-trainable params: 0 13 | Total mult-adds: 910 14 | ========================================================================================== 15 | Input size: 432 16 | Forward/backward pass size: 80 17 | Params size (kB): 3.64 18 | Estimated Total Size (kB): 4.15 19 | ========================================================================================== 20 | ========================================================================================== 21 | Layer (type:depth-idx) Output Shape Param # 22 | ========================================================================================== 23 | ModuleDictModel [1, 10, 1, 1] -- 24 | ├─ModuleDict: 1-1 -- 910 25 | │ └─MaxPool2d: 2-1 [1, 10, 1, 1] -- 26 | ├─ModuleDict: 1-2 -- -- 27 | │ └─PReLU: 2-2 [1, 10, 1, 1] 1 28 | ========================================================================================== 29 | Total params: 911 30 | Trainable params: 911 31 | Non-trainable params: 0 32 | Total mult-adds: 1 33 | ========================================================================================== 34 | Input size: 432 35 | Forward/backward pass size: 80 36 | Params size: 4 37 | Estimated Total Size: 516 38 | ========================================================================================== 39 | -------------------------------------------------------------------------------- /tests/test_output/multiple_input_tensor_args.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | MultipleInputNetDifferentDtypes [2, 10] -- 5 | ├─Linear: 1-1 [1, 50] 15,050 6 | ├─Linear: 1-2 [1, 10] 510 7 | ├─Linear: 1-3 [1, 50] 15,050 8 | ├─Linear: 1-4 [1, 10] 510 9 | ========================================================================================== 10 | Total params: 31,120 11 | Trainable params: 31,120 12 | Non-trainable params: 0 13 | Total mult-adds (k): 31.12 14 | ========================================================================================== 15 | Input size (kB): 1.27 16 | Forward/backward pass size: 960 17 | Params size (kB): 124.48 18 | Estimated Total Size (kB): 126.71 19 | ========================================================================================== 20 | -------------------------------------------------------------------------------- /tests/test_output/multiple_input_tensor_dict.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | MultipleInputNetDifferentDtypes [2, 10] -- 5 | ├─Linear: 1-1 [1, 50] 15,050 6 | ├─Linear: 1-2 [1, 10] 510 7 | ├─Linear: 1-3 [1, 50] 15,050 8 | ├─Linear: 1-4 [1, 10] 510 9 | ========================================================================================== 10 | Total params: 31,120 11 | Trainable params: 31,120 12 | Non-trainable params: 0 13 | Total mult-adds (k): 31.12 14 | ========================================================================================== 15 | Input size (kB): 3.74 16 | Forward/backward pass size: 960 17 | Params size (kB): 124.48 18 | Estimated Total Size (kB): 129.18 19 | ========================================================================================== 20 | -------------------------------------------------------------------------------- /tests/test_output/multiple_input_tensor_list.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | MultipleInputNetDifferentDtypes [2, 10] -- 5 | ├─Linear: 1-1 [1, 50] 15,050 6 | ├─Linear: 1-2 [1, 10] 510 7 | ├─Linear: 1-3 [1, 50] 15,050 8 | ├─Linear: 1-4 [1, 10] 510 9 | ========================================================================================== 10 | Total params: 31,120 11 | Trainable params: 31,120 12 | Non-trainable params: 0 13 | Total mult-adds (k): 31.12 14 | ========================================================================================== 15 | Input size (kB): 3.74 16 | Forward/backward pass size: 960 17 | Params size (kB): 124.48 18 | Estimated Total Size (kB): 129.18 19 | ========================================================================================== 20 | -------------------------------------------------------------------------------- /tests/test_output/multiple_input_types.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | MultipleInputNetDifferentDtypes [2, 10] -- 5 | ├─Linear: 1-1 [1, 50] 15,050 6 | ├─Linear: 1-2 [1, 10] 510 7 | ├─Linear: 1-3 [1, 50] 15,050 8 | ├─Linear: 1-4 [1, 10] 510 9 | ========================================================================================== 10 | Total params: 31,120 11 | Trainable params: 31,120 12 | Non-trainable params: 0 13 | Total mult-adds (k): 31.12 14 | ========================================================================================== 15 | Input size (kB): 3.74 16 | Forward/backward pass size: 960 17 | Params size (kB): 124.48 18 | Estimated Total Size (kB): 129.18 19 | ========================================================================================== 20 | -------------------------------------------------------------------------------- /tests/test_output/namedtuple.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | NamedTuple [2, 1, 28, 28] -- 5 | ========================================================================================== 6 | Total params: 0 7 | Trainable params: 0 8 | Non-trainable params: 0 9 | Total mult-adds: 0 10 | ========================================================================================== 11 | Input size (kB): 12.69 12 | Forward/backward pass size: 0 13 | Params size: 0 14 | Estimated Total Size (kB): 12.69 15 | ========================================================================================== 16 | -------------------------------------------------------------------------------- /tests/test_output/nested_leftover_params.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type (var_name)) Output Shape Param # 3 | ========================================================================================== 4 | InsideModel (InsideModel) [100, 2] 2 5 | ├─Linear (l_0) [100, 1] 3 6 | ├─Inside (inside) [100, 1] 1 7 | │ └─Linear (l_1) [100, 1] 2 8 | ========================================================================================== 9 | Total params: 8 10 | Trainable params: 8 11 | Non-trainable params: 0 12 | Total mult-adds: 500 13 | ========================================================================================== 14 | Input size: 872 15 | Forward/backward pass size (kB): 1.60 16 | Params size: 20 17 | Estimated Total Size (kB): 2.49 18 | ========================================================================================== 19 | -------------------------------------------------------------------------------- /tests/test_output/numpy_model.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | NumpyModel [3] -- 5 | ├─Linear: 1-1 [3] 12 6 | ========================================================================================== 7 | Total params: 12 8 | Trainable params: 12 9 | Non-trainable params: 0 10 | Total mult-adds (M): 0.00 11 | ========================================================================================== 12 | Input size (MB): 0.00 13 | Forward/backward pass size (MB): 0.00 14 | Params size (MB): 0.00 15 | Estimated Total Size (MB): 0.00 16 | ========================================================================================== 17 | -------------------------------------------------------------------------------- /tests/test_output/pack_padded.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | PackPaddedLSTM [128, 18] -- 5 | ├─Embedding: 1-1 [20, 128, 128] 7,680 6 | ├─LSTM: 1-2 [906, 32] 20,736 7 | ├─Dropout: 1-3 [128, 32] -- 8 | ├─Linear: 1-4 [128, 18] 594 9 | ========================================================================================== 10 | Total params: 29,010 11 | Trainable params: 29,010 12 | Non-trainable params: 0 13 | Total mult-adds (M): 601.41 14 | ========================================================================================== 15 | Input size (kB): 20.55 16 | Forward/backward pass size (MB): 2.87 17 | Params size (kB): 116.04 18 | Estimated Total Size (MB): 3.01 19 | ========================================================================================== 20 | -------------------------------------------------------------------------------- /tests/test_output/parameter_list.out: -------------------------------------------------------------------------------- 1 | ======================================================================================================================================================================================================== 2 | Layer (type:depth-idx) Kernel Shape Groups Input Shape Output Shape Param # Param % Mult-Adds Trainable 3 | ======================================================================================================================================================================================================== 4 | ParameterListModel -- -- [100, 100] [100, 100] 30,000 100.00% -- True 5 | ======================================================================================================================================================================================================== 6 | Total params: 30,000 7 | Trainable params: 30,000 8 | Non-trainable params: 0 9 | Total mult-adds: 0 10 | ======================================================================================================================================================================================================== 11 | Input size (kB): 40.07 12 | Forward/backward pass size: 0 13 | Params size: 0 14 | Estimated Total Size (kB): 40.07 15 | ======================================================================================================================================================================================================== 16 | -------------------------------------------------------------------------------- /tests/test_output/parameters_with_other_layers.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | ParameterFCNet [3, 32] 8,256 5 | ├─a ├─8,192 6 | ├─b └─64 7 | ├─Linear: 1-1 [3, 32] 2,080 8 | │ └─weight ├─2,048 9 | │ └─bias └─32 10 | ========================================================================================== 11 | Total params: 10,336 12 | Trainable params: 10,336 13 | Non-trainable params: 0 14 | Total mult-adds (k): 6.24 15 | ========================================================================================== 16 | Input size (kB): 1.61 17 | Forward/backward pass size: 768 18 | Params size (kB): 8.32 19 | Estimated Total Size (kB): 10.70 20 | ========================================================================================== 21 | ========================================================================================== 22 | Layer (type:depth-idx) Output Shape Param # 23 | ========================================================================================== 24 | ParameterFCNet [3, 64] 8,256 25 | ├─a ├─8,192 26 | ├─b └─64 27 | ========================================================================================== 28 | Total params: 8,256 29 | Trainable params: 8,256 30 | Non-trainable params: 0 31 | Total mult-adds: 0 32 | ========================================================================================== 33 | Input size (kB): 1.61 34 | Forward/backward pass size (kB): 1.54 35 | Params size (kB): 33.02 36 | Estimated Total Size (kB): 36.17 37 | ========================================================================================== 38 | -------------------------------------------------------------------------------- /tests/test_output/partial_jit.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | PartialJITModel -- -- 5 | ├─Conv2d: 1-1 -- 260 6 | ├─Conv2d: 1-2 -- 5,020 7 | ├─Dropout2d: 1-3 -- -- 8 | ├─Linear: 1-4 -- 16,050 9 | ├─Linear: 1-5 -- 510 10 | ========================================================================================== 11 | Total params: 21,840 12 | Trainable params: 21,840 13 | Non-trainable params: 0 14 | Total mult-adds: 0 15 | ========================================================================================== 16 | Input size (kB): 6.34 17 | Forward/backward pass size: 0 18 | Params size (kB): 87.36 19 | Estimated Total Size (kB): 93.70 20 | ========================================================================================== 21 | -------------------------------------------------------------------------------- /tests/test_output/pruned_adversary.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | PrunedLayerNameModel [4, 8] 32 5 | ========================================================================================== 6 | Total params: 32 7 | Trainable params: 32 8 | Non-trainable params: 0 9 | Total mult-adds (k): 1.02 10 | ========================================================================================== 11 | Input size: 76 12 | Forward/backward pass size: 256 13 | Params size: 128 14 | Estimated Total Size: 460 15 | ========================================================================================== 16 | ========================================================================================== 17 | Layer (type:depth-idx) Output Shape Param # 18 | ========================================================================================== 19 | FakePrunedLayerModel [4, 8] 32 20 | ========================================================================================== 21 | Total params: 32 22 | Trainable params: 32 23 | Non-trainable params: 0 24 | Total mult-adds (k): 1.02 25 | ========================================================================================== 26 | Input size: 76 27 | Forward/backward pass size: 256 28 | Params size: 128 29 | Estimated Total Size: 460 30 | ========================================================================================== 31 | -------------------------------------------------------------------------------- /tests/test_output/pruning.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | SingleInputNet [16, 10] -- 5 | ├─Conv2d: 1-1 [16, 10, 24, 24] 135 6 | ├─Conv2d: 1-2 [16, 20, 8, 8] 2,520 7 | ├─Dropout2d: 1-3 [16, 20, 8, 8] -- 8 | ├─Linear: 1-4 [16, 50] 8,050 9 | ├─Linear: 1-5 [16, 10] 260 10 | ========================================================================================== 11 | Total params: 10,965 12 | Trainable params: 10,965 13 | Non-trainable params: 0 14 | Total mult-adds (M): 3.96 15 | ========================================================================================== 16 | Input size (kB): 50.25 17 | Forward/backward pass size (kB): 908.80 18 | Params size (kB): 43.86 19 | Estimated Total Size (MB): 1.00 20 | ========================================================================================== 21 | -------------------------------------------------------------------------------- /tests/test_output/recursive.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | RecursiveNet [1, 64, 28, 28] -- 5 | ├─Conv2d: 1-1 [1, 64, 28, 28] 36,928 6 | ├─Conv2d: 1-2 [1, 64, 28, 28] (recursive) 7 | ├─Conv2d: 1-3 [1, 64, 28, 28] (recursive) 8 | ├─Conv2d: 1-4 [1, 64, 28, 28] (recursive) 9 | ├─Conv2d: 1-5 [1, 64, 28, 28] (recursive) 10 | ├─Conv2d: 1-6 [1, 64, 28, 28] (recursive) 11 | ========================================================================================== 12 | Total params: 36,928 13 | Trainable params: 36,928 14 | Non-trainable params: 0 15 | Total mult-adds (M): 173.71 16 | ========================================================================================== 17 | Input size (kB): 200.78 18 | Forward/backward pass size (MB): 2.41 19 | Params size (kB): 147.71 20 | Estimated Total Size (MB): 2.76 21 | ========================================================================================== 22 | -------------------------------------------------------------------------------- /tests/test_output/recursive_with_missing_layers.out: -------------------------------------------------------------------------------- 1 | ============================================================================================================== 2 | Layer (type (var_name):depth-idx) Output Shape Param # 3 | ============================================================================================================== 4 | RecursiveWithMissingLayers (RecursiveWithMissingLayers) [2, 1, 128, 128] -- 5 | ├─Conv2d (out_conv0): 1-1 [2, 8, 128, 128] 608 6 | ├─BatchNorm2d (out_bn0): 1-2 [2, 8, 128, 128] 16 7 | ├─ModuleDict (block0): 1-3 -- -- 8 | │ └─Conv2d (in_conv1): 2-1 [2, 8, 128, 128] 584 9 | │ └─BatchNorm2d (in_bn1): 2-2 [2, 8, 128, 128] 16 10 | │ └─Conv2d (in_conv2): 2-3 [2, 8, 128, 128] 584 11 | │ └─BatchNorm2d (in_bn2): 2-4 [2, 8, 128, 128] 16 12 | │ └─Conv2d (in_conv3): 2-5 [2, 8, 128, 128] 584 13 | │ └─BatchNorm2d (in_bn3): 2-6 [2, 8, 128, 128] 16 14 | ├─ModuleDict (block1): 1-4 -- -- 15 | │ └─Conv2d (in_conv4): 2-7 [2, 8, 128, 128] 584 16 | │ └─BatchNorm2d (in_bn4): 2-8 [2, 8, 128, 128] 16 17 | │ └─Conv2d (in_conv5): 2-9 [2, 8, 128, 128] 584 18 | │ └─BatchNorm2d (in_bn5): 2-10 [2, 8, 128, 128] 16 19 | │ └─Conv2d (in_conv6): 2-11 [2, 8, 128, 128] 584 20 | │ └─BatchNorm2d (in_bn6): 2-12 [2, 8, 128, 128] 16 21 | ├─Conv2d (out_conv7): 1-5 [2, 1, 128, 128] 9 22 | ├─BatchNorm2d (out_bn7): 1-6 [2, 1, 128, 128] 2 23 | ============================================================================================================== 24 | Total params: 4,235 25 | Trainable params: 4,235 26 | Non-trainable params: 0 27 | Total mult-adds (M): 135.04 28 | ============================================================================================================== 29 | Input size (kB): 393.29 30 | Forward/backward pass size (MB): 29.88 31 | Params size (kB): 16.94 32 | Estimated Total Size (MB): 30.29 33 | ============================================================================================================== 34 | -------------------------------------------------------------------------------- /tests/test_output/register_parameter.out: -------------------------------------------------------------------------------- 1 | ================================================================= 2 | Layer (type:depth-idx) Param # 3 | ================================================================= 4 | RegisterParameter 2 5 | ├─Linear: 1-1 6 6 | ├─Linear: 1-2 6 7 | ================================================================= 8 | Total params: 14 9 | Trainable params: 14 10 | Non-trainable params: 0 11 | ================================================================= 12 | -------------------------------------------------------------------------------- /tests/test_output/resnet18_depth_consistency.out: -------------------------------------------------------------------------------- 1 | =================================================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # Param % 3 | =================================================================================================================== 4 | ResNet [1, 1000] -- -- 5 | ├─Conv2d: 1-1 [1, 64, 32, 32] 9,408 0.08% 6 | ├─BatchNorm2d: 1-2 [1, 64, 32, 32] 128 0.00% 7 | ├─ReLU: 1-3 [1, 64, 32, 32] -- -- 8 | ├─MaxPool2d: 1-4 [1, 64, 16, 16] -- -- 9 | ├─Sequential: 1-5 [1, 64, 16, 16] 147,968 1.27% 10 | ├─Sequential: 1-6 [1, 128, 8, 8] 525,568 4.50% 11 | ├─Sequential: 1-7 [1, 256, 4, 4] 2,099,712 17.96% 12 | ├─Sequential: 1-8 [1, 512, 2, 2] 8,393,728 71.81% 13 | ├─AdaptiveAvgPool2d: 1-9 [1, 512, 1, 1] -- -- 14 | ├─Linear: 1-10 [1, 1000] 513,000 4.39% 15 | =================================================================================================================== 16 | Total params: 11,689,512 17 | Trainable params: 11,689,512 18 | Non-trainable params: 0 19 | Total mult-adds (M): 148.57 20 | =================================================================================================================== 21 | Input size (kB): 49.22 22 | Forward/backward pass size (MB): 3.25 23 | Params size (MB): 46.76 24 | Estimated Total Size (MB): 50.06 25 | =================================================================================================================== 26 | =================================================================================================================== 27 | Layer (type:depth-idx) Output Shape Param # Param % 28 | =================================================================================================================== 29 | ResNet [1, 1000] -- -- 30 | ├─Conv2d: 1-1 [1, 64, 32, 32] 9,408 0.08% 31 | ├─BatchNorm2d: 1-2 [1, 64, 32, 32] 128 0.00% 32 | ├─ReLU: 1-3 [1, 64, 32, 32] -- -- 33 | ├─MaxPool2d: 1-4 [1, 64, 16, 16] -- -- 34 | ├─Sequential: 1-5 [1, 64, 16, 16] -- -- 35 | │ └─BasicBlock: 2-1 [1, 64, 16, 16] 73,984 0.63% 36 | │ └─BasicBlock: 2-2 [1, 64, 16, 16] 73,984 0.63% 37 | ├─Sequential: 1-6 [1, 128, 8, 8] -- -- 38 | │ └─BasicBlock: 2-3 [1, 128, 8, 8] 230,144 1.97% 39 | │ └─BasicBlock: 2-4 [1, 128, 8, 8] 295,424 2.53% 40 | ├─Sequential: 1-7 [1, 256, 4, 4] -- -- 41 | │ └─BasicBlock: 2-5 [1, 256, 4, 4] 919,040 7.86% 42 | │ └─BasicBlock: 2-6 [1, 256, 4, 4] 1,180,672 10.10% 43 | ├─Sequential: 1-8 [1, 512, 2, 2] -- -- 44 | │ └─BasicBlock: 2-7 [1, 512, 2, 2] 3,673,088 31.42% 45 | │ └─BasicBlock: 2-8 [1, 512, 2, 2] 4,720,640 40.38% 46 | ├─AdaptiveAvgPool2d: 1-9 [1, 512, 1, 1] -- -- 47 | ├─Linear: 1-10 [1, 1000] 513,000 4.39% 48 | =================================================================================================================== 49 | Total params: 11,689,512 50 | Trainable params: 11,689,512 51 | Non-trainable params: 0 52 | Total mult-adds (M): 148.57 53 | =================================================================================================================== 54 | Input size (kB): 49.22 55 | Forward/backward pass size (MB): 3.25 56 | Params size (MB): 46.76 57 | Estimated Total Size (MB): 50.06 58 | =================================================================================================================== 59 | -------------------------------------------------------------------------------- /tests/test_output/return_dict.out: -------------------------------------------------------------------------------- 1 | ========================================================================================================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================================================================================================== 4 | ReturnDict [[1, 10]] -- 5 | ├─ReturnDictLayer: 1-1 [[1, 10]] -- 6 | │ └─Conv2d: 2-1 [1, 10, 24, 24] 260 7 | │ └─Conv2d: 2-2 [1, 20, 8, 8] 5,020 8 | │ └─Linear: 2-3 [1, 50] 16,050 9 | │ └─Linear: 2-4 [1, 10] 510 10 | ========================================================================================================================================================================== 11 | Total params: 21,840 12 | Trainable params: 21,840 13 | Non-trainable params: 0 14 | Total mult-adds (k): 487.60 15 | ========================================================================================================================================================================== 16 | Input size (kB): 3.33 17 | Forward/backward pass size (kB): 56.80 18 | Params size (kB): 87.36 19 | Estimated Total Size (kB): 147.49 20 | ========================================================================================================================================================================== 21 | -------------------------------------------------------------------------------- /tests/test_output/reusing_activation_layers.out: -------------------------------------------------------------------------------- 1 | ================================================================= 2 | Layer (type:depth-idx) Param # 3 | ================================================================= 4 | Sequential -- 5 | ├─LeakyReLU: 1-1 -- 6 | ├─Identity: 1-2 -- 7 | ├─LeakyReLU: 1-3 -- 8 | ├─Identity: 1-4 -- 9 | ├─LeakyReLU: 1-5 -- 10 | ================================================================= 11 | Total params: 0 12 | Trainable params: 0 13 | Non-trainable params: 0 14 | ================================================================= 15 | ================================================================= 16 | Layer (type:depth-idx) Param # 17 | ================================================================= 18 | Sequential -- 19 | ├─LeakyReLU: 1-1 -- 20 | ├─Identity: 1-2 -- 21 | ├─LeakyReLU: 1-3 -- 22 | ├─Identity: 1-4 -- 23 | ├─LeakyReLU: 1-5 -- 24 | ================================================================= 25 | Total params: 0 26 | Trainable params: 0 27 | Non-trainable params: 0 28 | ================================================================= 29 | -------------------------------------------------------------------------------- /tests/test_output/row_settings.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type (var_name)) Output Shape Param # 3 | ========================================================================================== 4 | SingleInputNet (SingleInputNet) [16, 10] -- 5 | ├─Conv2d (conv1) [16, 10, 24, 24] 260 6 | ├─Conv2d (conv2) [16, 20, 8, 8] 5,020 7 | ├─Dropout2d (conv2_drop) [16, 20, 8, 8] -- 8 | ├─Linear (fc1) [16, 50] 16,050 9 | ├─Linear (fc2) [16, 10] 510 10 | ========================================================================================== 11 | Total params: 21,840 12 | Trainable params: 21,840 13 | Non-trainable params: 0 14 | Total mult-adds (M): 7.80 15 | ========================================================================================== 16 | Input size (kB): 50.25 17 | Forward/backward pass size (kB): 908.80 18 | Params size (kB): 87.36 19 | Estimated Total Size (MB): 1.05 20 | ========================================================================================== 21 | -------------------------------------------------------------------------------- /tests/test_output/siamese_net.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | SiameseNets [1, 1] -- 5 | ├─Conv2d: 1-1 [1, 64, 79, 79] 6,464 6 | ├─MaxPool2d: 1-2 [1, 64, 39, 39] -- 7 | ├─Conv2d: 1-3 [1, 128, 33, 33] 401,536 8 | ├─MaxPool2d: 1-4 [1, 128, 16, 16] -- 9 | ├─Conv2d: 1-5 [1, 128, 13, 13] 262,272 10 | ├─MaxPool2d: 1-6 [1, 128, 6, 6] -- 11 | ├─Conv2d: 1-7 [1, 256, 3, 3] 524,544 12 | ├─MaxPool2d: 1-8 [1, 256, 1, 1] -- 13 | ├─Conv2d: 1-9 [1, 64, 79, 79] (recursive) 14 | ├─MaxPool2d: 1-10 [1, 64, 39, 39] -- 15 | ├─Conv2d: 1-11 [1, 128, 33, 33] (recursive) 16 | ├─MaxPool2d: 1-12 [1, 128, 16, 16] -- 17 | ├─Conv2d: 1-13 [1, 128, 13, 13] (recursive) 18 | ├─MaxPool2d: 1-14 [1, 128, 6, 6] -- 19 | ├─Conv2d: 1-15 [1, 256, 3, 3] (recursive) 20 | ├─MaxPool2d: 1-16 [1, 256, 1, 1] -- 21 | ├─Linear: 1-17 [1, 4096] 1,052,672 22 | ├─Linear: 1-18 [1, 4096] (recursive) 23 | ├─Dropout: 1-19 [1, 4096] -- 24 | ├─Linear: 1-20 [1, 1] 4,097 25 | ========================================================================================== 26 | Total params: 2,251,585 27 | Trainable params: 2,251,585 28 | Non-trainable params: 0 29 | Total mult-adds (G): 1.06 30 | ========================================================================================== 31 | Input size (kB): 62.10 32 | Forward/backward pass size (MB): 9.07 33 | Params size (MB): 9.01 34 | Estimated Total Size (MB): 18.14 35 | ========================================================================================== 36 | -------------------------------------------------------------------------------- /tests/test_output/single_input.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | SingleInputNet [2, 10] -- 5 | ├─Conv2d: 1-1 [2, 10, 24, 24] 260 6 | ├─Conv2d: 1-2 [2, 20, 8, 8] 5,020 7 | ├─Dropout2d: 1-3 [2, 20, 8, 8] -- 8 | ├─Linear: 1-4 [2, 50] 16,050 9 | ├─Linear: 1-5 [2, 10] 510 10 | ========================================================================================== 11 | Total params: 21,840 12 | Trainable params: 21,840 13 | Non-trainable params: 0 14 | Total mult-adds (k): 975.20 15 | ========================================================================================== 16 | Input size (kB): 6.34 17 | Forward/backward pass size (kB): 113.60 18 | Params size (kB): 87.36 19 | Estimated Total Size (kB): 207.30 20 | ========================================================================================== 21 | -------------------------------------------------------------------------------- /tests/test_output/single_input_all_cols.out: -------------------------------------------------------------------------------- 1 | ======================================================================================================================================================================================================== 2 | Layer (type:depth-idx) Kernel Shape Groups Input Shape Output Shape Param # Param % Mult-Adds Trainable 3 | ======================================================================================================================================================================================================== 4 | SingleInputNet -- -- [7, 1, 28, 28] [7, 10] -- -- -- True 5 | ├─Conv2d: 1-1 [5, 5] 1 [7, 1, 28, 28] [7, 10, 24, 24] 260 1.19% 1,048,320 True 6 | ├─Conv2d: 1-2 [5, 5] 1 [7, 10, 12, 12] [7, 20, 8, 8] 5,020 22.99% 2,248,960 True 7 | ├─Dropout2d: 1-3 -- -- [7, 20, 8, 8] [7, 20, 8, 8] -- -- -- -- 8 | ├─Linear: 1-4 -- -- [7, 320] [7, 50] 16,050 73.49% 112,350 True 9 | ├─Linear: 1-5 -- -- [7, 50] [7, 10] 510 2.34% 3,570 True 10 | ======================================================================================================================================================================================================== 11 | Total params: 21,840 12 | Trainable params: 21,840 13 | Non-trainable params: 0 14 | Total mult-adds (M): 3.41 15 | ======================================================================================================================================================================================================== 16 | Input size (kB): 22.02 17 | Forward/backward pass size (kB): 397.60 18 | Params size (kB): 87.36 19 | Estimated Total Size (kB): 506.98 20 | ======================================================================================================================================================================================================== 21 | -------------------------------------------------------------------------------- /tests/test_output/single_input_batch_dim.out: -------------------------------------------------------------------------------- 1 | ============================================================================================================================================ 2 | Layer (type:depth-idx) Kernel Shape Input Shape Output Shape Param # Mult-Adds 3 | ============================================================================================================================================ 4 | SingleInputNet -- [1, 1, 28, 28] [1, 10] -- -- 5 | ├─Conv2d: 1-1 [5, 5] [1, 1, 28, 28] [1, 10, 24, 24] 260 149,760 6 | ├─Conv2d: 1-2 [5, 5] [1, 10, 12, 12] [1, 20, 8, 8] 5,020 321,280 7 | ├─Dropout2d: 1-3 -- [1, 20, 8, 8] [1, 20, 8, 8] -- -- 8 | ├─Linear: 1-4 -- [1, 320] [1, 50] 16,050 16,050 9 | ├─Linear: 1-5 -- [1, 50] [1, 10] 510 510 10 | ============================================================================================================================================ 11 | Total params: 21,840 12 | Trainable params: 21,840 13 | Non-trainable params: 0 14 | Total mult-adds (k): 487.60 15 | ============================================================================================================================================ 16 | Input size (kB): 3.21 17 | Forward/backward pass size (kB): 56.80 18 | Params size (kB): 87.36 19 | Estimated Total Size (kB): 147.37 20 | ============================================================================================================================================ 21 | -------------------------------------------------------------------------------- /tests/test_output/single_input_half.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | SingleInputNet [2, 10] -- 5 | ├─Conv2d: 1-1 [2, 10, 24, 24] 260 6 | ├─Conv2d: 1-2 [2, 20, 8, 8] 5,020 7 | ├─Dropout2d: 1-3 [2, 20, 8, 8] -- 8 | ├─Linear: 1-4 [2, 50] 16,050 9 | ├─Linear: 1-5 [2, 10] 510 10 | ========================================================================================== 11 | Total params: 21,840 12 | Trainable params: 21,840 13 | Non-trainable params: 0 14 | Total mult-adds (M): 0.98 15 | ========================================================================================== 16 | Input size (MB): 0.00 17 | Forward/backward pass size (MB): 0.06 18 | Params size (MB): 0.04 19 | Estimated Total Size (MB): 0.10 20 | ========================================================================================== 21 | -------------------------------------------------------------------------------- /tests/test_output/single_layer_network_on_gpu.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | Linear [1, 5] 15 5 | ========================================================================================== 6 | Total params: 15 7 | Trainable params: 15 8 | Non-trainable params: 0 9 | Total mult-adds (M): 0.00 10 | ========================================================================================== 11 | Input size (MB): 0.00 12 | Forward/backward pass size (MB): 0.00 13 | Params size (MB): 0.00 14 | Estimated Total Size (MB): 0.00 15 | ========================================================================================== 16 | -------------------------------------------------------------------------------- /tests/test_output/single_layer_network_on_gpu_device.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | Linear [1, 5] 15 5 | ========================================================================================== 6 | Total params: 15 7 | Trainable params: 15 8 | Non-trainable params: 0 9 | Total mult-adds (M): 0.00 10 | ========================================================================================== 11 | Input size (MB): 0.00 12 | Forward/backward pass size (MB): 0.00 13 | Params size (MB): 0.00 14 | Estimated Total Size (MB): 0.00 15 | ========================================================================================== 16 | -------------------------------------------------------------------------------- /tests/test_output/single_linear_layer.out: -------------------------------------------------------------------------------- 1 | ================================================================= 2 | Layer (type:depth-idx) Param # 3 | ================================================================= 4 | Linear 15 5 | ================================================================= 6 | Total params: 15 7 | Trainable params: 15 8 | Non-trainable params: 0 9 | ================================================================= 10 | ========================================================================================== 11 | Layer (type:depth-idx) Output Shape Param # 12 | ========================================================================================== 13 | Linear [1, 5] 15 14 | ========================================================================================== 15 | Total params: 15 16 | Trainable params: 15 17 | Non-trainable params: 0 18 | Total mult-adds: 15 19 | ========================================================================================== 20 | Input size: 80 21 | Forward/backward pass size: 40 22 | Params size: 60 23 | Estimated Total Size: 180 24 | ========================================================================================== 25 | -------------------------------------------------------------------------------- /tests/test_output/single_parameter_model.out: -------------------------------------------------------------------------------- 1 | ================================================================= 2 | Layer (type:depth-idx) Param # 3 | ================================================================= 4 | ParameterA 1,024 5 | ================================================================= 6 | Total params: 1,024 7 | Trainable params: 1,024 8 | Non-trainable params: 0 9 | ================================================================= 10 | ================================================================= 11 | Layer (type:depth-idx) Param # 12 | ================================================================= 13 | ParameterB 1,024 14 | ├─Conv2d: 1-1 168 15 | ================================================================= 16 | Total params: 1,192 17 | Trainable params: 1,192 18 | Non-trainable params: 0 19 | ================================================================= 20 | ================================================================= 21 | Layer (type:depth-idx) Param # 22 | ================================================================= 23 | ParameterA 1,024 24 | ├─w └─1,024 25 | ================================================================= 26 | Total params: 1,024 27 | Trainable params: 1,024 28 | Non-trainable params: 0 29 | ================================================================= 30 | ================================================================= 31 | Layer (type:depth-idx) Param # 32 | ================================================================= 33 | ParameterB 1,024 34 | ├─w └─1,024 35 | ├─Conv2d: 1-1 168 36 | │ └─weight ├─162 37 | │ └─bias └─6 38 | ================================================================= 39 | Total params: 1,192 40 | Trainable params: 1,192 41 | Non-trainable params: 0 42 | ================================================================= 43 | -------------------------------------------------------------------------------- /tests/test_output/string_result.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | SingleInputNet [16, 10] -- 5 | ├─Conv2d: 1-1 [16, 10, 24, 24] 260 6 | ├─Conv2d: 1-2 [16, 20, 8, 8] 5,020 7 | ├─Dropout2d: 1-3 [16, 20, 8, 8] -- 8 | ├─Linear: 1-4 [16, 50] 16,050 9 | ├─Linear: 1-5 [16, 10] 510 10 | ========================================================================================== 11 | Total params: 21,840 12 | Trainable params: 21,840 13 | Non-trainable params: 0 14 | Total mult-adds (M): 7.80 15 | ========================================================================================== 16 | Input size (kB): 50.25 17 | Forward/backward pass size (kB): 908.80 18 | Params size (kB): 87.36 19 | Estimated Total Size (MB): 1.05 20 | ========================================================================================== 21 | -------------------------------------------------------------------------------- /tests/test_output/too_many_linear.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | ReuseLinear [2, 10] -- 5 | ├─Sequential: 1-1 [2, 10] -- 6 | │ └─Linear: 2-1 [2, 10] 110 7 | │ └─ReLU: 2-2 [2, 10] -- 8 | │ └─Linear: 2-3 [2, 10] (recursive) 9 | │ └─ReLU: 2-4 [2, 10] -- 10 | │ └─Linear: 2-5 [2, 10] (recursive) 11 | │ └─ReLU: 2-6 [2, 10] -- 12 | │ └─Linear: 2-7 [2, 10] (recursive) 13 | │ └─ReLU: 2-8 [2, 10] -- 14 | ========================================================================================== 15 | Total params: 110 16 | Trainable params: 110 17 | Non-trainable params: 0 18 | Total mult-adds: 880 19 | ========================================================================================== 20 | Input size: 152 21 | Forward/backward pass size: 640 22 | Params size: 440 23 | Estimated Total Size (kB): 1.23 24 | ========================================================================================== 25 | -------------------------------------------------------------------------------- /tests/test_output/too_many_linear_plus_existing_hooks.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | ReuseLinearExtended [2, 10] -- 5 | ├─Sequential: 1-1 [2, 10] -- 6 | │ └─Linear: 2-1 [2, 10] 110 7 | │ └─ReLU: 2-2 [2, 10] -- 8 | │ └─Linear: 2-3 [2, 10] (recursive) 9 | │ └─ReLU: 2-4 [2, 10] -- 10 | │ └─Linear: 2-5 [2, 10] (recursive) 11 | │ └─ReLU: 2-6 [2, 10] -- 12 | │ └─Linear: 2-7 [2, 10] (recursive) 13 | │ └─ReLU: 2-8 [2, 10] -- 14 | ========================================================================================== 15 | Total params: 110 16 | Trainable params: 110 17 | Non-trainable params: 0 18 | Total mult-adds: 880 19 | ========================================================================================== 20 | Input size: 152 21 | Forward/backward pass size: 640 22 | Params size: 440 23 | Estimated Total Size (kB): 1.23 24 | ========================================================================================== 25 | ========================================================================================== 26 | Layer (type:depth-idx) Output Shape Param # 27 | ========================================================================================== 28 | ReuseLinearExtended [2, 10] -- 29 | ├─Sequential: 1-1 [2, 10] -- 30 | │ └─Linear: 2-1 [2, 10] 110 31 | │ └─ReLU: 2-2 [2, 10] -- 32 | │ └─Linear: 2-3 [2, 10] (recursive) 33 | │ └─ReLU: 2-4 [2, 10] -- 34 | │ └─Linear: 2-5 [2, 10] (recursive) 35 | │ └─ReLU: 2-6 [2, 10] -- 36 | │ └─Linear: 2-7 [2, 10] (recursive) 37 | │ └─ReLU: 2-8 [2, 10] -- 38 | ========================================================================================== 39 | Total params: 110 40 | Trainable params: 110 41 | Non-trainable params: 0 42 | Total mult-adds: 880 43 | ========================================================================================== 44 | Input size: 152 45 | Forward/backward pass size: 640 46 | Params size: 440 47 | Estimated Total Size (kB): 1.23 48 | ========================================================================================== 49 | -------------------------------------------------------------------------------- /tests/test_output/too_many_relus.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | ReuseReLU [4, 8, 11, 11] -- 5 | ├─Sequential: 1-1 [4, 8, 11, 11] -- 6 | │ └─ReflectionPad2d: 2-1 [4, 4, 70, 70] -- 7 | │ └─Conv2d: 2-2 [4, 1, 70, 70] 5 8 | │ └─BatchNorm2d: 2-3 [4, 1, 70, 70] 2 9 | │ └─ReLU: 2-4 [4, 1, 70, 70] -- 10 | │ └─Conv2d: 2-5 [4, 2, 36, 36] 4 11 | │ └─BatchNorm2d: 2-6 [4, 2, 36, 36] 4 12 | │ └─ReLU: 2-7 [4, 2, 36, 36] -- 13 | │ └─Conv2d: 2-8 [4, 4, 19, 19] 12 14 | │ └─BatchNorm2d: 2-9 [4, 4, 19, 19] 8 15 | │ └─ReLU: 2-10 [4, 4, 19, 19] -- 16 | │ └─Conv2d: 2-11 [4, 8, 11, 11] 40 17 | │ └─BatchNorm2d: 2-12 [4, 8, 11, 11] 16 18 | │ └─ReLU: 2-13 [4, 8, 11, 11] -- 19 | ========================================================================================== 20 | Total params: 91 21 | Trainable params: 91 22 | Non-trainable params: 0 23 | Total mult-adds (k): 155.54 24 | ========================================================================================== 25 | Input size (kB): 262.22 26 | Forward/backward pass size (kB): 633.86 27 | Params size: 364 28 | Estimated Total Size (kB): 896.44 29 | ========================================================================================== 30 | -------------------------------------------------------------------------------- /tests/test_output/trainable_column.out: -------------------------------------------------------------------------------- 1 | ============================================================================================================================================ 2 | Layer (type:depth-idx) Kernel Shape Input Shape Output Shape Trainable 3 | ============================================================================================================================================ 4 | MixedTrainable -- [1, 1, 1] [1, 1, 1] Partial 5 | ├─Conv1d: 1-1 [1] [1, 1, 1] [1, 1, 1] True 6 | ├─Conv1d: 1-2 [1] [1, 1, 1] [1, 1, 1] Partial 7 | ├─Conv1d: 1-3 [1] [1, 1, 1] [1, 1, 1] False 8 | ├─Dropout: 1-4 -- [1, 1, 1] [1, 1, 1] -- 9 | ============================================================================================================================================ 10 | Total params: 6 11 | Trainable params: 3 12 | Non-trainable params: 3 13 | Total mult-adds: 6 14 | ============================================================================================================================================ 15 | Input size: 76 16 | Forward/backward pass size: 24 17 | Params size: 24 18 | Estimated Total Size: 124 19 | ============================================================================================================================================ 20 | -------------------------------------------------------------------------------- /tests/test_output/training_mode.out: -------------------------------------------------------------------------------- 1 | ========================================================================================== 2 | Layer (type:depth-idx) Output Shape Param # 3 | ========================================================================================== 4 | Linear [1, 5] 15 5 | ========================================================================================== 6 | Total params: 15 7 | Trainable params: 15 8 | Non-trainable params: 0 9 | Total mult-adds: 15 10 | ========================================================================================== 11 | Input size: 80 12 | Forward/backward pass size: 40 13 | Params size: 60 14 | Estimated Total Size: 180 15 | ========================================================================================== 16 | -------------------------------------------------------------------------------- /tests/test_output/uninitialized_tensor.out: -------------------------------------------------------------------------------- 1 | ================================================================= 2 | Layer (type:depth-idx) Param # 3 | ================================================================= 4 | UninitializedParameterModel -- 5 | ================================================================= 6 | Total params: 0 7 | Trainable params: 0 8 | Non-trainable params: 0 9 | ================================================================= 10 | ========================================================================================== 11 | Layer (type:depth-idx) Output Shape Param # 12 | ========================================================================================== 13 | UninitializedParameterModel [2, 2] 128 14 | ========================================================================================== 15 | Total params: 128 16 | Trainable params: 128 17 | Non-trainable params: 0 18 | Total mult-adds: 0 19 | ========================================================================================== 20 | Input size: 88 21 | Forward/backward pass size: 32 22 | Params size: 512 23 | Estimated Total Size: 632 24 | ========================================================================================== 25 | -------------------------------------------------------------------------------- /tests/torchinfo_xl_test.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import torch 3 | import torchvision # type: ignore[import-untyped] 4 | from packaging import version 5 | 6 | from tests.fixtures.genotype import GenotypeNetwork # type: ignore[attr-defined] 7 | from tests.fixtures.tmva_net import TMVANet # type: ignore[attr-defined] 8 | from torchinfo import summary 9 | from torchinfo.enums import ColumnSettings 10 | 11 | if version.parse(torch.__version__) >= version.parse("1.8"): 12 | from transformers import ( # type: ignore[import-untyped] 13 | AutoModelForSeq2SeqLM, 14 | ) 15 | 16 | 17 | def test_ascii_only() -> None: 18 | result = summary( 19 | torchvision.models.resnet18(), 20 | depth=3, 21 | input_size=(1, 3, 64, 64), 22 | row_settings=["ascii_only"], 23 | ) 24 | 25 | assert str(result).encode("ascii").decode("ascii") 26 | 27 | 28 | def test_frozen_layers() -> None: 29 | model = torchvision.models.resnet18() 30 | for ind, param in enumerate(model.parameters()): 31 | if ind < 30: 32 | param.requires_grad = False 33 | 34 | summary( 35 | model, 36 | input_size=(1, 3, 64, 64), 37 | depth=3, 38 | col_names=("output_size", "num_params", "kernel_size", "mult_adds"), 39 | ) 40 | 41 | 42 | def test_eval_order_doesnt_matter() -> None: 43 | # prevent mixed device if cuda is available 44 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 45 | 46 | input_size = (1, 3, 224, 224) 47 | input_tensor = torch.ones(input_size).to(device) 48 | 49 | model1 = torchvision.models.resnet18( 50 | weights=torchvision.models.ResNet18_Weights.DEFAULT 51 | ) 52 | model1.eval() 53 | summary(model1, input_size=input_size) 54 | with torch.inference_mode(): 55 | output1 = model1(input_tensor) 56 | 57 | model2 = torchvision.models.resnet18( 58 | weights=torchvision.models.ResNet18_Weights.DEFAULT 59 | ) 60 | summary(model2, input_size=input_size, mode="eval") 61 | model2.eval() 62 | with torch.inference_mode(): 63 | output2 = model2(input_tensor) 64 | 65 | assert torch.all(torch.eq(output1, output2)) 66 | 67 | 68 | def test_resnet18_depth_consistency() -> None: 69 | model = torchvision.models.resnet18() 70 | 71 | for depth in range(1, 3): 72 | summary( 73 | model, 74 | (1, 3, 64, 64), 75 | col_names=( 76 | ColumnSettings.OUTPUT_SIZE, 77 | ColumnSettings.NUM_PARAMS, 78 | ColumnSettings.PARAMS_PERCENT, 79 | ), 80 | depth=depth, 81 | cache_forward_pass=True, 82 | ) 83 | 84 | 85 | def test_resnet50() -> None: 86 | # According to https://arxiv.org/abs/1605.07146, 87 | # resnet50 has ~25.6 M trainable params. 88 | model = torchvision.models.resnet50() 89 | results = summary(model, input_size=(2, 3, 224, 224)) 90 | 91 | assert results.total_params == 25557032 # close to 25.6e6 92 | assert results.total_mult_adds == sum( 93 | layer.macs for layer in results.summary_list if layer.is_leaf_layer 94 | ) 95 | 96 | 97 | def test_resnet152() -> None: 98 | model = torchvision.models.resnet152() 99 | 100 | summary(model, (1, 3, 224, 224), depth=3) 101 | 102 | 103 | @pytest.mark.skip(reason="nondeterministic output") 104 | def test_fasterrcnn() -> None: 105 | model = torchvision.models.detection.fasterrcnn_resnet50_fpn( 106 | pretrained_backbone=False 107 | ) 108 | results = summary(model, input_size=(1, 3, 112, 112)) 109 | 110 | assert results.total_params == 41755286 111 | 112 | 113 | def test_genotype() -> None: 114 | model = GenotypeNetwork() 115 | 116 | x = summary(model, (2, 3, 32, 32), depth=3, cache_forward_pass=True) 117 | y = summary(model, (2, 3, 32, 32), depth=7, cache_forward_pass=True) 118 | 119 | assert x.total_params == y.total_params, (x, y) 120 | 121 | 122 | def test_tmva_net_column_totals() -> None: 123 | for depth in (1, 3, 5): 124 | results = summary( 125 | TMVANet(n_classes=4, n_frames=5), 126 | input_data=[ 127 | torch.randn(1, 1, 5, 256, 64), 128 | torch.randn(1, 1, 5, 256, 256), 129 | torch.randn(1, 1, 5, 256, 64), 130 | ], 131 | col_names=["output_size", "num_params", "mult_adds"], 132 | depth=depth, 133 | cache_forward_pass=True, 134 | ) 135 | 136 | assert results.total_params == sum( 137 | layer.num_params for layer in results.summary_list if layer.is_leaf_layer 138 | ) 139 | assert results.total_mult_adds == sum( 140 | layer.macs for layer in results.summary_list if layer.is_leaf_layer 141 | ) 142 | 143 | 144 | def test_google() -> None: 145 | google_net = torchvision.models.googlenet(init_weights=False) 146 | 147 | summary(google_net, (1, 3, 112, 112), depth=7, mode="eval") 148 | 149 | # Check googlenet in training mode since InceptionAux layers are used in 150 | # forward-prop in train mode but not in eval mode. 151 | summary(google_net, (1, 3, 112, 112), depth=7, mode="train") 152 | 153 | 154 | @pytest.mark.skipif( 155 | version.parse(torch.__version__) < version.parse("1.8"), 156 | reason="FlanT5Small only works for PyTorch v1.8 and above", 157 | ) 158 | def test_flan_t5_small() -> None: 159 | model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-small") 160 | inputs = { 161 | "input_ids": torch.zeros(3, 100).long(), 162 | "attention_mask": torch.zeros(3, 100).long(), 163 | "labels": torch.zeros(3, 100).long(), 164 | } 165 | summary(model, input_data=inputs) 166 | -------------------------------------------------------------------------------- /torchinfo/__init__.py: -------------------------------------------------------------------------------- 1 | from .enums import ColumnSettings, Mode, RowSettings, Units, Verbosity 2 | from .model_statistics import ModelStatistics 3 | from .torchinfo import summary 4 | 5 | __all__ = ( 6 | "ColumnSettings", 7 | "Mode", 8 | "ModelStatistics", 9 | "RowSettings", 10 | "Units", 11 | "Verbosity", 12 | "summary", 13 | ) 14 | __version__ = "1.8.0" 15 | -------------------------------------------------------------------------------- /torchinfo/enums.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | from enum import Enum, IntEnum, unique 4 | 5 | 6 | @unique 7 | class Mode(str, Enum): 8 | """Enum containing all model modes.""" 9 | 10 | __slots__ = () 11 | 12 | TRAIN = "train" 13 | EVAL = "eval" 14 | SAME = "same" 15 | 16 | 17 | @unique 18 | class RowSettings(str, Enum): 19 | """Enum containing all available row settings.""" 20 | 21 | __slots__ = () 22 | 23 | DEPTH = "depth" 24 | VAR_NAMES = "var_names" 25 | ASCII_ONLY = "ascii_only" 26 | HIDE_RECURSIVE_LAYERS = "hide_recursive_layers" 27 | 28 | 29 | @unique 30 | class ColumnSettings(str, Enum): 31 | """Enum containing all available column settings.""" 32 | 33 | __slots__ = () 34 | 35 | KERNEL_SIZE = "kernel_size" 36 | GROUPS = "groups" 37 | INPUT_SIZE = "input_size" 38 | OUTPUT_SIZE = "output_size" 39 | NUM_PARAMS = "num_params" 40 | PARAMS_PERCENT = "params_percent" 41 | MULT_ADDS = "mult_adds" 42 | TRAINABLE = "trainable" 43 | 44 | 45 | @unique 46 | class Units(str, Enum): 47 | """Enum containing all available bytes units.""" 48 | 49 | __slots__ = () 50 | 51 | AUTO = "auto" 52 | KILOBYTES = "k" 53 | MEGABYTES = "M" 54 | GIGABYTES = "G" 55 | TERABYTES = "T" 56 | NONE = "" 57 | 58 | 59 | @unique 60 | class Verbosity(IntEnum): 61 | """Contains verbosity levels.""" 62 | 63 | QUIET, DEFAULT, VERBOSE = 0, 1, 2 64 | -------------------------------------------------------------------------------- /torchinfo/formatting.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | import math 4 | from typing import TYPE_CHECKING, Any 5 | 6 | from .enums import ColumnSettings, RowSettings, Units, Verbosity 7 | 8 | if TYPE_CHECKING: 9 | from .layer_info import LayerInfo 10 | 11 | HEADER_TITLES = { 12 | ColumnSettings.KERNEL_SIZE: "Kernel Shape", 13 | ColumnSettings.GROUPS: "Groups", 14 | ColumnSettings.INPUT_SIZE: "Input Shape", 15 | ColumnSettings.OUTPUT_SIZE: "Output Shape", 16 | ColumnSettings.NUM_PARAMS: "Param #", 17 | ColumnSettings.PARAMS_PERCENT: "Param %", 18 | ColumnSettings.MULT_ADDS: "Mult-Adds", 19 | ColumnSettings.TRAINABLE: "Trainable", 20 | } 21 | CONVERSION_FACTORS = { 22 | Units.TERABYTES: 1e12, 23 | Units.GIGABYTES: 1e9, 24 | Units.MEGABYTES: 1e6, 25 | Units.KILOBYTES: 1e3, 26 | Units.NONE: 1, 27 | } 28 | 29 | 30 | class FormattingOptions: 31 | """Class that holds information about formatting the table output.""" 32 | 33 | def __init__( 34 | self, 35 | max_depth: int, 36 | verbose: int, 37 | col_names: tuple[ColumnSettings, ...], 38 | col_width: int, 39 | row_settings: set[RowSettings], 40 | ) -> None: 41 | self.max_depth = max_depth 42 | self.verbose = verbose 43 | self.col_names = col_names 44 | self.col_width = col_width 45 | self.row_settings = row_settings 46 | self.params_count_units = Units.NONE 47 | self.params_size_units = Units.AUTO 48 | self.macs_units = Units.AUTO 49 | 50 | self.layer_name_width = 40 51 | self.ascii_only = RowSettings.ASCII_ONLY in self.row_settings 52 | self.show_var_name = RowSettings.VAR_NAMES in self.row_settings 53 | self.show_depth = RowSettings.DEPTH in self.row_settings 54 | self.hide_recursive_layers = ( 55 | RowSettings.HIDE_RECURSIVE_LAYERS in self.row_settings 56 | ) 57 | 58 | @staticmethod 59 | def str_(val: Any) -> str: 60 | return str(val) if val else "--" 61 | 62 | def get_start_str(self, depth: int) -> str: 63 | """This function should handle all ascii/non-ascii-related characters.""" 64 | if depth == 0: 65 | return "" 66 | if depth == 1: 67 | return "+ " if self.ascii_only else "├─" 68 | return ("| " if self.ascii_only else "│ ") * (depth - 1) + ( 69 | "+ " if self.ascii_only else "└─" 70 | ) 71 | 72 | def set_layer_name_width( 73 | self, summary_list: list[LayerInfo], align_val: int = 5 74 | ) -> None: 75 | """ 76 | Set layer name width by taking the longest line length and rounding up to 77 | the nearest multiple of align_val. 78 | """ 79 | max_length = 0 80 | for info in summary_list: 81 | depth_indent = info.depth * align_val + 1 82 | layer_title = info.get_layer_name(self.show_var_name, self.show_depth) 83 | max_length = max(max_length, len(layer_title) + depth_indent) 84 | if max_length >= self.layer_name_width: 85 | self.layer_name_width = math.ceil(max_length / align_val) * align_val 86 | 87 | def get_total_width(self) -> int: 88 | """Calculate the total width of all lines in the table.""" 89 | return len(tuple(self.col_names)) * self.col_width + self.layer_name_width 90 | 91 | def format_row(self, layer_name: str, row_values: dict[ColumnSettings, str]) -> str: 92 | """Get the string representation of a single layer of the model.""" 93 | info_to_use = [row_values.get(row_type, "") for row_type in self.col_names] 94 | new_line = f"{layer_name:<{self.layer_name_width}} " 95 | for info in info_to_use: 96 | new_line += f"{info:<{self.col_width}} " 97 | return new_line.rstrip() + "\n" 98 | 99 | def header_row(self) -> str: 100 | layer_header = "" 101 | if self.show_var_name: 102 | layer_header += " (var_name)" 103 | if self.show_depth: 104 | layer_header += ":depth-idx" 105 | return self.format_row(f"Layer (type{layer_header})", HEADER_TITLES) 106 | 107 | def layer_info_to_row( 108 | self, layer_info: LayerInfo, reached_max_depth: bool, total_params: int 109 | ) -> str: 110 | """Convert layer_info to string representation of a row.""" 111 | values_for_row = { 112 | ColumnSettings.KERNEL_SIZE: self.str_(layer_info.kernel_size), 113 | ColumnSettings.GROUPS: self.str_(layer_info.groups), 114 | ColumnSettings.INPUT_SIZE: self.str_(layer_info.input_size), 115 | ColumnSettings.OUTPUT_SIZE: self.str_(layer_info.output_size), 116 | ColumnSettings.NUM_PARAMS: layer_info.num_params_to_str(reached_max_depth), 117 | ColumnSettings.PARAMS_PERCENT: layer_info.params_percent( 118 | total_params, reached_max_depth 119 | ), 120 | ColumnSettings.MULT_ADDS: layer_info.macs_to_str(reached_max_depth), 121 | ColumnSettings.TRAINABLE: self.str_(layer_info.trainable), 122 | } 123 | start_str = self.get_start_str(layer_info.depth) 124 | layer_name = layer_info.get_layer_name(self.show_var_name, self.show_depth) 125 | new_line = self.format_row(f"{start_str}{layer_name}", values_for_row) 126 | 127 | if self.verbose == Verbosity.VERBOSE: 128 | for inner_name, inner_layer_info in layer_info.inner_layers.items(): 129 | prefix = self.get_start_str(layer_info.depth + 1) 130 | new_line += self.format_row(f"{prefix}{inner_name}", inner_layer_info) 131 | return new_line 132 | 133 | def layers_to_str(self, summary_list: list[LayerInfo], total_params: int) -> str: 134 | """ 135 | Print each layer of the model using only current layer info. 136 | Container modules are already dealt with in add_missing_container_layers. 137 | """ 138 | new_str = "" 139 | for layer_info in summary_list: 140 | if layer_info.depth > self.max_depth or ( 141 | self.hide_recursive_layers and layer_info.is_recursive 142 | ): 143 | continue 144 | 145 | reached_max_depth = layer_info.depth == self.max_depth 146 | new_str += self.layer_info_to_row( 147 | layer_info, reached_max_depth, total_params 148 | ) 149 | return new_str 150 | -------------------------------------------------------------------------------- /torchinfo/model_statistics.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | from typing import TYPE_CHECKING, Any 4 | 5 | from .enums import Units 6 | from .formatting import CONVERSION_FACTORS, FormattingOptions 7 | 8 | if TYPE_CHECKING: 9 | from .layer_info import LayerInfo 10 | 11 | 12 | class ModelStatistics: 13 | """Class for storing results of the summary.""" 14 | 15 | def __init__( 16 | self, 17 | summary_list: list[LayerInfo], 18 | input_size: Any, 19 | total_input_size: int, 20 | formatting: FormattingOptions, 21 | ) -> None: 22 | self.summary_list = summary_list 23 | self.input_size = input_size 24 | self.formatting = formatting 25 | self.total_input = total_input_size 26 | self.total_mult_adds = 0 27 | self.total_params, self.trainable_params = 0, 0 28 | self.total_param_bytes, self.total_output_bytes = 0, 0 29 | 30 | # TODO: Figure out why the below functions using max() are ever 0 31 | # (they should always be non-negative), and remove the call to max(). 32 | # Investigation: https://github.com/TylerYep/torchinfo/pull/195 33 | for layer_info in summary_list: 34 | if layer_info.is_leaf_layer: 35 | self.total_mult_adds += layer_info.macs 36 | if layer_info.num_params > 0: 37 | # x2 for gradients 38 | self.total_output_bytes += layer_info.output_bytes * 2 39 | if layer_info.is_recursive: 40 | continue 41 | self.total_params += max(layer_info.num_params, 0) 42 | self.total_param_bytes += layer_info.param_bytes 43 | self.trainable_params += max(layer_info.trainable_params, 0) 44 | else: 45 | if layer_info.is_recursive: 46 | continue 47 | leftover_params = layer_info.leftover_params() 48 | leftover_trainable_params = layer_info.leftover_trainable_params() 49 | self.total_params += max(leftover_params, 0) 50 | self.trainable_params += max(leftover_trainable_params, 0) 51 | self.formatting.set_layer_name_width(summary_list) 52 | 53 | def __repr__(self) -> str: 54 | """Print results of the summary.""" 55 | divider = "=" * self.formatting.get_total_width() 56 | total_params = ModelStatistics.format_output_num( 57 | self.total_params, self.formatting.params_count_units, False 58 | ) 59 | trainable_params = ModelStatistics.format_output_num( 60 | self.trainable_params, self.formatting.params_count_units, False 61 | ) 62 | non_trainable_params = ModelStatistics.format_output_num( 63 | self.total_params - self.trainable_params, 64 | self.formatting.params_count_units, 65 | False, 66 | ) 67 | all_layers = self.formatting.layers_to_str(self.summary_list, self.total_params) 68 | summary_str = ( 69 | f"{divider}\n" 70 | f"{self.formatting.header_row()}{divider}\n" 71 | f"{all_layers}{divider}\n" 72 | f"Total params{total_params}\n" 73 | f"Trainable params{trainable_params}\n" 74 | f"Non-trainable params{non_trainable_params}\n" 75 | ) 76 | if self.input_size: 77 | macs = ModelStatistics.format_output_num( 78 | self.total_mult_adds, self.formatting.macs_units, False 79 | ) 80 | input_size = ModelStatistics.format_output_num( 81 | self.total_input, self.formatting.params_size_units, True 82 | ) 83 | output_bytes = ModelStatistics.format_output_num( 84 | self.total_output_bytes, self.formatting.params_size_units, True 85 | ) 86 | param_bytes = ModelStatistics.format_output_num( 87 | self.total_param_bytes, self.formatting.params_size_units, True 88 | ) 89 | total_bytes = ModelStatistics.format_output_num( 90 | self.total_input + self.total_output_bytes + self.total_param_bytes, 91 | self.formatting.params_size_units, 92 | True, 93 | ) 94 | summary_str += ( 95 | f"Total mult-adds{macs}\n{divider}\n" 96 | f"Input size{input_size}\n" 97 | f"Forward/backward pass size{output_bytes}\n" 98 | f"Params size{param_bytes}\n" 99 | f"Estimated Total Size{total_bytes}\n" 100 | ) 101 | summary_str += divider 102 | return summary_str 103 | 104 | @staticmethod 105 | def float_to_megabytes(num: int) -> float: 106 | """Converts a number (assume floats, 4 bytes each) to megabytes.""" 107 | return num * 4 / 1e6 108 | 109 | @staticmethod 110 | def to_megabytes(num: int) -> float: 111 | """Converts bytes to megabytes.""" 112 | return num / 1e6 113 | 114 | @staticmethod 115 | def to_readable(num: float, units: Units = Units.AUTO) -> tuple[Units, float]: 116 | """Converts a number to millions, billions, or trillions.""" 117 | if units == Units.AUTO: 118 | if num >= 1e12: 119 | return Units.TERABYTES, num / 1e12 120 | if num >= 1e9: 121 | return Units.GIGABYTES, num / 1e9 122 | if num >= 1e6: 123 | return Units.MEGABYTES, num / 1e6 124 | if num >= 1e3: 125 | return Units.KILOBYTES, num / 1e3 126 | return Units.NONE, num 127 | return units, num / CONVERSION_FACTORS[units] 128 | 129 | @staticmethod 130 | def format_output_num(num: int, units: Units, is_bytes: bool) -> str: 131 | units_used, converted_num = ModelStatistics.to_readable(num, units) 132 | if isinstance(converted_num, float) and converted_num.is_integer(): 133 | converted_num = int(converted_num) 134 | units_display = ( 135 | "" 136 | if units_used == Units.NONE 137 | else f" ({units_used.value}{'B' if is_bytes else ''})" 138 | ) 139 | fmt = "d" if isinstance(converted_num, int) else ".2f" 140 | return f"{units_display}: {converted_num:,{fmt}}" 141 | -------------------------------------------------------------------------------- /torchinfo/py.typed: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TylerYep/torchinfo/e67e74885ce103d7508b71c20e279aee1d9eb4ce/torchinfo/py.typed --------------------------------------------------------------------------------