How to deploy Grad-CAM using TensorRT

AI & Data Science Deep Learning (Training & Inference) TensorRT
,
1

Description

Hi all, I’m trying to deploy Grad-CAM with TRT C++ API. The deployment is divided into two steps, and I’m having trouble in step2.

  • Step1: forward inference, which is easy to do using TRT.
  • Step2: perform gradient backward, I didn’t find a reference example or API in TRT.

The step2 in pytorch:

  1. use register_backward_hook() to mark the desired gradient position
  2. call loss.backward() function for gradient backward
  3. use hook to get the gradient as a heatmap
    how can I use TRT to do this?
    Thanks in advance.

Part of the code, the full implement is in Relevant Files

class GradCAM:
    def __init__(self, arch: torch.nn.Module, target_layer: torch.nn.Module):
        self.gradients = dict()

        def backward_hook(module, grad_input, grad_output):
            self.gradients['value'] = grad_output[0]

        target_layer.register_backward_hook(backward_hook)

    def forward(self, input, class_idx=None, retain_graph=False):
        logit = self.model_arch(input)
        score = logit[:, class_idx].squeeze()

        score.backward(retain_graph=retain_graph)
        gradients = self.gradients['value']

    def __call__(self, input, class_idx=None, retain_graph=False):
        return self.forward(input, class_idx, retain_graph)

Environment

TensorRT Version: 10.9.0
GPU Type: A100
Nvidia Driver Version: 535.104.12
CUDA Version: 12.2
CUDNN Version: 8.9.7
Operating System + Version: CentOS
Python Version (if applicable): N.A.
TensorFlow Version (if applicable): N.A.
PyTorch Version (if applicable): 2.4.1
Baremetal or Container (if container which image + tag): N.A.

Relevant Files

Similar topic in the forum, but no solution for step2：

Grad-CAM paper:
https://arxiv.org/pdf/1610.02391

Steps To Reproduce

Please include:

  • Exact steps/commands to build your repro
  • Exact steps/commands to run your repro
  • Full traceback of errors encountered