Skip to content

Commit

Permalink
Export Verbosity as an IntEnum
Browse files Browse the repository at this point in the history
  • Loading branch information
TylerYep committed Dec 24, 2021
1 parent 2b82275 commit d6a9af8
Show file tree
Hide file tree
Showing 6 changed files with 10 additions and 9 deletions.
2 changes: 1 addition & 1 deletion LICENSE
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
MIT License

Copyright (c) 2021 Tyler Yep
Copyright (c) 2022 Tyler Yep

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
Expand Down
3 changes: 2 additions & 1 deletion tests/torchinfo_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
)
from tests.fixtures.tmva_net import TMVANet # type: ignore[attr-defined]
from torchinfo import ColumnSettings, summary
from torchinfo.enums import Verbosity


def test_basic_summary() -> None:
Expand Down Expand Up @@ -253,7 +254,7 @@ def test_lstm() -> None:
LSTMNet(),
input_size=(1, 100),
dtypes=[torch.long],
verbose=2,
verbose=Verbosity.VERBOSE,
col_width=20,
col_names=("kernel_size", "output_size", "num_params", "mult_adds"),
row_settings=("var_names",),
Expand Down
6 changes: 3 additions & 3 deletions torchinfo/__init__.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
""" torchinfo """
from .enums import ColumnSettings, RowSettings
from .enums import ColumnSettings, RowSettings, Verbosity
from .model_statistics import ModelStatistics
from .torchinfo import summary

__all__ = ("ModelStatistics", "summary", "ColumnSettings", "RowSettings")
__version__ = "1.6.0"
__all__ = ("summary", "ColumnSettings", "ModelStatistics", "RowSettings", "Verbosity")
__version__ = "1.6.1"
4 changes: 2 additions & 2 deletions torchinfo/enums.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
""" constants.py """
from __future__ import annotations

from enum import Enum, unique
from enum import Enum, IntEnum, unique


@unique
Expand All @@ -25,7 +25,7 @@ class ColumnSettings(Enum):


@unique
class Verbosity(Enum):
class Verbosity(IntEnum):
"""Contains verbosity levels."""

QUIET, DEFAULT, VERBOSE = 0, 1, 2
2 changes: 1 addition & 1 deletion torchinfo/formatting.py
Original file line number Diff line number Diff line change
Expand Up @@ -119,7 +119,7 @@ def layer_info_to_row(
layer_name = layer_info.get_layer_name(self.show_var_name, self.show_depth)
new_line = self.format_row(f"{start_str}{layer_name}", values_for_row)

if self.verbose == Verbosity.VERBOSE.value:
if self.verbose == Verbosity.VERBOSE:
for inner_name, inner_layer_info in layer_info.inner_layers.items():
prefix = self.get_start_str(layer_info.depth + 1)
new_line += self.format_row(f"{prefix}{inner_name}", inner_layer_info)
Expand Down
2 changes: 1 addition & 1 deletion torchinfo/torchinfo.py
Original file line number Diff line number Diff line change
Expand Up @@ -205,7 +205,7 @@ class name as the key. If the forward pass is an expensive operation,
results = ModelStatistics(
summary_list, correct_input_size, get_total_memory_used(x), formatting
)
if verbose > Verbosity.QUIET.value:
if verbose > Verbosity.QUIET:
print(results)
return results

Expand Down

0 comments on commit d6a9af8

Please sign in to comment.