Allow TensorRT inference without metadata (#10152)

Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Nazmi 2024-04-19 02:27:03 +01:00 committed by GitHub
parent 5321d58594
commit 5e99506c0f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -234,8 +234,11 @@ class AutoBackend(nn.Module):
logger = trt.Logger(trt.Logger.INFO)
# Read file
with open(w, "rb") as f, trt.Runtime(logger) as runtime:
try:
meta_len = int.from_bytes(f.read(4), byteorder="little") # read metadata length
metadata = json.loads(f.read(meta_len).decode("utf-8")) # read metadata
except UnicodeDecodeError:
f.seek(0) # engine file may lack embedded Ultralytics metadata
model = runtime.deserialize_cuda_engine(f.read()) # read engine
# Model context