Main forward method: Builds TRT engine if not available yet. Tries to run TRT engine If exception thrown and self.callback==True: falls back to original Pytorch Args: Passing through whatever args wrapped module's forward() has Returns: Passing th
(self, model, argv, kwargs)
| 426 | self.logger.info(f"Exception while loading the engine:\n{e}") |
| 427 | |
| 428 | def forward(self, model, argv, kwargs): |
| 429 | """ |
| 430 | Main forward method: |
| 431 | Builds TRT engine if not available yet. |
| 432 | Tries to run TRT engine |
| 433 | If exception thrown and self.callback==True: falls back to original Pytorch |
| 434 | |
| 435 | Args: Passing through whatever args wrapped module's forward() has |
| 436 | Returns: Passing through wrapped module's forward() return value(s) |
| 437 | |
| 438 | """ |
| 439 | args = self.defaults |
| 440 | args.update(kwargs) |
| 441 | if len(argv) > 0: |
| 442 | args.update(self._inputs_to_dict(argv)) |
| 443 | |
| 444 | if self.engine is None and not self.disabled: |
| 445 | # Restore original forward for export |
| 446 | new_forward = model.forward |
| 447 | model.forward = self.old_forward |
| 448 | try: |
| 449 | self._load_engine() |
| 450 | if self.engine is None: |
| 451 | build_args = args.copy() |
| 452 | with torch.no_grad(): |
| 453 | self._build_and_save(model, build_args) |
| 454 | # This will reassign input_names from the engine |
| 455 | self._load_engine() |
| 456 | assert self.engine is not None |
| 457 | except Exception as e: |
| 458 | if self.fallback: |
| 459 | self.logger.info(f"Failed to build engine: {e}") |
| 460 | self.disabled = True |
| 461 | else: |
| 462 | raise e |
| 463 | if not self.disabled and not self.fallback: |
| 464 | # Delete all parameters |
| 465 | for param in model.parameters(): |
| 466 | del param |
| 467 | # Call empty_cache to release GPU memory |
| 468 | torch.cuda.empty_cache() |
| 469 | # restore TRT hook |
| 470 | model.forward = new_forward |
| 471 | # Run the engine |
| 472 | try: |
| 473 | if self.engine is not None: |
| 474 | # forward_trt is not thread safe as we do not use per-thread execution contexts |
| 475 | with lock_sm: |
| 476 | device = torch.cuda.current_device() |
| 477 | stream = torch.cuda.Stream(device=device) |
| 478 | self.engine.set_inputs(unroll_input(self.input_names, args), stream.cuda_stream) |
| 479 | self.engine.allocate_buffers(device=device) |
| 480 | # Need this to synchronize with Torch stream |
| 481 | stream.wait_stream(torch.cuda.current_stream()) |
| 482 | ret = self.engine.infer(stream.cuda_stream, use_cuda_graph=self.use_cuda_graph) |
| 483 | # if output_names is not None, return dictionary |
| 484 | if not self.return_dict: |
| 485 | ret = list(ret.values()) |
no test coverage detected