xclele/yolov10
1
0
Fork 0
mirror of https://github.com/THU-MIG/yolov10.git synced 2025-05-22 21:04:21 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix OpenVINO missing batch-size bug (#9264)

Browse Source
This commit is contained in:
Glenn Jocher 2024-03-24 05:56:37 +01:00 committed by GitHub
parent 8d31c50856
commit fc6c66a4a4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
ultralytics/nn/autobackend.py
View File

@ -202,13 +202,10 @@ class AutoBackend(nn.Module):
ov_model = core.read_model(model=str(w), weights=w.with_suffix(".bin"))
if ov_model.get_parameters()[0].get_layout().empty:
ov_model.get_parameters()[0].set_layout(ov.Layout("NCHW"))
batch_dim = ov.get_batch(ov_model)
if batch_dim.is_static:
batch_size = batch_dim.get_length()
# OpenVINO inference modes are 'LATENCY', 'THROUGHPUT' (not recommended), or 'CUMULATIVE_THROUGHPUT'
inference_mode = "CUMULATIVE_THROUGHPUT" if batch > 1 else "LATENCY"
LOGGER.info(f"Using OpenVINO {inference_mode} mode for batch-size={batch_size} inference...")
LOGGER.info(f"Using OpenVINO {inference_mode} mode for batch={batch} inference...")
ov_compiled_model = core.compile_model(
ov_model,
device_name="AUTO", # AUTO selects best available device, do not modify