函数 torch_tensorrt::torchscript::compile

函数文档

TORCHTRT_API torch::jit::Module torch_tensorrt::torchscript::compile(const torch::jit::Module &module, CompileSpec info)

使用 TensorRT 为 NVIDIA GPU 编译 TorchScript 模块。

接受现有的 TorchScript 模块和一组设置来配置编译器，并将方法转换为调用等效 TensorRT 引擎的 JIT 图表

具体转换 TorchScript 模块的 forward 方法

参数

• module, : – torch::jit::Module - 现有的 TorchScript 模块

• info, : – torch_tensorrt::CompileSpec - 编译设置

返回值

: 一个针对 TensorRT 引擎的新模块