Skip to content

GitLab

Unverified Commit b2225436 authored by Bill Wu's avatar Bill Wu Committed by GitHub
Browse files

flops counter formatting fix (#3837)

parent 91ef554d
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 2 deletions
nni/compression/pytorch/utils/counter.py
View file @ b2225436
...@@ -313,12 +313,13 @@ class ModelProfiler: ...@@ -313,12 +313,13 @@ class ModelProfiler:
table.field_names = headers table.field_names = headers
for i, result in enumerate(self.results): for i, result in enumerate(self.results):
flops_count = int(result['flops'].item()) if isinstance(result['flops'], torch.Tensor) else int(result['flops'])
row_values = [ row_values = [
i, i,
result['name'], result['name'],
result['module_type'], result['module_type'],
str(result['weight_shape']), str(result['weight_shape']),
result['flops'], flops_count,
result['params'], result['params'],
] ]
name_counter[result['name']] += 1 name_counter[result['name']] += 1
...@@ -408,4 +409,4 @@ def count_flops_params(model, x, custom_ops=None, verbose=True, mode='default'): ...@@ -408,4 +409,4 @@ def count_flops_params(model, x, custom_ops=None, verbose=True, mode='default'):
print(f'FLOPs total: {profiler.sum_flops()}') print(f'FLOPs total: {profiler.sum_flops()}')
print(f'#Params total: {profiler.sum_params()}') print(f'#Params total: {profiler.sum_params()}')
return profiler.sum_flops(), profiler.sum_params(), profiler.results return profiler.sum_flops(), profiler.sum_params(), profiler.results
\ No newline at end of file
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