PyTorch Forward Backward Hook

sabiha_afroz: <p>I want to collect activation tensors from forward hook and gradients of activation tensors from backward hook. Here is the model.</p>
<pre><code class="lang-auto">class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10),
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

def hook_pre_forward_function(module, input):
    print(f"Forward Pre hook: Module Name: {module}")
    for tensor in input:
        if isinstance(tensor, torch.Tensor):
            print(f"Forward Pre hook: Module Input Tensor:  {tensor.shape}  id {id(tensor)}")

def hook_forward_function(module, input, output):
    print(f"Forward hook: Module Name: {module}")
    for tensor in input:
        if isinstance(tensor, torch.Tensor):
            print(f"Forward hook: Module Input Tensor:  {tensor.shape} id {id(tensor)}")
    print(f"Forward hook: Module Output Tensor: {output.data.shape} type {type(output.data)} id {id(output)}")

def hook_pre_backward_function(module, grad_output):
    print(f"Backward Pre hook: Module Name: {module}")
    for tensor in grad_output:
        if isinstance(tensor, torch.Tensor):
            print(f"Backward Pre hook: Module Input Tensor:  {tensor.shape} id {id(tensor)}")

def hook_backward_function(module, grad_input, grad_output):
    print(f"Backward hook: Module Name: {module}")
    for tensor in grad_input:
        if isinstance(tensor, torch.Tensor):
            print(f"Backward hook: Module Input Tensor:  {tensor.shape} id {id(tensor)}")
    for tensor in grad_output:
        if isinstance(tensor, torch.Tensor):
            print(f"Backward hook: Module Output Tensor:  {tensor.shape} id {id(tensor)}")

def set_hook(model):
    count = 1
    for name, layer in model._modules.items():
        print(f"Name: {name}, Module: {layer}")
        layer.sid = count
        count = count + 1
        if isinstance(layer, nn.Sequential):
            set_hook(layer)
        else:
            handle1 = layer.register_forward_pre_hook(hook_pre_forward_function)
            handle2 = layer.register_forward_hook(hook_forward_function)
            handle3 = layer.register_full_backward_pre_hook(hook_pre_backward_function)
            handle4 = layer.register_full_backward_hook(hook_backward_function)
            hooks.append(handle1)
            hooks.append(handle2)
            hooks.append(handle3)
            hooks.append(handle4)

</code></pre>
<p>I am printing the id of the tensors to observe if that activation output tensors are used as input tensors in forward pass for next layers. Similarly, the gradients output tensors are whether used in backward pass as inputs for next layers. While I am printing the outputs, I am getting the following output.</p>
<pre><code class="lang-auto">Forward Pre hook: Module Name: Flatten(start_dim=1, end_dim=-1)
Forward Pre hook: Module Input Tensor:  torch.Size([64, 784])  id 139940400429200
Forward hook: Module Name: Flatten(start_dim=1, end_dim=-1)
Forward hook: Module Input Tensor:  torch.Size([64, 784]) id 139940400430256
Forward hook: Module Output Tensor: torch.Size([64, 784]) type &lt;class 'torch.Tensor'&gt; id 139940400430256
Forward Pre hook: Module Name: Linear(in_features=784, out_features=512, bias=True)
Forward Pre hook: Module Input Tensor:  torch.Size([64, 784])  id 139940400430352
Forward hook: Module Name: Linear(in_features=784, out_features=512, bias=True)
Forward hook: Module Input Tensor:  torch.Size([64, 784]) id 139940400430160
Forward hook: Module Output Tensor: torch.Size([64, 512]) type &lt;class 'torch.Tensor'&gt; id 139940400430448
Forward Pre hook: Module Name: ReLU()
Forward Pre hook: Module Input Tensor:  torch.Size([64, 512])  id 139940400430544
Forward hook: Module Name: ReLU()
Forward hook: Module Input Tensor:  torch.Size([64, 512]) id 139940400430640
Forward hook: Module Output Tensor: torch.Size([64, 512]) type &lt;class 'torch.Tensor'&gt; id 139940400430736
Forward Pre hook: Module Name: Linear(in_features=512, out_features=512, bias=True)
Forward Pre hook: Module Input Tensor:  torch.Size([64, 512])  id 139940400430832
Forward hook: Module Name: Linear(in_features=512, out_features=512, bias=True)
Forward hook: Module Input Tensor:  torch.Size([64, 512]) id 139940400430544
Forward hook: Module Output Tensor: torch.Size([64, 512]) type &lt;class 'torch.Tensor'&gt; id 139940400430640
Forward Pre hook: Module Name: ReLU()
Forward Pre hook: Module Input Tensor:  torch.Size([64, 512])  id 139940400430928
Forward hook: Module Name: ReLU()
Forward hook: Module Input Tensor:  torch.Size([64, 512]) id 139940400430832
Forward hook: Module Output Tensor: torch.Size([64, 512]) type &lt;class 'torch.Tensor'&gt; id 139940400431024
Forward Pre hook: Module Name: Linear(in_features=512, out_features=10, bias=True)
Forward Pre hook: Module Input Tensor:  torch.Size([64, 512])  id 139940400693328
Forward hook: Module Name: Linear(in_features=512, out_features=10, bias=True)
Forward hook: Module Input Tensor:  torch.Size([64, 512]) id 139940400420176
Forward hook: Module Output Tensor: torch.Size([64, 10]) type &lt;class 'torch.Tensor'&gt; id 139940400430928
Backward Pre hook: Module Name: Linear(in_features=512, out_features=10, bias=True)
Backward Pre hook: Module Input Tensor:  torch.Size([64, 10]) id 139940400693424
Backward hook: Module Name: Linear(in_features=512, out_features=10, bias=True)
Backward hook: Module Input Tensor:  torch.Size([64, 512]) id 139940400420176
Backward hook: Module Output Tensor:  torch.Size([64, 10]) id 139940400693424
Backward Pre hook: Module Name: ReLU()
Backward Pre hook: Module Input Tensor:  torch.Size([64, 512]) id 139940400420176
Backward hook: Module Name: ReLU()
Backward hook: Module Input Tensor:  torch.Size([64, 512]) id 139940400693424
Backward hook: Module Output Tensor:  torch.Size([64, 512]) id 139940400420176
Backward Pre hook: Module Name: Linear(in_features=512, out_features=512, bias=True)
Backward Pre hook: Module Input Tensor:  torch.Size([64, 512]) id 139940400693424
Backward hook: Module Name: Linear(in_features=512, out_features=512, bias=True)
Backward hook: Module Input Tensor:  torch.Size([64, 512]) id 139940400430544
Backward hook: Module Output Tensor:  torch.Size([64, 512]) id 139940400693424
Backward Pre hook: Module Name: ReLU()
Backward Pre hook: Module Input Tensor:  torch.Size([64, 512]) id 139940400430544
Backward hook: Module Name: ReLU()
Backward hook: Module Input Tensor:  torch.Size([64, 512]) id 139940400420176
Backward hook: Module Output Tensor:  torch.Size([64, 512]) id 139940400430544
Backward Pre hook: Module Name: Linear(in_features=784, out_features=512, bias=True)
Backward Pre hook: Module Input Tensor:  torch.Size([64, 512]) id 139940400420176
Backward hook: Module Name: Linear(in_features=784, out_features=512, bias=True)
Backward hook: Module Input Tensor:  torch.Size([64, 784]) id 139940400430160
Backward hook: Module Output Tensor:  torch.Size([64, 512]) id 139940400420176
Backward Pre hook: Module Name: Flatten(start_dim=1, end_dim=-1)
Backward Pre hook: Module Input Tensor:  torch.Size([64, 784]) id 139940400430160
Backward hook: Module Name: Flatten(start_dim=1, end_dim=-1)
Backward hook: Module Input Tensor:  torch.Size([64, 784]) id 139940400430160
Backward hook: Module Output Tensor:  torch.Size([64, 784]) id 139940400430160
</code></pre>
<p>I expected that flatten layer output tensor will be used as an input to linear layer. But the tensor IDs are different. Same is for backward propagation - next layer output gradient is not used for current layer input gradient (tensor IDs are not same). Why the IDs are different? Is there any other way to get the tensor ID? How to get the activation tensors properly and their gradients?</p>