Differentiable affine transforms with grid_sample

harrycb: <p>I’m trying to create a model takes two images of the same size, pushes them through an affine transformation matrix and computes a loss value based on their overlap.</p>
<p>I want the optimiser to change the affine transformations so that they are overlapping.</p>
<p>It doesn’t seem that the gradient is being computed back through to the values in the affine transform.</p>
<p>Here is the module that contains the optimisable parameters</p>
<pre data-code-wrap="py"><code class="lang-py">tt = torch.tensor
class AffineTransform(nn.Module):
    def __init__(self, scalew=1, scaleh=1, transX=0., transY=0.):
        super().__init__()
        
        def makep(x):
            x = tt(x).float()
            return nn.Parameter(x)
        
        self.scalew = makep(scalew)
        self.scaleh = makep(scaleh)
        self.transX = makep(transX) 
        self.transY = makep(transY)
        
    def forward(self, x):
        theta = tt([
            [self.scalew, 0, self.transX],
            [0, self.scaleh, self.transY]
        ])[None]
        grid = F.affine_grid(theta, x.size(), align_corners=False)
        return F.grid_sample(x, grid, align_corners=False)
</code></pre>
<h2><a name="generate-some-random-boxes-1" class="anchor" href="#generate-some-random-boxes-1"></a>Generate some random boxes</h2>
<pre data-code-wrap="py"><code class="lang-py">x = torch.zeros(1, 1,200,200)
target = torch.zeros(1, 1,200,200)
target[0, 0, 90:110, 90:110] = 1
x[0, 0, 75:125,  75:125] = .5
stn = AffineTransform(transX=.2, transY=-.6)
stn2 = AffineTransform()
display(to_pil_image(stn(x)[0]))
display(to_pil_image(stn2(target)[0]))
</code></pre>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/e/9/e93ddd590ee9783d0ca9ecaef8ddd072b86a988f.png" alt="image" data-base62-sha1="xhlNhXxSbL99PKB25wgNrfziZzx" width="200" height="200"></p>
<h2><a name="training-loop-2" class="anchor" href="#training-loop-2"></a>Training loop</h2>
<pre><code class="lang-auto">optim = torch.optim.SGD(stn.parameters(), 1e-3)
display(to_pil_image(stn(x)[0]+stn2(target)[0]))

losses = []
grad_losses = []
overlaps = []

for i in range(100):
    optim.zero_grad()
    box1 = stn(x)
    box2 = stn2(target)
    
    overlap = stn(x).flatten() * stn2(target).flatten()
    
    overlap[overlap!=0] = 1
    overlap = overlap.sum()/maxele
    overlap = 1 - overlap/maxele
    loss = overlap
    
    losses.append(loss.item())
    grad_losses.append(grad_loss.item())
    overlaps.append(overlap)
    loss.backward()
    optim.step()
    
print("After:", stn.transX.item())
display(to_pil_image(stn(x)[0]+target[0]))
plt.plot(losses)
plt.show()
plt.plot(grad_losses)
plt.show()
plt.plot(overlaps)
plt.show()
</code></pre>
<p>The problem:</p>
<pre><code class="lang-auto">RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
</code></pre>

ptrblck: <p>You are recreating the <code>theta</code> tensor in the <code>forward</code>, which will detach the parameters.<br>
You could either create <code>theta</code> in the <code>__init__</code> method as:</p>
<pre><code class="lang-python">        self.theta = nn.Parameter(torch.tensor([
            [self.scalew, 0, self.transX],
            [0, self.scaleh, self.transY]
        ])[None])
</code></pre>
<p>or use <code>torch.cat</code> or <code>torch.stack</code> to create <code>theta</code> in the <code>forward</code> method from the parameters.</p>

harrycb: <p>Thanks for you response, yes that has solved the gradient error but the boxes aren’t being moved around into positions where their overlap is lower.</p>

harrycb: <p>Scratch that, your suggestion completely solved it thank you <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

harrycb: <p>In order to prevent the transformations shearing I attempted to turn off the gradient for the <code>stn.theta[0][1]</code> and <code>stn.theta[1][0]</code> parts of the affine tranformation but this wasn’t possible due to them not being leaf nodes.</p>
<p>Instead I’ve created a tensor to prevent the shear parts of the affine transform from having any effect. Is this the correct approach?</p>
<pre><code class="lang-auto">T = torch.tensor
class AffineTransform(nn.Module):
    def __init__(self, scalew=1, scaleh=1, transX=0., transY=0.):
        super().__init__()
        self.theta = nn.Parameter(T([
            [scalew, 0, transX],
            [0, scaleh, transY]
        ], dtype=torch.float)[None])
        self._stop_shear = T([[1,0,1],[0,1,1]], requires_grad=False)
        
        
    def forward(self, x):
        theta = self.theta * self._stop_shear
        grid = F.affine_grid(theta, x.size(), align_corners=True)
        return F.grid_sample(x, grid, align_corners=True)
</code></pre>

ptrblck: <p>Zeroing out the values in each iteration will reset these parameters and would thus probably work.<br>
However, your first approach of zeroing out the gradients should also work, I think.<br>
Did <code>self.theta.grad[0, 1] = 0.</code> give you any errors?<br>
You could also use <code>self.theta.register_hook</code> and zero out the gradients using the mask.</p>

harrycb: <p>Setting the grads directly seems to work yes thank you</p>

Elior_Malul: <aside class="quote no-group" data-username="harrycb" data-post="7" data-topic="79305">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/harrycb/48/23610_2.png" class="avatar"> harrycb:</div>
<blockquote>
<p>to work yes thank you</p>
</blockquote>
</aside>
<p>Isn’t possible to just maintain a <a href="https://pytorch.org/docs/stable/generated/torch.nn.Linear.html" rel="noopener nofollow ugc">Linear Layer</a> of 2X2 for 2D affine transform,<br>
or 3X3 for a 3D affine transform?</p>