Register
Login
Resources
Docs Blog Datasets Glossary Case Studies Tutorials & Webinars
Product
Data Engine LLMs Platform Enterprise
Pricing Explore
Connect to our Discord channel

#20413 YOLOE: Fix visual prompt training

Merged
Ghost merged 1 commits into Ultralytics:main from ultralytics:yoloe-vp-fix
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
  1. # Ultralytics πŸš€ AGPL-3.0 License - https://ultralytics.com/license
  2. """
  3. Benchmark a YOLO model formats for speed and accuracy.
  4. Usage:
  5. from ultralytics.utils.benchmarks import ProfileModels, benchmark
  6. ProfileModels(['yolo11n.yaml', 'yolov8s.yaml']).profile()
  7. benchmark(model='yolo11n.pt', imgsz=160)
  8. Format | `format=argument` | Model
  9. --- | --- | ---
  10. PyTorch | - | yolo11n.pt
  11. TorchScript | `torchscript` | yolo11n.torchscript
  12. ONNX | `onnx` | yolo11n.onnx
  13. OpenVINO | `openvino` | yolo11n_openvino_model/
  14. TensorRT | `engine` | yolo11n.engine
  15. CoreML | `coreml` | yolo11n.mlpackage
  16. TensorFlow SavedModel | `saved_model` | yolo11n_saved_model/
  17. TensorFlow GraphDef | `pb` | yolo11n.pb
  18. TensorFlow Lite | `tflite` | yolo11n.tflite
  19. TensorFlow Edge TPU | `edgetpu` | yolo11n_edgetpu.tflite
  20. TensorFlow.js | `tfjs` | yolo11n_web_model/
  21. PaddlePaddle | `paddle` | yolo11n_paddle_model/
  22. MNN | `mnn` | yolo11n.mnn
  23. NCNN | `ncnn` | yolo11n_ncnn_model/
  24. RKNN | `rknn` | yolo11n_rknn_model/
  25. """
  26. import glob
  27. import os
  28. import platform
  29. import re
  30. import shutil
  31. import time
  32. from pathlib import Path
  33. import numpy as np
  34. import torch.cuda
  35. import yaml
  36. from ultralytics import YOLO, YOLOWorld
  37. from ultralytics.cfg import TASK2DATA, TASK2METRIC
  38. from ultralytics.engine.exporter import export_formats
  39. from ultralytics.utils import ARM64, ASSETS, LINUX, LOGGER, MACOS, TQDM, WEIGHTS_DIR
  40. from ultralytics.utils.checks import IS_PYTHON_3_13, check_imgsz, check_requirements, check_yolo, is_rockchip
  41. from ultralytics.utils.downloads import safe_download
  42. from ultralytics.utils.files import file_size
  43. from ultralytics.utils.torch_utils import get_cpu_info, select_device
  44. def benchmark(
  45. model=WEIGHTS_DIR / "yolo11n.pt",
  46. data=None,
  47. imgsz=160,
  48. half=False,
  49. int8=False,
  50. device="cpu",
  51. verbose=False,
  52. eps=1e-3,
  53. format="",
  54. ):
  55. """
  56. Benchmark a YOLO model across different formats for speed and accuracy.
  57. Args:
  58. model (str | Path): Path to the model file or directory.
  59. data (str | None): Dataset to evaluate on, inherited from TASK2DATA if not passed.
  60. imgsz (int): Image size for the benchmark.
  61. half (bool): Use half-precision for the model if True.
  62. int8 (bool): Use int8-precision for the model if True.
  63. device (str): Device to run the benchmark on, either 'cpu' or 'cuda'.
  64. verbose (bool | float): If True or a float, assert benchmarks pass with given metric.
  65. eps (float): Epsilon value for divide by zero prevention.
  66. format (str): Export format for benchmarking. If not supplied all formats are benchmarked.
  67. Returns:
  68. (pandas.DataFrame): A pandas DataFrame with benchmark results for each format, including file size, metric,
  69. and inference time.
  70. Examples:
  71. Benchmark a YOLO model with default settings:
  72. >>> from ultralytics.utils.benchmarks import benchmark
  73. >>> benchmark(model="yolo11n.pt", imgsz=640)
  74. """
  75. imgsz = check_imgsz(imgsz)
  76. assert imgsz[0] == imgsz[1] if isinstance(imgsz, list) else True, "benchmark() only supports square imgsz."
  77. import pandas as pd # scope for faster 'import ultralytics'
  78. pd.options.display.max_columns = 10
  79. pd.options.display.width = 120
  80. device = select_device(device, verbose=False)
  81. if isinstance(model, (str, Path)):
  82. model = YOLO(model)
  83. is_end2end = getattr(model.model.model[-1], "end2end", False)
  84. data = data or TASK2DATA[model.task] # task to dataset, i.e. coco8.yaml for task=detect
  85. key = TASK2METRIC[model.task] # task to metric, i.e. metrics/mAP50-95(B) for task=detect
  86. y = []
  87. t0 = time.time()
  88. format_arg = format.lower()
  89. if format_arg:
  90. formats = frozenset(export_formats()["Argument"])
  91. assert format in formats, f"Expected format to be one of {formats}, but got '{format_arg}'."
  92. for i, (name, format, suffix, cpu, gpu, _) in enumerate(zip(*export_formats().values())):
  93. emoji, filename = "❌", None # export defaults
  94. try:
  95. if format_arg and format_arg != format:
  96. continue
  97. # Checks
  98. if i == 7: # TF GraphDef
  99. assert model.task != "obb", "TensorFlow GraphDef not supported for OBB task"
  100. elif i == 9: # Edge TPU
  101. assert LINUX and not ARM64, "Edge TPU export only supported on non-aarch64 Linux"
  102. elif i in {5, 10}: # CoreML and TF.js
  103. assert MACOS or (LINUX and not ARM64), (
  104. "CoreML and TF.js export only supported on macOS and non-aarch64 Linux"
  105. )
  106. if i in {5}: # CoreML
  107. assert not IS_PYTHON_3_13, "CoreML not supported on Python 3.13"
  108. if i in {6, 7, 8, 9, 10}: # TF SavedModel, TF GraphDef, and TFLite, TF EdgeTPU and TF.js
  109. assert not isinstance(model, YOLOWorld), "YOLOWorldv2 TensorFlow exports not supported by onnx2tf yet"
  110. # assert not IS_PYTHON_MINIMUM_3_12, "TFLite exports not supported on Python>=3.12 yet"
  111. if i == 11: # Paddle
  112. assert not isinstance(model, YOLOWorld), "YOLOWorldv2 Paddle exports not supported yet"
  113. assert model.task != "obb", "Paddle OBB bug https://github.com/PaddlePaddle/Paddle/issues/72024"
  114. assert not is_end2end, "End-to-end models not supported by PaddlePaddle yet"
  115. assert LINUX or MACOS, "Windows Paddle exports not supported yet"
  116. if i == 12: # MNN
  117. assert not isinstance(model, YOLOWorld), "YOLOWorldv2 MNN exports not supported yet"
  118. if i == 13: # NCNN
  119. assert not isinstance(model, YOLOWorld), "YOLOWorldv2 NCNN exports not supported yet"
  120. if i == 14: # IMX
  121. assert not is_end2end
  122. assert not isinstance(model, YOLOWorld), "YOLOWorldv2 IMX exports not supported"
  123. assert model.task == "detect", "IMX only supported for detection task"
  124. assert "C2f" in model.__str__(), "IMX only supported for YOLOv8" # TODO: enable for YOLO11
  125. if i == 15: # RKNN
  126. assert not isinstance(model, YOLOWorld), "YOLOWorldv2 RKNN exports not supported yet"
  127. assert not is_end2end, "End-to-end models not supported by RKNN yet"
  128. assert LINUX, "RKNN only supported on Linux"
  129. assert not is_rockchip(), "RKNN Inference only supported on Rockchip devices"
  130. if "cpu" in device.type:
  131. assert cpu, "inference not supported on CPU"
  132. if "cuda" in device.type:
  133. assert gpu, "inference not supported on GPU"
  134. # Export
  135. if format == "-":
  136. filename = model.pt_path or model.ckpt_path or model.model_name
  137. exported_model = model # PyTorch format
  138. else:
  139. filename = model.export(
  140. imgsz=imgsz, format=format, half=half, int8=int8, data=data, device=device, verbose=False
  141. )
  142. exported_model = YOLO(filename, task=model.task)
  143. assert suffix in str(filename), "export failed"
  144. emoji = "❎" # indicates export succeeded
  145. # Predict
  146. assert model.task != "pose" or i != 7, "GraphDef Pose inference is not supported"
  147. assert i not in {9, 10}, "inference not supported" # Edge TPU and TF.js are unsupported
  148. assert i != 5 or platform.system() == "Darwin", "inference only supported on macOS>=10.13" # CoreML
  149. if i in {13}:
  150. assert not is_end2end, "End-to-end torch.topk operation is not supported for NCNN prediction yet"
  151. exported_model.predict(ASSETS / "bus.jpg", imgsz=imgsz, device=device, half=half, verbose=False)
  152. # Validate
  153. results = exported_model.val(
  154. data=data, batch=1, imgsz=imgsz, plots=False, device=device, half=half, int8=int8, verbose=False
  155. )
  156. metric, speed = results.results_dict[key], results.speed["inference"]
  157. fps = round(1000 / (speed + eps), 2) # frames per second
  158. y.append([name, "βœ…", round(file_size(filename), 1), round(metric, 4), round(speed, 2), fps])
  159. except Exception as e:
  160. if verbose:
  161. assert type(e) is AssertionError, f"Benchmark failure for {name}: {e}"
  162. LOGGER.error(f"Benchmark failure for {name}: {e}")
  163. y.append([name, emoji, round(file_size(filename), 1), None, None, None]) # mAP, t_inference
  164. # Print results
  165. check_yolo(device=device) # print system info
  166. df = pd.DataFrame(y, columns=["Format", "Status❔", "Size (MB)", key, "Inference time (ms/im)", "FPS"])
  167. name = model.model_name
  168. dt = time.time() - t0
  169. legend = "Benchmarks legend: - βœ… Success - ❎ Export passed but validation failed - ❌️ Export failed"
  170. s = f"\nBenchmarks complete for {name} on {data} at imgsz={imgsz} ({dt:.2f}s)\n{legend}\n{df.fillna('-')}\n"
  171. LOGGER.info(s)
  172. with open("benchmarks.log", "a", errors="ignore", encoding="utf-8") as f:
  173. f.write(s)
  174. if verbose and isinstance(verbose, float):
  175. metrics = df[key].array # values to compare to floor
  176. floor = verbose # minimum metric floor to pass, i.e. = 0.29 mAP for YOLOv5n
  177. assert all(x > floor for x in metrics if pd.notna(x)), f"Benchmark failure: metric(s) < floor {floor}"
  178. return df
  179. class RF100Benchmark:
  180. """
  181. Benchmark YOLO model performance across various formats for speed and accuracy.
  182. This class provides functionality to benchmark YOLO models on the RF100 dataset collection.
  183. Attributes:
  184. ds_names (List[str]): Names of datasets used for benchmarking.
  185. ds_cfg_list (List[Path]): List of paths to dataset configuration files.
  186. rf (Roboflow): Roboflow instance for accessing datasets.
  187. val_metrics (List[str]): Metrics used for validation.
  188. Methods:
  189. set_key: Set Roboflow API key for accessing datasets.
  190. parse_dataset: Parse dataset links and download datasets.
  191. fix_yaml: Fix train and validation paths in YAML files.
  192. evaluate: Evaluate model performance on validation results.
  193. """
  194. def __init__(self):
  195. """Initialize the RF100Benchmark class for benchmarking YOLO model performance across various formats."""
  196. self.ds_names = []
  197. self.ds_cfg_list = []
  198. self.rf = None
  199. self.val_metrics = ["class", "images", "targets", "precision", "recall", "map50", "map95"]
  200. def set_key(self, api_key):
  201. """
  202. Set Roboflow API key for processing.
  203. Args:
  204. api_key (str): The API key.
  205. Examples:
  206. Set the Roboflow API key for accessing datasets:
  207. >>> benchmark = RF100Benchmark()
  208. >>> benchmark.set_key("your_roboflow_api_key")
  209. """
  210. check_requirements("roboflow")
  211. from roboflow import Roboflow
  212. self.rf = Roboflow(api_key=api_key)
  213. def parse_dataset(self, ds_link_txt="datasets_links.txt"):
  214. """
  215. Parse dataset links and download datasets.
  216. Args:
  217. ds_link_txt (str): Path to the file containing dataset links.
  218. Returns:
  219. ds_names (List[str]): List of dataset names.
  220. ds_cfg_list (List[Path]): List of paths to dataset configuration files.
  221. Examples:
  222. >>> benchmark = RF100Benchmark()
  223. >>> benchmark.set_key("api_key")
  224. >>> benchmark.parse_dataset("datasets_links.txt")
  225. """
  226. (shutil.rmtree("rf-100"), os.mkdir("rf-100")) if os.path.exists("rf-100") else os.mkdir("rf-100")
  227. os.chdir("rf-100")
  228. os.mkdir("ultralytics-benchmarks")
  229. safe_download("https://github.com/ultralytics/assets/releases/download/v0.0.0/datasets_links.txt")
  230. with open(ds_link_txt, encoding="utf-8") as file:
  231. for line in file:
  232. try:
  233. _, url, workspace, project, version = re.split("/+", line.strip())
  234. self.ds_names.append(project)
  235. proj_version = f"{project}-{version}"
  236. if not Path(proj_version).exists():
  237. self.rf.workspace(workspace).project(project).version(version).download("yolov8")
  238. else:
  239. LOGGER.info("Dataset already downloaded.")
  240. self.ds_cfg_list.append(Path.cwd() / proj_version / "data.yaml")
  241. except Exception:
  242. continue
  243. return self.ds_names, self.ds_cfg_list
  244. @staticmethod
  245. def fix_yaml(path):
  246. """Fix the train and validation paths in a given YAML file."""
  247. with open(path, encoding="utf-8") as file:
  248. yaml_data = yaml.safe_load(file)
  249. yaml_data["train"] = "train/images"
  250. yaml_data["val"] = "valid/images"
  251. with open(path, "w", encoding="utf-8") as file:
  252. yaml.safe_dump(yaml_data, file)
  253. def evaluate(self, yaml_path, val_log_file, eval_log_file, list_ind):
  254. """
  255. Evaluate model performance on validation results.
  256. Args:
  257. yaml_path (str): Path to the YAML configuration file.
  258. val_log_file (str): Path to the validation log file.
  259. eval_log_file (str): Path to the evaluation log file.
  260. list_ind (int): Index of the current dataset in the list.
  261. Returns:
  262. (float): The mean average precision (mAP) value for the evaluated model.
  263. Examples:
  264. Evaluate a model on a specific dataset
  265. >>> benchmark = RF100Benchmark()
  266. >>> benchmark.evaluate("path/to/data.yaml", "path/to/val_log.txt", "path/to/eval_log.txt", 0)
  267. """
  268. skip_symbols = ["πŸš€", "⚠️", "πŸ’‘", "❌"]
  269. with open(yaml_path, encoding="utf-8") as stream:
  270. class_names = yaml.safe_load(stream)["names"]
  271. with open(val_log_file, encoding="utf-8") as f:
  272. lines = f.readlines()
  273. eval_lines = []
  274. for line in lines:
  275. if any(symbol in line for symbol in skip_symbols):
  276. continue
  277. entries = line.split(" ")
  278. entries = list(filter(lambda val: val != "", entries))
  279. entries = [e.strip("\n") for e in entries]
  280. eval_lines.extend(
  281. {
  282. "class": entries[0],
  283. "images": entries[1],
  284. "targets": entries[2],
  285. "precision": entries[3],
  286. "recall": entries[4],
  287. "map50": entries[5],
  288. "map95": entries[6],
  289. }
  290. for e in entries
  291. if e in class_names or (e == "all" and "(AP)" not in entries and "(AR)" not in entries)
  292. )
  293. map_val = 0.0
  294. if len(eval_lines) > 1:
  295. LOGGER.info("Multiple dicts found")
  296. for lst in eval_lines:
  297. if lst["class"] == "all":
  298. map_val = lst["map50"]
  299. else:
  300. LOGGER.info("Single dict found")
  301. map_val = [res["map50"] for res in eval_lines][0]
  302. with open(eval_log_file, "a", encoding="utf-8") as f:
  303. f.write(f"{self.ds_names[list_ind]}: {map_val}\n")
  304. class ProfileModels:
  305. """
  306. ProfileModels class for profiling different models on ONNX and TensorRT.
  307. This class profiles the performance of different models, returning results such as model speed and FLOPs.
  308. Attributes:
  309. paths (List[str]): Paths of the models to profile.
  310. num_timed_runs (int): Number of timed runs for the profiling.
  311. num_warmup_runs (int): Number of warmup runs before profiling.
  312. min_time (float): Minimum number of seconds to profile for.
  313. imgsz (int): Image size used in the models.
  314. half (bool): Flag to indicate whether to use FP16 half-precision for TensorRT profiling.
  315. trt (bool): Flag to indicate whether to profile using TensorRT.
  316. device (torch.device): Device used for profiling.
  317. Methods:
  318. profile: Profiles the models and prints the result.
  319. get_files: Gets all relevant model files.
  320. get_onnx_model_info: Extracts metadata from an ONNX model.
  321. iterative_sigma_clipping: Applies sigma clipping to remove outliers.
  322. profile_tensorrt_model: Profiles a TensorRT model.
  323. profile_onnx_model: Profiles an ONNX model.
  324. generate_table_row: Generates a table row with model metrics.
  325. generate_results_dict: Generates a dictionary of profiling results.
  326. print_table: Prints a formatted table of results.
  327. Examples:
  328. Profile models and print results
  329. >>> from ultralytics.utils.benchmarks import ProfileModels
  330. >>> profiler = ProfileModels(["yolo11n.yaml", "yolov8s.yaml"], imgsz=640)
  331. >>> profiler.profile()
  332. """
  333. def __init__(
  334. self,
  335. paths: list,
  336. num_timed_runs=100,
  337. num_warmup_runs=10,
  338. min_time=60,
  339. imgsz=640,
  340. half=True,
  341. trt=True,
  342. device=None,
  343. ):
  344. """
  345. Initialize the ProfileModels class for profiling models.
  346. Args:
  347. paths (List[str]): List of paths of the models to be profiled.
  348. num_timed_runs (int): Number of timed runs for the profiling.
  349. num_warmup_runs (int): Number of warmup runs before the actual profiling starts.
  350. min_time (float): Minimum time in seconds for profiling a model.
  351. imgsz (int): Size of the image used during profiling.
  352. half (bool): Flag to indicate whether to use FP16 half-precision for TensorRT profiling.
  353. trt (bool): Flag to indicate whether to profile using TensorRT.
  354. device (torch.device | None): Device used for profiling. If None, it is determined automatically.
  355. Notes:
  356. FP16 'half' argument option removed for ONNX as slower on CPU than FP32.
  357. Examples:
  358. Initialize and profile models
  359. >>> from ultralytics.utils.benchmarks import ProfileModels
  360. >>> profiler = ProfileModels(["yolo11n.yaml", "yolov8s.yaml"], imgsz=640)
  361. >>> profiler.profile()
  362. """
  363. self.paths = paths
  364. self.num_timed_runs = num_timed_runs
  365. self.num_warmup_runs = num_warmup_runs
  366. self.min_time = min_time
  367. self.imgsz = imgsz
  368. self.half = half
  369. self.trt = trt # run TensorRT profiling
  370. self.device = device or torch.device(0 if torch.cuda.is_available() else "cpu")
  371. def profile(self):
  372. """
  373. Profile YOLO models for speed and accuracy across various formats including ONNX and TensorRT.
  374. Returns:
  375. (List[Dict]): List of dictionaries containing profiling results for each model.
  376. Examples:
  377. Profile models and print results
  378. >>> from ultralytics.utils.benchmarks import ProfileModels
  379. >>> profiler = ProfileModels(["yolo11n.yaml", "yolov8s.yaml"])
  380. >>> results = profiler.profile()
  381. """
  382. files = self.get_files()
  383. if not files:
  384. LOGGER.warning("No matching *.pt or *.onnx files found.")
  385. return
  386. table_rows = []
  387. output = []
  388. for file in files:
  389. engine_file = file.with_suffix(".engine")
  390. if file.suffix in {".pt", ".yaml", ".yml"}:
  391. model = YOLO(str(file))
  392. model.fuse() # to report correct params and GFLOPs in model.info()
  393. model_info = model.info()
  394. if self.trt and self.device.type != "cpu" and not engine_file.is_file():
  395. engine_file = model.export(
  396. format="engine",
  397. half=self.half,
  398. imgsz=self.imgsz,
  399. device=self.device,
  400. verbose=False,
  401. )
  402. onnx_file = model.export(
  403. format="onnx",
  404. imgsz=self.imgsz,
  405. device=self.device,
  406. verbose=False,
  407. )
  408. elif file.suffix == ".onnx":
  409. model_info = self.get_onnx_model_info(file)
  410. onnx_file = file
  411. else:
  412. continue
  413. t_engine = self.profile_tensorrt_model(str(engine_file))
  414. t_onnx = self.profile_onnx_model(str(onnx_file))
  415. table_rows.append(self.generate_table_row(file.stem, t_onnx, t_engine, model_info))
  416. output.append(self.generate_results_dict(file.stem, t_onnx, t_engine, model_info))
  417. self.print_table(table_rows)
  418. return output
  419. def get_files(self):
  420. """
  421. Return a list of paths for all relevant model files given by the user.
  422. Returns:
  423. (List[Path]): List of Path objects for the model files.
  424. """
  425. files = []
  426. for path in self.paths:
  427. path = Path(path)
  428. if path.is_dir():
  429. extensions = ["*.pt", "*.onnx", "*.yaml"]
  430. files.extend([file for ext in extensions for file in glob.glob(str(path / ext))])
  431. elif path.suffix in {".pt", ".yaml", ".yml"}: # add non-existing
  432. files.append(str(path))
  433. else:
  434. files.extend(glob.glob(str(path)))
  435. LOGGER.info(f"Profiling: {sorted(files)}")
  436. return [Path(file) for file in sorted(files)]
  437. @staticmethod
  438. def get_onnx_model_info(onnx_file: str):
  439. """Extracts metadata from an ONNX model file including parameters, GFLOPs, and input shape."""
  440. return 0.0, 0.0, 0.0, 0.0 # return (num_layers, num_params, num_gradients, num_flops)
  441. @staticmethod
  442. def iterative_sigma_clipping(data, sigma=2, max_iters=3):
  443. """
  444. Apply iterative sigma clipping to data to remove outliers.
  445. Args:
  446. data (numpy.ndarray): Input data array.
  447. sigma (float): Number of standard deviations to use for clipping.
  448. max_iters (int): Maximum number of iterations for the clipping process.
  449. Returns:
  450. (numpy.ndarray): Clipped data array with outliers removed.
  451. """
  452. data = np.array(data)
  453. for _ in range(max_iters):
  454. mean, std = np.mean(data), np.std(data)
  455. clipped_data = data[(data > mean - sigma * std) & (data < mean + sigma * std)]
  456. if len(clipped_data) == len(data):
  457. break
  458. data = clipped_data
  459. return data
  460. def profile_tensorrt_model(self, engine_file: str, eps: float = 1e-3):
  461. """
  462. Profile YOLO model performance with TensorRT, measuring average run time and standard deviation.
  463. Args:
  464. engine_file (str): Path to the TensorRT engine file.
  465. eps (float): Small epsilon value to prevent division by zero.
  466. Returns:
  467. mean_time (float): Mean inference time in milliseconds.
  468. std_time (float): Standard deviation of inference time in milliseconds.
  469. """
  470. if not self.trt or not Path(engine_file).is_file():
  471. return 0.0, 0.0
  472. # Model and input
  473. model = YOLO(engine_file)
  474. input_data = np.zeros((self.imgsz, self.imgsz, 3), dtype=np.uint8) # use uint8 for Classify
  475. # Warmup runs
  476. elapsed = 0.0
  477. for _ in range(3):
  478. start_time = time.time()
  479. for _ in range(self.num_warmup_runs):
  480. model(input_data, imgsz=self.imgsz, verbose=False)
  481. elapsed = time.time() - start_time
  482. # Compute number of runs as higher of min_time or num_timed_runs
  483. num_runs = max(round(self.min_time / (elapsed + eps) * self.num_warmup_runs), self.num_timed_runs * 50)
  484. # Timed runs
  485. run_times = []
  486. for _ in TQDM(range(num_runs), desc=engine_file):
  487. results = model(input_data, imgsz=self.imgsz, verbose=False)
  488. run_times.append(results[0].speed["inference"]) # Convert to milliseconds
  489. run_times = self.iterative_sigma_clipping(np.array(run_times), sigma=2, max_iters=3) # sigma clipping
  490. return np.mean(run_times), np.std(run_times)
  491. def profile_onnx_model(self, onnx_file: str, eps: float = 1e-3):
  492. """
  493. Profile an ONNX model, measuring average inference time and standard deviation across multiple runs.
  494. Args:
  495. onnx_file (str): Path to the ONNX model file.
  496. eps (float): Small epsilon value to prevent division by zero.
  497. Returns:
  498. mean_time (float): Mean inference time in milliseconds.
  499. std_time (float): Standard deviation of inference time in milliseconds.
  500. """
  501. check_requirements("onnxruntime")
  502. import onnxruntime as ort
  503. # Session with either 'TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider'
  504. sess_options = ort.SessionOptions()
  505. sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
  506. sess_options.intra_op_num_threads = 8 # Limit the number of threads
  507. sess = ort.InferenceSession(onnx_file, sess_options, providers=["CPUExecutionProvider"])
  508. input_tensor = sess.get_inputs()[0]
  509. input_type = input_tensor.type
  510. dynamic = not all(isinstance(dim, int) and dim >= 0 for dim in input_tensor.shape) # dynamic input shape
  511. input_shape = (1, 3, self.imgsz, self.imgsz) if dynamic else input_tensor.shape
  512. # Mapping ONNX datatype to numpy datatype
  513. if "float16" in input_type:
  514. input_dtype = np.float16
  515. elif "float" in input_type:
  516. input_dtype = np.float32
  517. elif "double" in input_type:
  518. input_dtype = np.float64
  519. elif "int64" in input_type:
  520. input_dtype = np.int64
  521. elif "int32" in input_type:
  522. input_dtype = np.int32
  523. else:
  524. raise ValueError(f"Unsupported ONNX datatype {input_type}")
  525. input_data = np.random.rand(*input_shape).astype(input_dtype)
  526. input_name = input_tensor.name
  527. output_name = sess.get_outputs()[0].name
  528. # Warmup runs
  529. elapsed = 0.0
  530. for _ in range(3):
  531. start_time = time.time()
  532. for _ in range(self.num_warmup_runs):
  533. sess.run([output_name], {input_name: input_data})
  534. elapsed = time.time() - start_time
  535. # Compute number of runs as higher of min_time or num_timed_runs
  536. num_runs = max(round(self.min_time / (elapsed + eps) * self.num_warmup_runs), self.num_timed_runs)
  537. # Timed runs
  538. run_times = []
  539. for _ in TQDM(range(num_runs), desc=onnx_file):
  540. start_time = time.time()
  541. sess.run([output_name], {input_name: input_data})
  542. run_times.append((time.time() - start_time) * 1000) # Convert to milliseconds
  543. run_times = self.iterative_sigma_clipping(np.array(run_times), sigma=2, max_iters=5) # sigma clipping
  544. return np.mean(run_times), np.std(run_times)
  545. def generate_table_row(self, model_name, t_onnx, t_engine, model_info):
  546. """
  547. Generate a table row string with model performance metrics.
  548. Args:
  549. model_name (str): Name of the model.
  550. t_onnx (tuple): ONNX model inference time statistics (mean, std).
  551. t_engine (tuple): TensorRT engine inference time statistics (mean, std).
  552. model_info (tuple): Model information (layers, params, gradients, flops).
  553. Returns:
  554. (str): Formatted table row string with model metrics.
  555. """
  556. layers, params, gradients, flops = model_info
  557. return (
  558. f"| {model_name:18s} | {self.imgsz} | - | {t_onnx[0]:.1f}Β±{t_onnx[1]:.1f} ms | {t_engine[0]:.1f}Β±"
  559. f"{t_engine[1]:.1f} ms | {params / 1e6:.1f} | {flops:.1f} |"
  560. )
  561. @staticmethod
  562. def generate_results_dict(model_name, t_onnx, t_engine, model_info):
  563. """
  564. Generate a dictionary of profiling results.
  565. Args:
  566. model_name (str): Name of the model.
  567. t_onnx (tuple): ONNX model inference time statistics (mean, std).
  568. t_engine (tuple): TensorRT engine inference time statistics (mean, std).
  569. model_info (tuple): Model information (layers, params, gradients, flops).
  570. Returns:
  571. (dict): Dictionary containing profiling results.
  572. """
  573. layers, params, gradients, flops = model_info
  574. return {
  575. "model/name": model_name,
  576. "model/parameters": params,
  577. "model/GFLOPs": round(flops, 3),
  578. "model/speed_ONNX(ms)": round(t_onnx[0], 3),
  579. "model/speed_TensorRT(ms)": round(t_engine[0], 3),
  580. }
  581. @staticmethod
  582. def print_table(table_rows):
  583. """
  584. Print a formatted table of model profiling results.
  585. Args:
  586. table_rows (List[str]): List of formatted table row strings.
  587. """
  588. gpu = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "GPU"
  589. headers = [
  590. "Model",
  591. "size<br><sup>(pixels)",
  592. "mAP<sup>val<br>50-95",
  593. f"Speed<br><sup>CPU ({get_cpu_info()}) ONNX<br>(ms)",
  594. f"Speed<br><sup>{gpu} TensorRT<br>(ms)",
  595. "params<br><sup>(M)",
  596. "FLOPs<br><sup>(B)",
  597. ]
  598. header = "|" + "|".join(f" {h} " for h in headers) + "|"
  599. separator = "|" + "|".join("-" * (len(h) + 2) for h in headers) + "|"
  600. LOGGER.info(f"\n\n{header}")
  601. LOGGER.info(separator)
  602. for row in table_rows:
  603. LOGGER.info(row)
Discard
Tip!

Press p or to see the previous file or, n or to see the next file