Function torch_tensorrt::torchscript::embed_engine_in_new_module

Function Documentation

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>())

Take a previously created TensorRT engine and embed it in in a TorchScript module.

Takes a pre-built serialized TensorRT engine and embeds it in a TorchScript module. Registers execution of the engine as the forward method of the module Forward is defined as: forward(Tensor[]) -> Tensor[]

If binding names not specified TensorRT bindings must have names with the following format:

• [symbol].[index in input / output array] ex.

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

Parameters

• engine, : – std::string - Pre-built serialized TensorRT engine

• device, : – CompileSepc::Device - Device information

• input_binding_names, : – std::vector<std::string> - Name of TensorRT bindings in order passed in by original PyTorch function (defaults to assuming convention below)

• output_binding_names, : – std::vector<std::string> - Name of TensorRT bindings in order returned by original PyTorch function (defaults to assuming convention below)

Returns

: A new module targeting a TensorRT engine