Back propagate masked output

pietz: <p>I’m fiddling around with a loss function that measures similarity between feature maps in pixel space. The target is a sparse map of 0s or 1s of size 224x224. The key aspect of what I’m trying is that I only want the value of local maxima of the prediction to be back propagated, meaning that I need to mask the predicted map.</p>
<p>Within the loss I do this:</p>
<pre><code class="lang-auto">def forward(self, inputs, targets):
    pooled = self.max_pool(inputs)
    mask = (inputs == pooled).float()
    maxima = inputs * mask
    return self.bce_loss(maxima, targets)
</code></pre>
<p>By using max pooling and comparing it to the initial mask, I create a mask that only highlights local maxima. Multiplying this to the initial values gets my wanted logits back.</p>
<p>What I want is to back propagate only those information that were part of the mask (have a 1.0) and not those values that are set to 0.0.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/0/e08900fd85a8f969f8382cfee6f8a835a1510148.png" data-download-href="https://discuss.pytorch.org/uploads/default/e08900fd85a8f969f8382cfee6f8a835a1510148" title="Screenshot 2022-07-07 at 09.28.44"><img src="https://discuss.pytorch.org/uploads/default/original/3X/e/0/e08900fd85a8f969f8382cfee6f8a835a1510148.png" alt="Screenshot 2022-07-07 at 09.28.44" data-base62-sha1="w2ktVWklgVwifWi21jyPJBd5KnS" width="506" height="500" data-dominant-color="030303"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot 2022-07-07 at 09.28.44</span><span class="informations">636×628 4.56 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>My direct model outputs show a very strong checkerboard pattern, which is why I suspect also the masked out values are somehow being back propagated. I expected the output to look a lot more smooth. E.g. I expected the values around local maxima to be closer to the maxima instead of being zero (on the grid lines).</p>
<p>I do use transposed convs with k=2 for upsampling.</p>
<p>Where am I going wrong? Is there a fundamental problem in my approach?</p>

pietz: <p>I tried this</p>
<pre><code class="lang-auto">def forward(self, inputs, targets):
    inp_copy = inputs.clone().detach()
    pooled = self.mp(inp_copy)
    mask = (inp_copy == pooled).float()
    inputs.register_hook(lambda grad: grad * mask)
    return self.loss_fn(inputs, targets)
</code></pre>
<p>but now I’m getting a <code>cannot register a hook on a tensor that doesn't require gradient</code> exception.</p>

pietz: <p>This fixed the hook problem:</p>
<pre><code class="lang-auto">def forward(self, inputs, targets):
    pooled = self.mp(inputs)
    mask = (inputs == pooled).float()
    if self.train and inputs.requires_grad:
        inputs.register_hook(lambda grad: grad * mask)
    return self.loss_fn(inputs, targets)
</code></pre>
<p>The problem came from pytorch lightning running a couple of eval batches before training.</p>
<p>However, the original problem with the checkerboard was fixed by switching from transposed convs to interpolate + regular conv. I thought a non overlapping transpose conv was enough to get rid of it, but that was not true.</p>