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

ultralytics 8.0.195 NVIDIA Triton Inference Server support (#5257)

Browse source 
Co-authored-by: TheConstant3 <46416203+TheConstant3@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Glenn Jocher 2023-10-07 19:26:35 +02:00 • committed by GitHub
parent 40e3923cfc
commit c7aa83da31
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
21 changed files with 349 additions and 98 deletions
Download patch file Download diff file Expand all files Collapse all files

20
ultralytics/nn/autobackend.py
View file

@ -7,7 +7,6 @@ import platform
import zipfile
from collections import OrderedDict, namedtuple
from pathlib import Path
from urllib.parse import urlparse
import cv2
import numpy as np
@ -32,8 +31,8 @@ def check_class_names(names):
raise KeyError(f'{n}-class dataset requires class indices 0-{n - 1}, but you have invalid class indices '
f'{min(names.keys())}-{max(names.keys())} defined in your dataset YAML.')
if isinstance(names[0], str) and names[0].startswith('n0'): # imagenet class codes, i.e. 'n01440764'
map = yaml_load(ROOT / 'cfg/datasets/ImageNet.yaml')['map'] # human-readable names
names = {k: map[v] for k, v in names.items()}
names_map = yaml_load(ROOT / 'cfg/datasets/ImageNet.yaml')['map'] # human-readable names
names = {k: names_map[v] for k, v in names.items()}
return names
@ -274,13 +273,9 @@ class AutoBackend(nn.Module):
net.load_model(str(w.with_suffix('.bin')))
metadata = w.parent / 'metadata.yaml'
elif triton: # NVIDIA Triton Inference Server
"""TODO
check_requirements('tritonclient[all]')
from utils.triton import TritonRemoteModel
model = TritonRemoteModel(url=w)
nhwc = model.runtime.startswith("tensorflow")
"""
raise NotImplementedError('Triton Inference Server is not currently supported.')
from ultralytics.utils.triton import TritonRemoteModel
model = TritonRemoteModel(w)
else:
from ultralytics.engine.exporter import export_formats
raise TypeError(f"model='{w}' is not a supported model format. "
@ -395,6 +390,7 @@ class AutoBackend(nn.Module):
ex.extract(output_name, mat_out)
y.append(np.array(mat_out)[None])
elif self.triton: # NVIDIA Triton Inference Server
im = im.cpu().numpy() # torch to numpy
y = self.model(im)
else: # TensorFlow (SavedModel, GraphDef, Lite, Edge TPU)
im = im.cpu().numpy()
@ -498,6 +494,8 @@ class AutoBackend(nn.Module):
if any(types):
triton = False
else:
url = urlparse(p) # if url may be Triton inference server
triton = all([any(s in url.scheme for s in ['http', 'grpc']), url.netloc])
from urllib.parse import urlsplit
url = urlsplit(p)
triton = url.netloc and url.path and url.scheme in {'http', 'grfc'}
return types + [triton]