Skip to content

GitLab

"test/vscode:/vscode.git/clone" did not exist on "5eb43ea7815d8a0cf93d0a4bb6f29faf57ae6553"
Unverified Commit 39d9a8d9 authored by BasicCoder's avatar BasicCoder Committed by GitHub
Browse files

add input and output shape information (#4539)

parent e483aa01
Show whitespace changes
Inline Side-by-side
Showing with 5 additions and 0 deletions
nni/compression/pytorch/utils/counter.py
View file @ 39d9a8d9
...@@ -293,6 +293,8 @@ class ModelProfiler: ...@@ -293,6 +293,8 @@ class ModelProfiler:
'Name', 'Name',
'Type', 'Type',
'Weight Shape', 'Weight Shape',
'Input Size',
'Output Size',
'FLOPs', 'FLOPs',
'#Params', '#Params',
] ]
...@@ -307,6 +309,8 @@ class ModelProfiler: ...@@ -307,6 +309,8 @@ class ModelProfiler:
result['name'], result['name'],
result['module_type'], result['module_type'],
str(result['weight_shape']), str(result['weight_shape']),
result['input_size'],
result['output_size'],
flops_count, flops_count,
result['params'], result['params'],
] ]
...@@ -314,6 +318,7 @@ class ModelProfiler: ...@@ -314,6 +318,7 @@ class ModelProfiler:
if has_multi_use: if has_multi_use:
row_values.append(name_counter[result['name']]) row_values.append(name_counter[result['name']])
table.add_row(row_values) table.add_row(row_values)
table.align["Name"] = "l"
return table return table
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment