From 8918e6347683e0f2a8a3d7ef93331001985f6560 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sat, 5 Dec 2020 11:41:34 +0100 Subject: [PATCH] Increase FLOPS robustness (#1608) --- utils/torch_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/utils/torch_utils.py b/utils/torch_utils.py index b330ca5..ce92749 100644 --- a/utils/torch_utils.py +++ b/utils/torch_utils.py @@ -1,12 +1,12 @@ # PyTorch utils import logging -import math import os import time from contextlib import contextmanager from copy import deepcopy +import math import torch import torch.backends.cudnn as cudnn import torch.nn as nn @@ -152,7 +152,7 @@ def model_info(model, verbose=False, img_size=640): try: # FLOPS from thop import profile - stride = int(model.stride.max()) + stride = int(model.stride.max()) if hasattr(model, 'stride') else 32 img = torch.zeros((1, 3, stride, stride), device=next(model.parameters()).device) # input flops = profile(deepcopy(model), inputs=(img,), verbose=False)[0] / 1E9 * 2 # stride FLOPS img_size = img_size if isinstance(img_size, list) else [img_size, img_size] # expand if int/float