torch_tensorrt::torchscript::embed_engine_in_new_module 函數

• 定義於 File torch_tensorrt.h

函數文件

TORCHTRT_API torch::jit::Module torch_tensorrt::torchscript::embed_engine_in_new_module(const std::string &engine, Device device, const std::vector<std::string> &input_binding_names = std::vector<std::string>(), const std::vector<std::string> &output_binding_names = std::vector<std::string>())

取得先前建立的 TensorRT 引擎，並將其嵌入到 TorchScript 模組中。

取得預先建置的序列化 TensorRT 引擎，並將其嵌入到 TorchScript 模組中。將引擎的執行註冊為模組的前向方法。前向方法定義為：forward(Tensor[]) -> Tensor[]

如果未指定綁定名稱，則 TensorRT 綁定必須具有以下格式的名稱

• [symbol].[輸入/輸出陣列中的索引] 例如：

• [x.0, x.1, x.2] -> [y.0]

參數

• engine – std::string - 預先建置的序列化 TensorRT 引擎

• device – CompileSepc::Device - 裝置 資訊

• input_binding_names – std::vector<std::string> - TensorRT 綁定的名稱，依原始 PyTorch 函數傳入的順序排列 (預設為假設以下慣例)

• output_binding_names – std::vector<std::string> - TensorRT 綁定的名稱，依原始 PyTorch 函數傳回的順序排列 (預設為假設以下慣例)

傳回

：一個以 TensorRT 引擎為目標的新模組