Skip to content

Commit

Permalink
Add File Size (MB) column to benchmarks (ultralytics#8359)
Browse files Browse the repository at this point in the history
* Add filesize to benchmarks.py

* Add filesize to benchmarks.py

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
  • Loading branch information
2 people authored and Clay Januhowski committed Sep 8, 2022
1 parent 3ec4114 commit b1e607f
Showing 1 changed file with 5 additions and 4 deletions.
9 changes: 5 additions & 4 deletions utils/benchmarks.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@
import export
import val
from utils import notebook_init
from utils.general import LOGGER, check_yaml, print_args
from utils.general import LOGGER, check_yaml, file_size, print_args
from utils.torch_utils import select_device


Expand Down Expand Up @@ -75,18 +75,19 @@ def run(
result = val.run(data, w, batch_size, imgsz, plots=False, device=device, task='benchmark', half=half)
metrics = result[0] # metrics (mp, mr, map50, map, *losses(box, obj, cls))
speeds = result[2] # times (preprocess, inference, postprocess)
y.append([name, round(metrics[3], 4), round(speeds[1], 2)]) # mAP, t_inference
y.append([name, round(file_size(w), 1), round(metrics[3], 4), round(speeds[1], 2)]) # MB, mAP, t_inference
except Exception as e:
LOGGER.warning(f'WARNING: Benchmark failure for {name}: {e}')
y.append([name, None, None]) # mAP, t_inference
y.append([name, None, None, None]) # mAP, t_inference
if pt_only and i == 0:
break # break after PyTorch

# Print results
LOGGER.info('\n')
parse_opt()
notebook_init() # print system info
py = pd.DataFrame(y, columns=['Format', 'mAP@0.5:0.95', 'Inference time (ms)'] if map else ['Format', 'Export', ''])
c = ['Format', 'Size (MB)', 'mAP@0.5:0.95', 'Inference time (ms)'] if map else ['Format', 'Export', '', '']
py = pd.DataFrame(y, columns=c)
LOGGER.info(f'\nBenchmarks complete ({time.time() - t:.2f}s)')
LOGGER.info(str(py if map else py.iloc[:, :2]))
return py
Expand Down

0 comments on commit b1e607f

Please sign in to comment.