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

ultralytics 8.0.202 sort Triton model outputs (#5945)

Browse source 
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Mike Tune <mtuneoff@gmail.com>
This commit is contained in:
Glenn Jocher 2023-10-26 20:32:51 +02:00 • committed by GitHub
parent a05edfbc27
commit e58db228c2
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
10 changed files with 138 additions and 133 deletions
Download patch file Download diff file Expand all files Collapse all files

7
ultralytics/utils/triton.py
View file

@ -53,10 +53,13 @@ class TritonRemoteModel:
self.triton_client = client.InferenceServerClient(url=self.url, verbose=False, ssl=False)
config = self.triton_client.get_model_config(endpoint, as_json=True)['config']
# Sort output names alphabetically, i.e. 'output0', 'output1', etc.
config['output'] = sorted(config['output'], key=lambda x: x.get('name'))
# Define model attributes
type_map = {'TYPE_FP32': np.float32, 'TYPE_FP16': np.float16, 'TYPE_UINT8': np.uint8}
self.InferRequestedOutput = client.InferRequestedOutput
self.InferInput = client.InferInput
type_map = {'TYPE_FP32': np.float32, 'TYPE_FP16': np.float16, 'TYPE_UINT8': np.uint8}
self.input_formats = [x['data_type'] for x in config['input']]
self.np_input_formats = [type_map[x] for x in self.input_formats]
self.input_names = [x['name'] for x in config['input']]