INNER CODE UNIT · Python

TensorRTBackendRep

onnx/onnx-tensorrt · onnx_tensorrt/backend.py:35

class TensorRTBackendRep(BackendRep):
    def __init__(self, model, device,
            max_workspace_size=None, serialize_engine=False, verbose=False, **kwargs):
        if not isinstance(device, Device):
            device = Device(device)
        self._set_device(device)
        self._logger = TRT_LOGGER
        self.builder = trt.Builder(self._logger)
        self.network = self.builder.create_network(flags=0)
        self.parser = trt.OnnxParser(self.network, self._logger)
        self.config = self.builder.create_builder_config()
        self.serialize_engine = serialize_engine
        self.verbose = verbose
        self.dynamic = False

        if self.verbose:
            print(f'\nRunning {model.graph.name}...')
            TRT_LOGGER.min_severity = trt.Logger.VERBOSE

View source record →