| 58 | |
| 59 | |
| 60 | def init_predictor(args): |
| 61 | if args.model_dir: |
| 62 | has_model = 0 |
| 63 | pdmodel_name = 0 |
| 64 | pdiparams_name = 0 |
| 65 | for file_name in os.listdir(args.model_dir): |
| 66 | if re.search("__model__", file_name): |
| 67 | has_model = 1 |
| 68 | if file_name.endswith(".pdmodel"): |
| 69 | pdmodel_name = os.path.join(args.model_dir, file_name) |
| 70 | if file_name.endswith(".pdiparams"): |
| 71 | pdiparams_name = os.path.join(args.model_dir, file_name) |
| 72 | if has_model == 1: |
| 73 | config = Config(args.model_dir) |
| 74 | elif pdmodel_name and pdiparams_name: |
| 75 | config = Config(pdmodel_name, pdiparams_name) |
| 76 | else: |
| 77 | raise ValueError( |
| 78 | "config setting error, please check your model path") |
| 79 | else: |
| 80 | config = Config(args.model_file, args.params_file) |
| 81 | |
| 82 | if args.use_gpu: |
| 83 | config.enable_use_gpu(1000, 0) |
| 84 | if args.enable_tensorRT: |
| 85 | config.enable_tensorrt_engine( |
| 86 | max_batch_size=args.batchsize, |
| 87 | min_subgraph_size=9, |
| 88 | precision_mode=paddle.inference.PrecisionType.Float32) |
| 89 | else: |
| 90 | config.disable_gpu() |
| 91 | config.set_cpu_math_library_num_threads(args.cpu_threads) |
| 92 | if args.enable_mkldnn: |
| 93 | config.enable_mkldnn() |
| 94 | predictor = create_predictor(config) |
| 95 | return predictor, config |
| 96 | |
| 97 | |
| 98 | def create_data_loader(args): |