mirrors
/
deep-person-reid
mirror of https://github.com/KaiyangZhou/deep-person-reid.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

add input_size to print

pull/201/head
KaiyangZhou 2019-05-22 21:46:48 +01:00
parent 575b456a92
commit 8ec474daa7
1 changed files with 6 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
torchreid/utils/model_complexity.py
View File

@ -331,15 +331,16 @@ def compute_model_complexity(model, input_size, verbose=False, only_conv_linear=
per_module_flops[layer.class_name].append(layer.flops)
if verbose:
print(' {}'.format('-'*30))
print(' Model Complexity')
print(' {}'.format('-'*30))
num_udscore = 55
print(' {}'.format('-'*num_udscore))
print(' Model complexity with input size {}'.format(input_size))
print(' {}'.format('-'*num_udscore))
for class_name in per_module_params:
params = int(np.sum(per_module_params[class_name]))
flops = int(np.sum(per_module_flops[class_name]))
print(' {} (params={:,}, flops={:,})'.format(class_name, params, flops))
print(' {}'.format('-'*30))
print(' {}'.format('-'*num_udscore))
print(' Total (params={:,}, flops={:,})'.format(total_params, total_flops))
print(' {}'.format('-'*30))
print(' {}'.format('-'*num_udscore))
return total_params, total_flops