Fix OpenVINO missing batch-size bug (#9264)

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

View File

@ -202,13 +202,10 @@ class AutoBackend(nn.Module):
ov_model = core.read_model(model=str(w), weights=w.with_suffix(".bin")) ov_model = core.read_model(model=str(w), weights=w.with_suffix(".bin"))
if ov_model.get_parameters()[0].get_layout().empty: if ov_model.get_parameters()[0].get_layout().empty:
ov_model.get_parameters()[0].set_layout(ov.Layout("NCHW")) 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' # OpenVINO inference modes are 'LATENCY', 'THROUGHPUT' (not recommended), or 'CUMULATIVE_THROUGHPUT'
inference_mode = "CUMULATIVE_THROUGHPUT" if batch > 1 else "LATENCY" 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_compiled_model = core.compile_model(
ov_model, ov_model,
device_name="AUTO", # AUTO selects best available device, do not modify device_name="AUTO", # AUTO selects best available device, do not modify