Create ColumnSettings and RowSettings enums

Author: TylerYep

tests/conftest.py

@@ -8,7 +8,7 @@
 from _pytest.config.argparsing import Parser
 
 from torchinfo import ModelStatistics
-from torchinfo.formatting import HEADER_TITLES
+from torchinfo.formatting import HEADER_TITLES, ColumnSettings
 from torchinfo.torchinfo import clear_cached_forward_pass
 
 
@@ -28,12 +28,6 @@ def verify_capsys(
         return
 
     test_name = request.node.name.replace("test_", "")
-    if sys.version_info < (3, 7) and test_name == "lstm":
-        warnings.warn(
-            "Verbose output is not determininstic because dictionaries "
-            "are not necessarily ordered in versions before Python 3.7."
-        )
-        return
     if sys.version_info < (3, 8) and test_name == "tmva_net_column_totals":
         warnings.warn(
             "sys.getsizeof can return different results on earlier Python versions."
@@ -67,7 +61,7 @@ def verify_output_str(output: str, filename: str) -> None:
     if output != expected:
         print(f"Expected:\n{expected}\nGot:\n{output}")
     assert output == expected
-    for category in ("num_params", "mult_adds"):
+    for category in (ColumnSettings.NUM_PARAMS, ColumnSettings.MULT_ADDS):
         assert_sum_column_totals_match(output, category)
 
 
@@ -85,7 +79,7 @@ def get_column_value_for_row(line: str, offset: int) -> int:
     return int(col_value.replace(",", "").replace("(", "").replace(")", ""))
 
 
-def assert_sum_column_totals_match(output: str, category: str) -> None:
+def assert_sum_column_totals_match(output: str, category: ColumnSettings) -> None:
     lines = output.replace("=", "").split("\n\n")
     header_row = lines[0].strip()
     offset = header_row.find(HEADER_TITLES[category])
@@ -95,10 +89,10 @@ def assert_sum_column_totals_match(output: str, category: str) -> None:
     calculated_total = sum(get_column_value_for_row(line, offset) for line in layers)
     results = lines[2].split("\n")
 
-    if category == "num_params":
+    if category == ColumnSettings.NUM_PARAMS:
         total_params = results[0].split(":")[1].replace(",", "")
         assert calculated_total == int(total_params)
-    elif category == "mult_adds":
+    elif category == ColumnSettings.MULT_ADDS:
         total_mult_adds = results[-1].split(":")[1].replace(",", "")
         assert float(
             f"{ModelStatistics.to_readable(calculated_total)[1]:0.2f}"

tests/torchinfo_test.py

@@ -28,7 +28,7 @@
     SingleInputNet,
 )
 from tests.fixtures.tmva_net import TMVANet  # type: ignore[attr-defined]
-from torchinfo import ALL_COLUMN_SETTINGS, summary
+from torchinfo import ColumnSettings, summary
 
 
 def test_basic_summary() -> None:
@@ -221,7 +221,7 @@ def test_parameter_list() -> None:
         model,
         input_size=(100, 100),
         verbose=2,
-        col_names=ALL_COLUMN_SETTINGS,
+        col_names=list(ColumnSettings),
         col_width=15,
     )
 

torchinfo/__init__.py

@@ -1,7 +1,7 @@
 """ torchinfo """
-from .formatting import ALL_COLUMN_SETTINGS, ALL_ROW_SETTINGS
+from .enums import ColumnSettings, RowSettings
 from .model_statistics import ModelStatistics
 from .torchinfo import summary
 
-__all__ = ("ModelStatistics", "summary", "ALL_COLUMN_SETTINGS", "ALL_ROW_SETTINGS")
+__all__ = ("ModelStatistics", "summary", "ColumnSettings", "RowSettings")
 __version__ = "1.6.0"

torchinfo/enums.py

@@ -0,0 +1,30 @@
+""" constants.py """
+from __future__ import annotations
+
+from enum import Enum, unique
+
+
+@unique
+class RowSettings(Enum):
+    """Enum containing all available row settings."""
+
+    DEPTH = "depth"
+    VAR_NAMES = "var_names"
+
+
+@unique
+class ColumnSettings(Enum):
+    """Enum containing all available column settings."""
+
+    KERNEL_SIZE = "kernel_size"
+    INPUT_SIZE = "input_size"
+    OUTPUT_SIZE = "output_size"
+    NUM_PARAMS = "num_params"
+    MULT_ADDS = "mult_adds"
+
+
+@unique
+class Verbosity(Enum):
+    """Contains verbosity levels."""
+
+    QUIET, DEFAULT, VERBOSE = 0, 1, 2

torchinfo/formatting.py

@@ -2,45 +2,30 @@
 from __future__ import annotations
 
 import math
-from enum import Enum, unique
-from typing import Any, Iterable
+from typing import Any
 
+from .enums import ColumnSettings, RowSettings, Verbosity
 from .layer_info import LayerInfo
 
-ALL_ROW_SETTINGS = ("depth", "var_names")
-ALL_COLUMN_SETTINGS = (
-    "kernel_size",
-    "input_size",
-    "output_size",
-    "num_params",
-    "mult_adds",
-)
 HEADER_TITLES = {
-    "kernel_size": "Kernel Shape",
-    "input_size": "Input Shape",
-    "output_size": "Output Shape",
-    "num_params": "Param #",
-    "mult_adds": "Mult-Adds",
+    ColumnSettings.KERNEL_SIZE: "Kernel Shape",
+    ColumnSettings.INPUT_SIZE: "Input Shape",
+    ColumnSettings.OUTPUT_SIZE: "Output Shape",
+    ColumnSettings.NUM_PARAMS: "Param #",
+    ColumnSettings.MULT_ADDS: "Mult-Adds",
 }
 
 
-@unique
-class Verbosity(Enum):
-    """Contains verbosity levels."""
-
-    QUIET, DEFAULT, VERBOSE = 0, 1, 2
-
-
 class FormattingOptions:
     """Class that holds information about formatting the table output."""
 
     def __init__(
         self,
         max_depth: int,
         verbose: int,
-        col_names: Iterable[str],
+        col_names: tuple[ColumnSettings, ...],
         col_width: int,
-        row_settings: Iterable[str],
+        row_settings: set[RowSettings],
     ) -> None:
         self.max_depth = max_depth
         self.verbose = verbose
@@ -49,8 +34,8 @@ def __init__(
         self.row_settings = row_settings
 
         self.layer_name_width = 40
-        self.show_var_name = "var_names" in self.row_settings
-        self.show_depth = "depth" in self.row_settings
+        self.show_var_name = RowSettings.VAR_NAMES in self.row_settings
+        self.show_depth = RowSettings.DEPTH in self.row_settings
 
     @staticmethod
     def get_start_str(depth: int) -> str:
@@ -95,7 +80,7 @@ def get_total_width(self) -> int:
         """Calculate the total width of all lines in the table."""
         return len(tuple(self.col_names)) * self.col_width + self.layer_name_width
 
-    def format_row(self, layer_name: str, row_values: dict[str, str]) -> str:
+    def format_row(self, layer_name: str, row_values: dict[ColumnSettings, str]) -> str:
         """Get the string representation of a single layer of the model."""
         info_to_use = [row_values.get(row_type, "") for row_type in self.col_names]
         new_line = f"{layer_name:<{self.layer_name_width}} "
@@ -118,16 +103,18 @@ def layer_info_to_row(
         children_layers: list[LayerInfo],
     ) -> str:
         """Convert layer_info to string representation of a row."""
-        row_values = {
-            "kernel_size": self.str_(layer_info.kernel_size),
-            "input_size": self.str_(layer_info.input_size),
-            "output_size": self.str_(layer_info.output_size),
-            "num_params": layer_info.num_params_to_str(reached_max_depth),
-            "mult_adds": layer_info.macs_to_str(reached_max_depth, children_layers),
+        values_for_row = {
+            ColumnSettings.KERNEL_SIZE: self.str_(layer_info.kernel_size),
+            ColumnSettings.INPUT_SIZE: self.str_(layer_info.input_size),
+            ColumnSettings.OUTPUT_SIZE: self.str_(layer_info.output_size),
+            ColumnSettings.NUM_PARAMS: layer_info.num_params_to_str(reached_max_depth),
+            ColumnSettings.MULT_ADDS: layer_info.macs_to_str(
+                reached_max_depth, children_layers
+            ),
         }
         start_str = self.get_start_str(layer_info.depth)
         layer_name = layer_info.get_layer_name(self.show_var_name, self.show_depth)
-        new_line = self.format_row(f"{start_str}{layer_name}", row_values)
+        new_line = self.format_row(f"{start_str}{layer_name}", values_for_row)
 
         if self.verbose == Verbosity.VERBOSE.value:
             for inner_name, inner_layer_info in layer_info.inner_layers.items():

torchinfo/layer_info.py

@@ -7,6 +7,8 @@
 from torch import nn
 from torch.jit import ScriptModule
 
+from .enums import ColumnSettings
+
 DETECTED_INPUT_OUTPUT_TYPES = Union[
     Sequence[Any], Dict[Any, torch.Tensor], torch.Tensor
 ]
@@ -39,8 +41,8 @@ def __init__(
             if isinstance(module, ScriptModule)
             else module.__class__.__name__
         )
-        # {layer name: {row_name: row_value}}
-        self.inner_layers: dict[str, dict[str, Any]] = {}
+        # {layer name: {col_name: value_for_row}}
+        self.inner_layers: dict[str, dict[ColumnSettings, Any]] = {}
         self.depth = depth
         self.depth_index = depth_index
         self.executed = False
@@ -163,13 +165,13 @@ def calculate_num_params(self) -> None:
 
             # RNN modules have inner weights such as weight_ih_l0
             self.inner_layers[name] = {
-                "kernel_size": str(ksize),
-                "num_params": f"├─{cur_params:,}",
+                ColumnSettings.KERNEL_SIZE: str(ksize),
+                ColumnSettings.NUM_PARAMS: f"├─{cur_params:,}",
             }
         if self.inner_layers:
             self.inner_layers[name][
-                "num_params"
-            ] = f"└─{self.inner_layers[name]['num_params'][2:]}"
+                ColumnSettings.NUM_PARAMS
+            ] = f"└─{self.inner_layers[name][ColumnSettings.NUM_PARAMS][2:]}"
 
     def calculate_macs(self) -> None:
         """

torchinfo/torchinfo.py

@@ -20,12 +20,8 @@
 from torch.jit import ScriptModule
 from torch.utils.hooks import RemovableHandle
 
-from .formatting import (
-    ALL_COLUMN_SETTINGS,
-    ALL_ROW_SETTINGS,
-    FormattingOptions,
-    Verbosity,
-)
+from .enums import ColumnSettings, RowSettings, Verbosity
+from .formatting import FormattingOptions
 from .layer_info import LayerInfo
 from .model_statistics import ModelStatistics
 
@@ -40,8 +36,13 @@
 INPUT_SIZE_TYPE = Sequence[Union[int, Sequence[Any], torch.Size]]
 CORRECTED_INPUT_SIZE_TYPE = List[Union[Sequence[Any], torch.Size]]
 
-DEFAULT_COLUMN_NAMES = ("output_size", "num_params")
-DEFAULT_ROW_SETTINGS = ("depth",)
+DEFAULT_COLUMN_NAMES = (ColumnSettings.OUTPUT_SIZE, ColumnSettings.NUM_PARAMS)
+DEFAULT_ROW_SETTINGS = {RowSettings.DEPTH}
+REQUIRES_INPUT = {
+    ColumnSettings.INPUT_SIZE,
+    ColumnSettings.OUTPUT_SIZE,
+    ColumnSettings.MULT_ADDS,
+}
 
 _cached_forward_pass: dict[str, list[LayerInfo]] = {}
 
@@ -166,12 +167,19 @@ class name as the key. If the forward pass is an expensive operation,
                 See torchinfo/model_statistics.py for more information.
     """
     input_data_specified = input_data is not None or input_size is not None
-
     if col_names is None:
-        col_names = DEFAULT_COLUMN_NAMES if input_data_specified else ("num_params",)
+        columns = (
+            DEFAULT_COLUMN_NAMES
+            if input_data_specified
+            else (ColumnSettings.NUM_PARAMS,)
+        )
+    else:
+        columns = tuple(ColumnSettings(name) for name in col_names)
 
     if row_settings is None:
-        row_settings = DEFAULT_ROW_SETTINGS
+        rows = DEFAULT_ROW_SETTINGS
+    else:
+        rows = {RowSettings(name) for name in row_settings}
 
     if verbose is None:
         # pylint: disable=no-member
@@ -184,17 +192,15 @@ class name as the key. If the forward pass is an expensive operation,
     if device is None:
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-    validate_user_params(
-        input_data, input_size, col_names, col_width, row_settings, verbose
-    )
+    validate_user_params(input_data, input_size, columns, col_width, verbose)
 
     x, correct_input_size = process_input(
         input_data, input_size, batch_dim, device, dtypes
     )
     summary_list = forward_pass(
         model, x, batch_dim, cache_forward_pass, device, **kwargs
     )
-    formatting = FormattingOptions(depth, verbose, col_names, col_width, row_settings)
+    formatting = FormattingOptions(depth, verbose, columns, col_width, rows)
     results = ModelStatistics(
         summary_list, correct_input_size, get_total_memory_used(x), formatting
     )
@@ -287,9 +293,8 @@ def forward_pass(
 def validate_user_params(
     input_data: INPUT_DATA_TYPE | None,
     input_size: INPUT_SIZE_TYPE | None,
-    col_names: Iterable[str],
+    col_names: tuple[ColumnSettings, ...],
     col_width: int,
-    row_settings: Iterable[str],
     verbose: int,
 ) -> None:
     """Raise exceptions if the user's input is invalid."""
@@ -304,16 +309,12 @@ def validate_user_params(
         raise RuntimeError("Only one of (input_data, input_size) should be specified.")
 
     neither_input_specified = input_data is None and input_size is None
-    for col in col_names:
-        if col not in ALL_COLUMN_SETTINGS:
-            raise ValueError(f"Column {col} is not a valid column name.")
-        if neither_input_specified and col not in ("num_params", "kernel_size"):
-            raise ValueError(
-                f"You must pass input_data or input_size in order to use column {col}"
-            )
-    for setting in row_settings:
-        if setting not in ALL_ROW_SETTINGS:
-            raise ValueError(f"Row setting {setting} is not a valid setting.")
+    not_allowed = set(col_names) & REQUIRES_INPUT
+    if neither_input_specified and not_allowed:
+        raise ValueError(
+            "You must pass input_data or input_size in order "
+            f"to use columns: {not_allowed}"
+        )
 
 
 def traverse_input_data(