def __init__(self,
module: torch.fx.GraphModule,
input_shapes: List[InputTensorSpec],
logger_level=trt.Logger.WARNING):
# Preprocess the model
module = copy.deepcopy(module)
module = module.cpu().float()
module = NormalizeArgs(module).transform()
super().__init__(module)
self.logger = trt.Logger(logger_level)
self.builder = trt.Builder(self.logger)
# TODO: explicit batching
# EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
# self.network = self.builder.create_network(EXPLICIT_BATCH)
self.network = self.builder.create_network()
self.input_shape_itr = iter(input_shapes)
self._cur_node_name: Optional[str] = None
self._input_names: List[str] = []
self._output_names: List[str] = []
def __init__(self,
module: torch.fx.GraphModule,
input_shapes: List[InputTensorSpec],
logger_level=trt.Logger.WARNING):
# Preprocess the model
module = copy.copy(module)
module = module.cpu()
module = NormalizeArgs(module).transform()
super().__init__(module)
self.logger = trt.Logger(logger_level)
self.builder = trt.Builder(self.logger)
self.network = self.builder.create_network()
self.input_shape_itr = iter(input_shapes)
self._cur_node_name: Optional[str] = None
self._input_names: List[str] = []
self._output_names: List[str] = []