| 7 | |
| 8 | |
| 9 | class InferenceModel(): |
| 10 | # 初始化函数 |
| 11 | def __init__(self, modelpath, use_gpu=False, gpu_id=0, use_mkldnn=False, cpu_threads=1): |
| 12 | ''' |
| 13 | init the inference model |
| 14 | |
| 15 | modelpath: inference model path |
| 16 | |
| 17 | use_gpu: use gpu or not |
| 18 | |
| 19 | use_mkldnn: use mkldnn or not |
| 20 | ''' |
| 21 | # 加载模型配置 |
| 22 | self.config = self.load_config(modelpath, use_gpu, gpu_id, use_mkldnn, cpu_threads) |
| 23 | |
| 24 | # 打印函数 |
| 25 | def __repr__(self): |
| 26 | ''' |
| 27 | get the numbers and name of inputs and outputs |
| 28 | ''' |
| 29 | return 'input_num: %d\ninput_names: %s\noutput_num: %d\noutput_names: %s' % ( |
| 30 | self.input_num, |
| 31 | str(self.input_names), |
| 32 | self.output_num, |
| 33 | str(self.output_names) |
| 34 | ) |
| 35 | |
| 36 | # 类调用函数 |
| 37 | def __call__(self, *input_datas, batch_size=1): |
| 38 | ''' |
| 39 | call function |
| 40 | ''' |
| 41 | return self.forward(*input_datas, batch_size=batch_size) |
| 42 | |
| 43 | # 模型参数加载函数 |
| 44 | def load_config(self, modelpath, use_gpu, gpu_id, use_mkldnn, cpu_threads): |
| 45 | ''' |
| 46 | load the model config |
| 47 | |
| 48 | modelpath: inference model path |
| 49 | |
| 50 | use_gpu: use gpu or not |
| 51 | |
| 52 | use_mkldnn: use mkldnn or not |
| 53 | ''' |
| 54 | # 对运行位置进行配置 |
| 55 | if use_gpu: |
| 56 | try: |
| 57 | int(os.environ.get('CUDA_VISIBLE_DEVICES')) |
| 58 | except Exception: |
| 59 | print( |
| 60 | '''Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU. Now switch to CPU to continue...''') |
| 61 | use_gpu = False |
| 62 | |
| 63 | if os.path.isdir(modelpath): |
| 64 | if os.path.exists(os.path.join(modelpath, "__params__")): |
| 65 | # __model__ + __params__ |
| 66 | model = os.path.join(modelpath, "__model__") |