-
-
Notifications
You must be signed in to change notification settings - Fork 16.5k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
* Add benchmarks.py * Update * Add requirements * Updates * Updates * Updates * Updates * Updates * Updates * dataset autodownload from root * Update * Redirect to /dev/null * sudo --help * Cleanup * Add exports pd df * Updates * Updates * Updates * Cleanup * dir handling fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Cleanup * Cleanup2 * Cleanup3 * Cleanup model_type Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
- Loading branch information
1 parent
96d8f86
commit a45e472
Showing
4 changed files
with
127 additions
and
6 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,92 @@ | ||
# YOLOv5 🚀 by Ultralytics, GPL-3.0 license | ||
""" | ||
Run YOLOv5 benchmarks on all supported export formats | ||
Format | `export.py --include` | Model | ||
--- | --- | --- | ||
PyTorch | - | yolov5s.pt | ||
TorchScript | `torchscript` | yolov5s.torchscript | ||
ONNX | `onnx` | yolov5s.onnx | ||
OpenVINO | `openvino` | yolov5s_openvino_model/ | ||
TensorRT | `engine` | yolov5s.engine | ||
CoreML | `coreml` | yolov5s.mlmodel | ||
TensorFlow SavedModel | `saved_model` | yolov5s_saved_model/ | ||
TensorFlow GraphDef | `pb` | yolov5s.pb | ||
TensorFlow Lite | `tflite` | yolov5s.tflite | ||
TensorFlow Edge TPU | `edgetpu` | yolov5s_edgetpu.tflite | ||
TensorFlow.js | `tfjs` | yolov5s_web_model/ | ||
Requirements: | ||
$ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime openvino-dev tensorflow-cpu # CPU | ||
$ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime-gpu openvino-dev tensorflow # GPU | ||
Usage: | ||
$ python utils/benchmarks.py --weights yolov5s.pt --img 640 | ||
""" | ||
|
||
import argparse | ||
import sys | ||
import time | ||
from pathlib import Path | ||
|
||
import pandas as pd | ||
|
||
FILE = Path(__file__).resolve() | ||
ROOT = FILE.parents[1] # YOLOv5 root directory | ||
if str(ROOT) not in sys.path: | ||
sys.path.append(str(ROOT)) # add ROOT to PATH | ||
# ROOT = ROOT.relative_to(Path.cwd()) # relative | ||
|
||
import export | ||
import val | ||
from utils import notebook_init | ||
from utils.general import LOGGER, print_args | ||
|
||
|
||
def run(weights=ROOT / 'yolov5s.pt', # weights path | ||
imgsz=640, # inference size (pixels) | ||
batch_size=1, # batch size | ||
data=ROOT / 'data/coco128.yaml', # dataset.yaml path | ||
): | ||
y, t = [], time.time() | ||
formats = export.export_formats() | ||
for i, (name, f, suffix) in formats.iterrows(): # index, (name, file, suffix) | ||
try: | ||
w = weights if f == '-' else export.run(weights=weights, imgsz=[imgsz], include=[f], device='cpu')[-1] | ||
assert suffix in str(w), 'export failed' | ||
result = val.run(data, w, batch_size, imgsz=imgsz, plots=False, device='cpu', task='benchmark') | ||
metrics = result[0] # metrics (mp, mr, map50, map, *losses(box, obj, cls)) | ||
speeds = result[2] # times (preprocess, inference, postprocess) | ||
y.append([name, metrics[3], speeds[1]]) # mAP, t_inference | ||
except Exception as e: | ||
LOGGER.warning(f'WARNING: Benchmark failure for {name}: {e}') | ||
y.append([name, None, None]) # mAP, t_inference | ||
|
||
# Print results | ||
LOGGER.info('\n') | ||
parse_opt() | ||
notebook_init() # print system info | ||
py = pd.DataFrame(y, columns=['Format', '[email protected]:0.95', 'Inference time (ms)']) | ||
LOGGER.info(f'\nBenchmarks complete ({time.time() - t:.2f}s)') | ||
LOGGER.info(str(py)) | ||
return py | ||
|
||
|
||
def parse_opt(): | ||
parser = argparse.ArgumentParser() | ||
parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='weights path') | ||
parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='inference size (pixels)') | ||
parser.add_argument('--batch-size', type=int, default=1, help='batch size') | ||
parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path') | ||
opt = parser.parse_args() | ||
print_args(FILE.stem, opt) | ||
return opt | ||
|
||
|
||
def main(opt): | ||
run(**vars(opt)) | ||
|
||
|
||
if __name__ == "__main__": | ||
opt = parse_opt() | ||
main(opt) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters