Customized linear layer gives issue during the backward computation

Torcione: <p>Hi everyone,<br>
I need to implement a customized version of a normalization layer.<br>
The idea is the following:</p>
<ul>
<li>after initialization we forward the dataset (or a single batch) through the network; similarly to batch norm, every time the module is called inside the model we compute the mean over the batch (<strong>first foundamental difference w.r.t. BatchNorm</strong>: we do not need the std in this case) and we save a copy of these mean tensors</li>
<li>we use, from that moment on, the  computed mean tensors to shift the forward propagation tensor, i.e. we subtract to the forwarding tensors the mean tensors computed after initialization, in the previous step(<strong>second foundamental difference w.r.t. BatchNorm</strong>: we use always the same constant, computed with the same forward, to offset the tensor instead of compute it with the statistics of the batch selected at each time)</li>
</ul>
<pre><code class="lang-auto">class CustomInitBatchNorm2d(nn.Module):
    def __init__(self, layer_id):
        super().__init__()
        self.layer_id = layer_id
        self.stored_means = {}  # Dictionary to store means for each layer
    
    def forward(self, x):
        if self.training and self.layer_id not in self.stored_means:
            self.stored_means[self.layer_id] = x.mean(dim=0, keepdim=True)
        
        if self.layer_id in self.stored_means:
            return x - self.stored_means[self.layer_id]
        else:
            return x
</code></pre>
<p>Using the above module inside my model I get an error during the backward pass; I don’t understand why as the function I am defining in my forward pass is a simple linear operation (just subtracting a constant tensor).</p>
<p>I tried to modify  the class adding a backward method as follows:</p>
<pre><code class="lang-auto">class CustomInitBatchNorm2d(nn.Module):
    def __init__(self, layer_id):
        super().__init__()
        self.layer_id = layer_id
        self.stored_means = {}  # Dictionary to store means for each layer
    
    def forward(self, x):
        if self.training and self.layer_id not in self.stored_means:
            self.stored_means[self.layer_id] = x.mean(dim=0, keepdim=True)
        
        if self.layer_id in self.stored_means:
            return x - self.stored_means[self.layer_id]
        else:
            return x

    def backward(self, grad_output):
        return grad_output  # Propagate gradients unchanged
</code></pre>
<p>but is not working as well.<br>
More precisely, the error I get is:</p>
<pre><code class="lang-auto">  File ".../.local/lib/python3.9/site-packages/torch/_tensor.py", line 487, in backward
    torch.autograd.backward(
  File ".../.local/lib/python3.9/site-packages/torch/autograd/__init__.py", line 197, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: Trying to backward through the graph a second time (or directly access saved tensors after they have already been freed). Saved intermediate values of the graph are freed when you call .backward() or autograd.grad(). Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved tensors after calling backward.
</code></pre>
<p>I am not an expert of pytorch, do you have any idea of what is going on and how to fix it?</p>

ptrblck: <p>Could you explain what exactly is failing and post the error message with the stacktrace here, please?</p>

Torcione: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a> , sure!<br>
It seems that during the <code>backward()</code> step something goes wrong during the gradient computation; I reported the error I get at the end of the question.</p>

ptrblck: <p>Thanks for the update!<br>
I guess this line:</p>
<pre><code class="lang-python">self.stored_means[self.layer_id] = x.mean(dim=0, keepdim=True)
</code></pre>
<p>is storing the gradient history in <code>self.stored_mean</code> so you might need to <code>.detach()</code> the <code>x.mean()</code> tensor before assigning it.</p>

Torcione: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> thank you for your prompt reply! it was indeed there the problem</p>