update printing of model complexity

2019-05-22 22:09:09 +01:00 · 2019-05-22 22:09:09 +01:00 · 2cfe7a9ecd
parent fa2a8b231b
commit 2cfe7a9ecd
1 changed files with 2 additions and 1 deletions
--- a/scripts/main.py
+++ b/scripts/main.py
@ -105,7 +105,8 @@ def main():
        pretrained=(not args.no_pretrained),
        use_gpu=use_gpu
    )
-    compute_model_complexity(model, (1, 3, args.height, args.width), verbose=True)
+    num_params, flops = compute_model_complexity(model, (1, 3, args.height, args.width))
+    print('Model complexity: params={:,} flops={:,}'.format(num_params, flops))

    if args.load_weights and check_isfile(args.load_weights):
        load_pretrained_weights(model, args.load_weights)