From 0bd9c48609522b455cf425578828c55b924de3d5 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 21 Dec 2020 13:29:52 -0800 Subject: [PATCH] Update torch_utils.py FLOPS to GFLOPS --- utils/torch_utils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/utils/torch_utils.py b/utils/torch_utils.py index 4638377..c76a318 100644 --- a/utils/torch_utils.py +++ b/utils/torch_utils.py @@ -86,7 +86,7 @@ def profile(x, ops, n=100, device=None): x = x.to(device) x.requires_grad = True print(torch.__version__, device.type, torch.cuda.get_device_properties(0) if device.type == 'cuda' else '') - print(f"\n{'Params':>12s}{'FLOPS':>12s}{'forward (ms)':>16s}{'backward (ms)':>16s}{'input':>24s}{'output':>24s}") + print(f"\n{'Params':>12s}{'GFLOPS':>12s}{'forward (ms)':>16s}{'backward (ms)':>16s}{'input':>24s}{'output':>24s}") for m in ops if isinstance(ops, list) else [ops]: m = m.to(device) if hasattr(m, 'to') else m dtf, dtb, t = 0., 0., [0., 0., 0.] # dt forward, backward @@ -197,9 +197,9 @@ def model_info(model, verbose=False, img_size=640): from thop import profile stride = int(model.stride.max()) if hasattr(model, 'stride') else 32 img = torch.zeros((1, model.yaml.get('ch', 3), stride, stride), device=next(model.parameters()).device) # input - flops = profile(deepcopy(model), inputs=(img,), verbose=False)[0] / 1E9 * 2 # stride FLOPS + flops = profile(deepcopy(model), inputs=(img,), verbose=False)[0] / 1E9 * 2 # stride GFLOPS img_size = img_size if isinstance(img_size, list) else [img_size, img_size] # expand if int/float - fs = ', %.1f GFLOPS' % (flops * img_size[0] / stride * img_size[1] / stride) # 640x640 FLOPS + fs = ', %.1f GFLOPS' % (flops * img_size[0] / stride * img_size[1] / stride) # 640x640 GFLOPS except (ImportError, Exception): fs = ''