Affine translation for registration task does not optimize

Nalaxon: <p>Hi,<br>
I am new to the forum and pytorch.<br>
I want to learn a registration transform using pytorch as optimizer.<br>
As first step I prepared a very simple example to shift an image based on other forum posts [1,2,3]<br>
After playing several days I still, fail to understand why the optimizer refuses to increase the shift parameters.</p>
<p>Applying the steps manually I see that proper values lead to an reduction of MSE.<br>
So I fail to understand why different parameters (searched with Ray) and optimizers (SGD, Adam) converge nearly always to the same result: <strong>x :1.0781, y -0.3473</strong></p>
<p>Any hints would be highly apreaciated.</p>
<p>Thanks a lot in advance,<br>
N.</p>
<p>Here is my very simple example</p>
<pre><code class="lang-auto">import torch
from torch import nn
import torch.nn.functional as TF
import torchvision.transforms.functional as TTF
import numpy as np
import matplotlib.pyplot as plt
import imageio
import cv2

class Translation2D(nn.Module):
    def __init__(self, translation_x, translation_y):
        super(Translation2D, self).__init__()
        self.translations = nn.Parameter(torch.stack([translation_x, translation_y], 1))

    def forward(self, input, train=True):
        b, w, h = input.shape

        # compute the transformation matrix
        theta = torch.eye(3).to(input.dtype)
        theta = theta.repeat((b,1,1))
        # if train:
        #     theta[:,0,2] = torch.sigmoid(self.translations[::2]* torch.Tensor([1./h]).repeat((b)).flatten())
        #     theta[:,1,2] = torch.sigmoid(self.translations[1::2] * torch.Tensor([1./w]).repeat((b)).flatten())
        # else:
        theta[:,0,2] = self.translations[...,0]* torch.Tensor([1./w]).repeat((b)).flatten()
        theta[:,1,2] = self.translations[...,1] * torch.Tensor([1./h]).repeat((b)).flatten()

        #apply transformation
        ddf = TF.affine_grid(theta[:,:-1], (b,1,h,w)).to(input.dtype) #2d
        input_permute = input.unsqueeze(0) #b,w,h-&gt; b,c,d,h,w
        input_warped = TF.grid_sample(input_permute, ddf) #needs b,c,d,h,w
        out = input_warped.squeeze(0)        
        return out

def training_loop(model, optimizer, x, y, n=200):
    "Training loop for torch model."
    losses = []
    for i in range(n):
        optimizer.zero_grad()
        pred = model(x, train=True)
        loss = TF.mse_loss(pred, y).sqrt()
        losses.append(loss.item())
        loss.backward()
        print(f'iteration {i}: Loss: {loss.item()}')
        print(f'train Grad: {model.translations.grad}')
        print(f'train Values: {model.translations.data}')
        optimizer.step()
    return losses

ex = imageio.imread(
    'https://nickelquilts.files.wordpress.com/2018/04/1-half-square-triangle-block-with-copyright.jpg') / 255.
sa = cv2.resize(ex, (100, 100)).astype(np.float32)
sa = cv2.cvtColor(sa, cv2.COLOR_BGR2GRAY)
#show(sa)
sa = torch.tensor(sa)

sa_in = TTF.crop(TF.pad(sa, [10, 0, 0, 0], 'constant', 0), 0, 0, 100, 100).unsqueeze(0)
sa = sa.unsqueeze(0)
#show(sa_in.numpy())

b, w, h = sa_in.shape
moving = sa_in.clone()
fixed = sa.clone()
#print(f'moving.size: {moving.shape}, fixed.size: {fixed.shape}')
tx = torch.tensor([0.3], dtype=torch.float32)
ty = torch.tensor([0.3], dtype=torch.float32)
    
# instantiate model
m = Translation2D(translation_x=tx, translation_y=ty)
# Instantiate optimizer

config = {'lr':30, 'momentum':0.9}
opt = torch.optim.SGD(m.parameters(), **config)
losses = training_loop(m, opt, moving, fixed, 50)

with torch.no_grad():
    warped_slice = m(sa_in, train=False)
    diff_slice = sa - warped_slice
    correct = len(torch.where(diff_slice.abs() &lt; .0011)[0])
    loss = TF.mse_loss(warped_slice, sa).sqrt()
    print(f'loss={loss.item()}, accuracy={correct/(h*w)}')
    print(f'translation: {m.translations.detach().data}')
    plt.imsave(f'../sa_in.png', sa_in.squeeze(0).numpy())
    plt.imsave(f'../sa_out.png', warped_slice.squeeze(0).numpy())
    plt.figure()
    plt.plot(losses)
    plt.savefig(f'../loss.png')
</code></pre>
<p>[1] <a href="https://discuss.pytorch.org/t/learning-translation-with-kornia/96314/9" class="inline-onebox">Learning Translation with Kornia - #9 by JuanFMontesinos</a><br>
[2] <a href="https://discuss.pytorch.org/t/trying-to-find-an-angle-use-of-trigonometric-functions-optimizer-not-updating/26780" class="inline-onebox">Trying to find an angle (use of trigonometric functions); optimizer not updating</a><br>
[3] <a href="https://gist.github.com/quantshah/c91543616d2d933e362308c199df184b" class="inline-onebox" rel="noopener nofollow ugc">Automatic differentiation and gradient based optimization with Autograd and Pytorch · GitHub</a></p>

Nalaxon: <p>Hi,<br>
in case someone finds this one I found my error (1) and a workaround (2):<br>
<strong>1. MSE was wrong</strong><br>
I manually tested the loss, but I somehow did it wrong. After doing it again, I found that the loss was not reduced by an correct transformation</p>
<p>It is working, if I swap the input like:</p>
<pre><code class="lang-auto">loss = TF.mse_loss(y, pred).sqrt()
</code></pre>
<p>instead of</p>
<pre><code class="lang-auto">loss = TF.mse_loss(pred, y).sqrt()
</code></pre>
<p><strong>2. Transformation use deformation field instead of affine transform</strong><br>
After accepting that I can’t make it work with <strong>affine transformation</strong> I build the deformation field by my own based on this answer[1]<br>
My class looks now like:</p>
<pre><code class="lang-auto">class Translation2D(nn.Module):
    def __init__(self, translation_x, translation_y):
        super(Translation2D, self).__init__()
        self.translation_x = nn.Parameter(translation_x)
        self.translation_y = nn.Parameter(translation_y)
    def forward(self, input):
        b, c, w, h = input.shape

        x, y = (torch.arange(w) + self.translation_x)/(w-1), (torch.arange(h) + self.translation_y)/(h-1)
        ddf = torch.dstack(torch.meshgrid(x, y, indexing='xy'))*2-1
        out = TF.grid_sample(input, ddf[None]) #needs b,c,d,h,w
        return out
</code></pre>
<p><strong>3. Keep Parameters separated doesn’t affect optimization</strong><br>
Now I keep the shift parameters seperated, but it seems to work also if they are stacked like:</p>
<pre><code class="lang-auto">class Translation2D(nn.Module):
    def __init__(self, translation_x, translation_y):
        super(Translation2D, self).__init__()
        self.translations = nn.Parameter(torch.stack([translation_x, translation_y], 1))

    def forward(self, input, train=True):
        x, y = (torch.arange(w) + self.translations[:,0])/(w-1), (torch.arange(h) + self.translations[:,1])/(h-1)
        ...
</code></pre>
<p>I don’t mark it as solution, because I think I did something wrong in my <strong>affine transformation</strong> approach, and if someone could guide me on that I would still appreciate it. <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
<p>Best regards,<br>
N.</p>
<p>[1] <a href="https://stackoverflow.com/questions/72769563/shifting-an-image-with-bilinear-interpolation-in-pytorch" class="inline-onebox" rel="noopener nofollow ugc">python - Shifting an image with bilinear interpolation in pytorch - Stack Overflow</a></p>