Skip to content

Commit e67e748

Browse files
TylerYepTylerYep
authored
Merge pull request #363 from Tanguy-ddv/main
Automatic conversion of displayed sizes, add support for kB and B.
2 parents 164597f + be7ea9f commit e67e748

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

67 files changed

+409
-352
lines changed

‎tests/conftest.py‎

Lines changed: 34 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import sys
22
import warnings
33
from pathlib import Path
4-
from typing import Iterator
4+
from typing import Iterator, Tuple
55

66
import pytest
77

@@ -60,10 +60,30 @@ def verify_output(capsys: pytest.CaptureFixture[str], filename: str) -> None:
6060

6161
def verify_output_str(output: str, filename: str) -> None:
6262
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)
6376
assert output == expected
6477
for category in (ColumnSettings.NUM_PARAMS, ColumnSettings.MULT_ADDS):
6578
assert_sum_column_totals_match(output, category)
6679

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
6787

6888
def get_column_value_for_row(line: str, offset: int) -> int:
6989
"""Helper function for getting the column totals."""
@@ -88,12 +108,23 @@ def assert_sum_column_totals_match(output: str, category: ColumnSettings) -> Non
88108
if offset == -1:
89109
return
90110
layers = lines[1].split("\n")
91-
calculated_total = sum(get_column_value_for_row(line, offset) for line in layers)
111+
calculated_total = float(sum(get_column_value_for_row(line, offset) for line in layers))
92112
results = lines[2].split("\n")
93113

94114
if category == ColumnSettings.NUM_PARAMS:
95115
total_params = results[0].split(":")[1].replace(",", "")
96-
assert calculated_total == int(total_params)
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)
97128
elif category == ColumnSettings.MULT_ADDS:
98129
total_mult_adds = results[-1].split(":")[1].replace(",", "")
99130
assert float(

‎tests/test_output/ascii_only.out‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,7 @@ Trainable params: 11,689,512
7777
Non-trainable params: 0
7878
Total mult-adds (M): 148.57
7979
==========================================================================================
80-
Input size (MB): 0.05
80+
Input size (kB): 49.22
8181
Forward/backward pass size (MB): 3.25
8282
Params size (MB): 46.76
8383
Estimated Total Size (MB): 50.06

‎tests/test_output/autoencoder.out‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,8 @@ Trainable params: 883
1616
Non-trainable params: 0
1717
Total mult-adds (M): 3.62
1818
===================================================================================================================
19-
Input size (MB): 0.05
20-
Forward/backward pass size (MB): 0.62
21-
Params size (MB): 0.00
22-
Estimated Total Size (MB): 0.68
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
2323
===================================================================================================================

‎tests/test_output/batch_size_optimization.out‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -11,10 +11,10 @@ SingleInputNet [1, 10] --
1111
Total params: 21,840
1212
Trainable params: 21,840
1313
Non-trainable params: 0
14-
Total mult-adds (M): 0.49
14+
Total mult-adds (k): 487.60
1515
==========================================================================================
16-
Input size (MB): 0.00
17-
Forward/backward pass size (MB): 0.06
18-
Params size (MB): 0.09
19-
Estimated Total Size (MB): 0.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
2020
==========================================================================================

‎tests/test_output/cnn_module_list.out‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -27,12 +27,12 @@ CNNModuleList [1, 1, 10] --
2727
Total params: 10
2828
Trainable params: 10
2929
Non-trainable params: 0
30-
Total mult-adds (M): 0.00
30+
Total mult-adds: 100
3131
==========================================================================================
32-
Input size (MB): 0.00
33-
Forward/backward pass size (MB): 0.00
34-
Params size (MB): 0.00
35-
Estimated Total Size (MB): 0.00
32+
Input size: 112
33+
Forward/backward pass size: 400
34+
Params size: 40
35+
Estimated Total Size: 552
3636
==========================================================================================
3737
==========================================================================================
3838
Layer (type:depth-idx) Output Shape Param #
@@ -63,10 +63,10 @@ CNNModuleList [1, 1, 10] --
6363
Total params: 10
6464
Trainable params: 10
6565
Non-trainable params: 0
66-
Total mult-adds (M): 0.00
66+
Total mult-adds: 100
6767
==========================================================================================
68-
Input size (MB): 0.00
69-
Forward/backward pass size (MB): 0.00
70-
Params size (MB): 0.00
71-
Estimated Total Size (MB): 0.00
68+
Input size: 112
69+
Forward/backward pass size: 400
70+
Params size: 40
71+
Estimated Total Size: 552
7272
==========================================================================================

‎tests/test_output/container.out‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -22,12 +22,12 @@ ContainerModule [1, 5] --
2222
Total params: 150
2323
Trainable params: 150
2424
Non-trainable params: 0
25-
Total mult-adds (M): 0.00
25+
Total mult-adds: 330
2626
==========================================================================================
27-
Input size (MB): 0.00
28-
Forward/backward pass size (MB): 0.00
29-
Params size (MB): 0.00
30-
Estimated Total Size (MB): 0.00
27+
Input size: 92
28+
Forward/backward pass size: 440
29+
Params size: 600
30+
Estimated Total Size (kB): 1.13
3131
==========================================================================================
3232
==========================================================================================
3333
Layer (type:depth-idx) Output Shape Param #
@@ -45,10 +45,10 @@ ContainerModule [5] --
4545
Total params: 150
4646
Trainable params: 150
4747
Non-trainable params: 0
48-
Total mult-adds (M): 0.00
48+
Total mult-adds: 330
4949
==========================================================================================
50-
Input size (MB): 0.00
51-
Forward/backward pass size (MB): 0.00
52-
Params size (MB): 0.00
53-
Estimated Total Size (MB): 0.00
50+
Input size: 92
51+
Forward/backward pass size: 440
52+
Params size: 600
53+
Estimated Total Size (kB): 1.13
5454
==========================================================================================

‎tests/test_output/custom_parameter.out‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,10 @@ CustomParameter [4, 8] 32
66
Total params: 32
77
Trainable params: 32
88
Non-trainable params: 0
9-
Total mult-adds (M): 0.00
9+
Total mult-adds: 128
1010
==========================================================================================
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
11+
Input size: 76
12+
Forward/backward pass size: 256
13+
Params size: 128
14+
Estimated Total Size: 460
1515
==========================================================================================

‎tests/test_output/device.out‎

Lines changed: 20 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -13,10 +13,10 @@ Trainable params: 21,840
1313
Non-trainable params: 0
1414
Total mult-adds (M): 2.44
1515
==========================================================================================
16-
Input size (MB): 0.02
17-
Forward/backward pass size (MB): 0.28
18-
Params size (MB): 0.09
19-
Estimated Total Size (MB): 0.39
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
2020
==========================================================================================
2121
==========================================================================================
2222
Layer (type:depth-idx) Output Shape Param #
@@ -33,10 +33,10 @@ Trainable params: 21,840
3333
Non-trainable params: 0
3434
Total mult-adds (M): 2.44
3535
==========================================================================================
36-
Input size (MB): 0.02
37-
Forward/backward pass size (MB): 0.28
38-
Params size (MB): 0.09
39-
Estimated Total Size (MB): 0.39
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
4040
==========================================================================================
4141
==========================================================================================
4242
Layer (type:depth-idx) Output Shape Param #
@@ -53,10 +53,10 @@ Trainable params: 21,840
5353
Non-trainable params: 0
5454
Total mult-adds (M): 2.44
5555
==========================================================================================
56-
Input size (MB): 0.02
57-
Forward/backward pass size (MB): 0.28
58-
Params size (MB): 0.09
59-
Estimated Total Size (MB): 0.39
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
6060
==========================================================================================
6161
==========================================================================================
6262
Layer (type:depth-idx) Output Shape Param #
@@ -73,10 +73,10 @@ Trainable params: 21,840
7373
Non-trainable params: 0
7474
Total mult-adds (M): 2.44
7575
==========================================================================================
76-
Input size (MB): 0.02
77-
Forward/backward pass size (MB): 0.28
78-
Params size (MB): 0.09
79-
Estimated Total Size (MB): 0.39
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
8080
==========================================================================================
8181
==========================================================================================
8282
Layer (type:depth-idx) Output Shape Param #
@@ -93,8 +93,8 @@ Trainable params: 21,840
9393
Non-trainable params: 0
9494
Total mult-adds (M): 2.44
9595
==========================================================================================
96-
Input size (MB): 0.02
97-
Forward/backward pass size (MB): 0.28
98-
Params size (MB): 0.09
99-
Estimated Total Size (MB): 0.39
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
100100
==========================================================================================

‎tests/test_output/dict_input.out‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -10,10 +10,10 @@ MultipleInputNetDifferentDtypes [2, 10] --
1010
Total params: 31,120
1111
Trainable params: 31,120
1212
Non-trainable params: 0
13-
Total mult-adds (M): 0.03
13+
Total mult-adds (k): 31.12
1414
==========================================================================================
15-
Input size (MB): 0.00
16-
Forward/backward pass size (MB): 0.00
17-
Params size (MB): 0.12
18-
Estimated Total Size (MB): 0.13
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
1919
==========================================================================================

‎tests/test_output/dict_parameters_1.out‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,10 @@ DictParameter [10, 1] --
66
Total params: 0
77
Trainable params: 0
88
Non-trainable params: 0
9-
Total mult-adds (M): 0
9+
Total mult-adds: 0
1010
==========================================================================================
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
11+
Input size: 229
12+
Forward/backward pass size: 0
13+
Params size: 0
14+
Estimated Total Size: 229
1515
==========================================================================================

0 commit comments

Comments
 (0)