| # from wangyang | |
| from tpu_perf.infer import SGInfer | |
| import numpy as np | |
| import time | |
| #import torch | |
| import os | |
| class EngineOV: | |
| def __init__(self, model_path="", batch=1, device_id=0) : | |
| # 如果环境变量中没有设置device_id，则使用默认值 | |
| if "DEVICE_ID" in os.environ: | |
| device_id = int(os.environ["DEVICE_ID"]) | |
| print(">>>> device_id is in os.environ. and device_id = ",device_id) | |
| self.model = SGInfer(model_path , batch=batch, devices=[device_id]) | |
| def __str__(self): | |
| return "EngineOV: model_path={}, device_id={}".format(self.model_path,self.device_id) | |
| def __call__(self, args): | |
| start = time.time() | |
| if isinstance(args, list): | |
| values = args | |
| elif isinstance(args, dict): | |
| values = list(args.values()) | |
| else: | |
| raise TypeError("args is not list or dict") | |
| # print(values) | |
| #print(time.time() - start) | |
| #start = time.time() | |
| task_id = self.model.put(*values) | |
| #print("put time : ",time.time() - start) | |
| task_id, results, valid = self.model.get() | |
| #print("get time : ",time.time() - start) | |
| return results | |