From 2cfe7a9ecd7e3bc0bccec80b35d986a9608fd0a5 Mon Sep 17 00:00:00 2001 From: KaiyangZhou Date: Wed, 22 May 2019 22:09:09 +0100 Subject: [PATCH] update printing of model complexity --- scripts/main.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/scripts/main.py b/scripts/main.py index 6c8f624..682ac04 100755 --- 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)