HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

Allow TensorRT inference without metadata (#10152)

Browse source 
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
1 changed files with 5 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

7
ultralytics/nn/autobackend.py
View file

@ -234,8 +234,11 @@ class AutoBackend(nn.Module):
logger = trt.Logger(trt.Logger.INFO) logger = trt.Logger(trt.Logger.INFO)
# Read file # Read file
with open(w, "rb") as f, trt.Runtime(logger) as runtime: with open(w, "rb") as f, trt.Runtime(logger) as runtime:
meta_len = int.from_bytes(f.read(4), byteorder="little") # read metadata length try:
metadata = json.loads(f.read(meta_len).decode("utf-8")) # read metadata 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 = runtime.deserialize_cuda_engine(f.read()) # read engine
# Model context # Model context