ONNX 加速模型推理

1.安装onnx

pip install onnx onnxruntime

2. 加载模型并转存模型为onnx格式，并测试

使用实例为sentence-transformers 预训练模型计算相似度

class Test(object):
    def init(self):
        # 加载预训练模型
        self.tokenizer = AutoTokenizer.from_pretrained("../all-MiniLM-L6-v2")
        self.model = AutoModel.from_pretrained("../all-MiniLM-L6-v2")
        # 模型设置为推理状态
        self.model.eval()
        # 转化成onnx模型
        self.onnxmodel = self.transfer_onnx(self.model)
        
    def transfer_onnx(self, model):
        #随机创建输入数据。因为模型的导出实际上是执行了一次推理过程。在执行的过程中记录使用到的操作。
        # 输入数据可拷贝至gpu，在cpu上直接写None就可
        inputs = {
            'input_ids':
            torch.randint(32, [2, 32], dtype=torch.long).to(
                None),
            'attention_mask':
            torch.ones([2, 32],
                       dtype=torch.long).to(None),
            'token_type_ids':
            torch.ones([2, 32],
                       dtype=torch.long).to(None), 
        }
        # 转换后模型存储地址
        onnx_model_path = "./temp_turbo_onn