29
29
30
30
31
31
def init_predictor (args ):
32
- infer_dir = args .trt_inference_dir
32
+ infer_dir = args .inference_dir
33
33
assert os .path .isdir (
34
34
infer_dir ), f'inference_dir = "{ infer_dir } " is not a directory'
35
35
pdiparams_path = glob .glob (os .path .join (infer_dir , '*.pdiparams' ))
@@ -41,7 +41,7 @@ def init_predictor(args):
41
41
predictor_config = Config (pdmodel_path [0 ], pdiparams_path [0 ])
42
42
predictor_config .enable_memory_optim ()
43
43
predictor_config .enable_use_gpu (0 , args .device )
44
- precision = args .trt_precision
44
+ precision = args .precision
45
45
max_batch_size = args .batch_size
46
46
assert precision in ['FP32' , 'FP16' , 'INT8' ], \
47
47
'precision should be FP32/FP16/INT8'
@@ -54,12 +54,17 @@ def init_predictor(args):
54
54
else :
55
55
raise NotImplementedError
56
56
predictor_config .enable_tensorrt_engine (
57
- workspace_size = args .trt_workspace_size ,
57
+ workspace_size = args .workspace_size ,
58
58
max_batch_size = max_batch_size ,
59
- min_subgraph_size = args .trt_min_subgraph_size ,
59
+ min_subgraph_size = args .min_subgraph_size ,
60
60
precision_mode = precision_mode ,
61
- use_static = args .trt_use_static ,
62
- use_calib_mode = args .trt_use_calib_mode )
61
+ use_static = args .use_static ,
62
+ use_calib_mode = args .use_calib_mode )
63
+ predictor_config .set_trt_dynamic_shape_info (
64
+ {"data" : (1 ,) + tuple (args .image_shape )},
65
+ {"data" : (args .batch_size ,) + tuple (args .image_shape )},
66
+ {"data" : (args .batch_size ,) + tuple (args .image_shape )},
67
+ )
63
68
predictor = create_predictor (predictor_config )
64
69
return predictor
65
70
@@ -140,7 +145,7 @@ def benchmark_dataset(args):
140
145
quantile = np .quantile (latency , [0.9 , 0.95 , 0.99 ])
141
146
142
147
statistics = {
143
- 'precision' : args .trt_precision ,
148
+ 'precision' : args .precision ,
144
149
'batch_size' : batch_size ,
145
150
'throughput' : total_images / (end - start ),
146
151
'accuracy' : correct_predict / total_images ,
@@ -189,7 +194,7 @@ def benchmark_synthetic(args):
189
194
quantile = np .quantile (latency , [0.9 , 0.95 , 0.99 ])
190
195
191
196
statistics = {
192
- 'precision' : args .trt_precision ,
197
+ 'precision' : args .precision ,
193
198
'batch_size' : batch_size ,
194
199
'throughput' : args .benchmark_steps * batch_size / (end - start ),
195
200
'eval_latency_avg' : np .mean (latency ),
@@ -200,11 +205,11 @@ def benchmark_synthetic(args):
200
205
return statistics
201
206
202
207
def main (args ):
203
- setup_dllogger (args .trt_log_path )
208
+ setup_dllogger (args .report_file )
204
209
if args .show_config :
205
210
print_args (args )
206
211
207
- if args .trt_use_synthetic :
212
+ if args .use_synthetic :
208
213
statistics = benchmark_synthetic (args )
209
214
else :
210
215
statistics = benchmark_dataset (args )
@@ -213,4 +218,4 @@ def main(args):
213
218
214
219
215
220
if __name__ == '__main__' :
216
- main (parse_args (including_trt = True ))
221
+ main (parse_args (script = 'inference' ))
0 commit comments