Learning Translation with Kornia

username: Niv_Nayman
cooked: <p>Hi,<br>
I have been trying to learn translation (x, y) parameters with Kornia in the following manner:</p>
<pre><code class="lang-auto">class DTranslation(nn.Module):
    def __init__(self, x_translation, y_translation):
        super(DTranslation, self).__init__()
        self.translations = torch.stack([x_translation, y_translation], 1)
        self.angle = torch.tensor([0])
        
        
    def forward(self, input):
        _, _, h, w = input.shape
        if self.angle.shape[0] != input.shape[0]:
            angle = self.angle.repeat(input.shape[0])
        else:
            angle = self.angle
            
        if self.translations.shape[0] != input.shape[0]:
            translations = self.translations.repeat([input.shape[0], 1])
        else:
            translations = self.translations
        
        translations[:, 0] *= h
        translations[:, 1] *= w
        
        # define the rotation center
        center = torch.ones(2)
        center[..., 0] = input.shape[3] / 2  # x
        center[..., 1] = input.shape[2] / 2  # y
        center = center.repeat(input.shape[0], 1)

        # define the scale factor
        scale = torch.ones(input.shape[0])

        # compute the transformation matrix
        M = kornia.get_rotation_matrix2d(center, -angle, scale)
        
        # Translate
        M[..., 2] += translations  # tx/ty
        
        # apply the transformation to original image
        out = kornia.warp_affine(input, M, dsize=(h, w))
        
        return out

tx = torch.tensor([0.3], dtype=torch.float32)
tx_p = Parameter(tx, requires_grad=True)
            
ty = torch.tensor([-0.2], dtype=torch.float32)
ty_p = Parameter(ty, requires_grad=True)

translation = DTranslation(x_translation=tx_p, y_translation=ty_p)

criterion = nn.MSELoss()
optimizer = optim.Adam([tx_p, ty_p], lr=1)

for x, y in dataloader: 
        optimizer.zero_grad()
        loss = criterion(x, translation(x))
        loss.backward()
        optimizer.step()
</code></pre>
<p>The first backward call passes, but the second one fails:</p>
<pre><code class="lang-auto">RuntimeError: Trying to backward through the graph a second time, but the saved intermediate results have already been freed. Specify retain_graph=True when calling backward the first time.
</code></pre>
<p>When I follow the instructions and instead I use:</p>
<pre><code class="lang-auto"> loss.backward(retain_graph=True)
</code></pre>
<p>I receive the following error:</p>
<pre><code class="lang-auto">RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.FloatTensor [1]] is at version 1; expected version 0 instead. Hint: enable anomaly detection to find the operation that failed to compute its gradient, with torch.autograd.set_detect_anomaly(True).
</code></pre>
<p>I also tried to avoid the inplace operation:</p>
<pre><code class="lang-auto">        M[..., 2] += translations  # tx/ty
</code></pre>
<p>and use instead:</p>
<pre><code class="lang-auto">shape = list(M.shape)
shape[-1] -= 1
M = M + torch.cat([torch.zeros(shape), translations.unsqueeze(-1)], 2)
</code></pre>
<p>but got the same errors.</p>
<p>Maybe another inplace operation causes this ? e.g. <a href="https://github.com/kornia/kornia/blob/5a736409a9a133da27c3dfa581bba2bc71f27286/kornia/geometry/conversions.py#L122" rel="nofollow noopener">https://github.com/kornia/kornia/blob/5a736409a9a133da27c3dfa581bba2bc71f27286/kornia/geometry/conversions.py#L122</a></p>
<p>Or is it something else ?</p>
<p>It is worth mentioning that I do manage to backpropagated through rotation angle and shear (x, y) but only translation seems to be the problem, i.e. when I comment out</p>
<pre><code class="lang-auto">        M[..., 2] += translations  # tx/ty
</code></pre>
<p>No errors occur, but of course I cannot learn the translation parameters either.</p>
<p>Any thoughts ?</p>

username: JuanFMontesinos
cooked: <p>My guess is that defining the constant outside nn.Module breaks the graph. Once you define them inside nn.Module they are gonna be registered as buffers and it sounds bad to me.</p>
<p>Why don’t you just define them as regular nn.Parameters and pass the nn.Module to the optimizer?<br>
Anyway, can you provide a runnable script?</p>
<p>WRT the inplace ops, indeed that M[…,2] is not allowed.</p>

username: Niv_Nayman
cooked: <aside class="quote no-group" data-username="JuanFMontesinos" data-post="2" data-topic="96314" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/juanfmontesinos/48/76115_2.png" class="avatar"> JuanFMontesinos:</div>
<blockquote>
<p>My guess is that defining the constant outside nn.Module breaks the graph. Once you define them inside nn.Module they are gonna be registered as buffers and it sounds bad to me.</p>
<p>Why don’t you just define them as regular nn.Parameters and pass the nn.Module to the optimizer?<br>
Anyway, can you provide a runnable script?</p>
<p>WRT the inplace ops, indeed that M[…,2] is not allowed.</p>
</blockquote>
</aside>
<p>Thank you.</p>
<p>Please note that I manage to learn rotation angle while the parameters are passed to the nn.Module so this doesn’t seem to be the problem.</p>
<p>Attached is a runbable code (originally a notebook):</p>
<pre><code class="lang-auto">import kornia; print(kornia.__version__)


# In[2]:


import torch
import torchvision
import torchvision.transforms as transforms

transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

bs= 4#096
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=bs, shuffle=True, num_workers=2)

testset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=bs, shuffle=False, num_workers=2)

classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')


# In[3]:


from torch import nn
class DRotation(nn.Module):
    def __init__(self, angle):
        super(DRotation, self).__init__()
        self.angle = angle
        
        
    def forward(self, input):
        if self.angle.shape[0] != input.shape[0]:
            angle = self.angle.repeat(input.shape[0])
        else:
            angle = self.angle
        
        # define the rotation center
        center = torch.ones(2)
        center[..., 0] = input.shape[3] / 2  # x
        center[..., 1] = input.shape[2] / 2  # y
        center = center.repeat(input.shape[0], 1)

        # define the scale factor
        scale = torch.ones(input.shape[0])

        # compute the transformation matrix
        M = kornia.get_rotation_matrix2d(center, -angle, scale)

        # apply the transformation to original image
        _, _, h, w = input.shape
        out = kornia.warp_affine(input, M, dsize=(h, w))
        
        return out


# In[4]:


class DTranslation(nn.Module):
    def __init__(self, x_translation, y_translation):
        super(DTranslation, self).__init__()
        self.translations = torch.stack([x_translation, y_translation], 1)
        self.angle = torch.tensor([0])
        
        
    def forward(self, input):
        _, _, h, w = input.shape
        if self.angle.shape[0] != input.shape[0]:
            angle = self.angle.repeat(input.shape[0])
        else:
            angle = self.angle
            
        if self.translations.shape[0] != input.shape[0]:
            translations = self.translations.repeat([input.shape[0], 1])
        else:
            translations = self.translations
        
        translations[:, 0] *= h
        translations[:, 1] *= w
        
        # define the rotation center
        center = torch.ones(2)
        center[..., 0] = input.shape[3] / 2  # x
        center[..., 1] = input.shape[2] / 2  # y
        center = center.repeat(input.shape[0], 1)

        # define the scale factor
        scale = torch.ones(input.shape[0])

        # compute the transformation matrix
        M = kornia.get_rotation_matrix2d(center, -angle, scale)
        

        # Translate
        M[..., 2] += translations  # tx/ty
#         shape = list(M.shape)
#         shape[-1] -= 1
#         M = M + torch.cat([torch.zeros(shape), translations.unsqueeze(-1)], 2)

        
        # apply the transformation to original image
        out = kornia.warp_affine(input, M, dsize=(h, w))
        
        return out


# In[5]:


from torch.nn.parameter import Parameter

class LearnableAug(nn.Module):
    def __init__(self, aug_names=[]):
        super(LearnableAug, self).__init__()
        
        self.fixed_transforms = []
        self.learnable_transforms = []
        self.have_entropy = []
        self._param_groups = []
        self.lrs = []
        self.bounds= {}
        self.eps = 0 #1e-6
        
        if 'translation' in aug_names:
            var = torch.tensor([0.3], dtype=torch.float32, requires_grad=False)
            self.register_buffer('_x_translation', var)
            var = torch.tensor([-0.2], dtype=torch.float32, requires_grad=False)
            self.register_buffer('_y_translation', var)
            
            self.fixed_transforms.append(DTranslation(
                x_translation=self._buffers['_x_translation'],
                y_translation=self._buffers['_y_translation'],
            ))

            var = torch.tensor([0], dtype=torch.float32)
            param = Parameter(var, requires_grad=True)
            self.register_parameter(name='x_translation', param=param)
            
            var = torch.tensor([0], dtype=torch.float32)
            param = Parameter(var, requires_grad=True)
            self.register_parameter(name='y_translation', param=param)
            
            self.learnable_transforms.append(DTranslation(
                x_translation=self._parameters['x_translation'],
                y_translation=self._parameters['y_translation'],
            ))
            
            self._param_groups.append(['x_translation'])
            self.lrs.append(1)
            self.bounds['x_translation'] = (-0.5, 0.5)
            
            self._param_groups.append(['y_translation'])
            self.lrs.append(1)
            self.bounds['y_translation'] = (-0.5, 0.5)
        
            
        if 'rotation' in aug_names and 'fixed_affine' not in aug_names:
            # Fixed Rotation
            var = torch.tensor([15], dtype=torch.float32, requires_grad=False)
            self.register_buffer('_angle', var)
            self.fixed_transforms.append(DRotation(var))

            var = torch.tensor([0], dtype=torch.float32)
            param = Parameter(var, requires_grad=True)
            self.register_parameter(name='angle', param=param)
            self.learnable_transforms.append(DRotation(param))
            self._param_groups.append(['angle'])
            self.lrs.append(1)
            self.bounds['angle'] = (-60, 60)
       
            
        self.fixed_transforms = torch.nn.Sequential(*self.fixed_transforms)
        self.learnable_transforms = torch.nn.Sequential(*self.learnable_transforms)

    def apply_constraints(self):
        for param_name, bounds in self.bounds.items():
            if self._parameters[param_name] &lt; bounds[0]:
                if isinstance(bounds[0], torch.Tensor):
                    self._parameters[param_name].data = bounds[0].data
                else:
                    bound_tensor = torch.ones_like(self._parameters[param_name]) * bounds[0] 
                    self._parameters[param_name].data = bound_tensor.data
                    
            if self._parameters[param_name] &gt; bounds[1]:
                if isinstance(bounds[1], torch.Tensor):
                    self._parameters[param_name].data = bounds[1].data
                else:
                    bound_tensor = torch.ones_like(self._parameters[param_name]) * bounds[1] 
                    self._parameters[param_name].data = bound_tensor.data
                    
                    
    def forward(self, x, fixed=False):
        if not fixed:
            self.apply_constraints()
            out = self.learnable_transforms(x)
        else:
            out = self.fixed_transforms(x)
        return out
    
    @property
    def param_groups(self):
        return self._param_groups
    
    def get_params_for_optimizer(self):
        param_groups = []
        for param_group, lr in zip(self._param_groups, self.lrs):
            param_groups.append(dict(params=[self._parameters[name] for name in param_group], lr=lr))
        
        return param_groups
    


# In[8]:


from IPython.display import clear_output

from time import sleep
from torch import optim
device = torch.device('cuda')

dataloader = trainloader

augs = LearnableAug(['rotation'])
# augs = LearnableAug(['translation'])
criterion = nn.MSELoss()
optimizer = optim.Adam(augs.get_params_for_optimizer())
scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=10)#, T_mult)

iters = len(dataloader)
epochs = 1

for epoch in range(epochs):
    for i, data in enumerate(dataloader):
        images, labels = data

        # perform the transforms
        ft_images = augs(images, fixed=True)
        dt_images = augs(images)

        optimizer.zero_grad()
        loss = criterion(ft_images, dt_images)
        print(loss)
        loss.backward()
            
        optimizer.step()
        scheduler.step((epoch * iters + i) / (epochs * iters))
</code></pre>

username: JuanFMontesinos
cooked: <p>Soo I did a minimal example:</p>
<pre><code class="lang-auto">from torch import nn
import torch
import kornia
import imageio

import cv2
import matplotlib.pyplot as plt
import numpy as np

ex = imageio.imread(
    'https://nickelquilts.files.wordpress.com/2018/04/1-half-square-triangle-block-with-copyright.jpg') / 255.

ex = cv2.resize(ex, (100, 100)).astype(np.float32)

torch.manual_seed(666)


def show(img):
    if isinstance(img, np.ndarray):
        pass
    else:
        img = img[0].permute(1, 2, 0).detach().cpu().numpy()
    plt.imshow(img)
    plt.show()


show(ex)


class DTranslation(nn.Module):
    def __init__(self, x_translation, y_translation):
        super(DTranslation, self).__init__()
        self.translations = nn.Parameter(torch.stack([x_translation, y_translation], 1))
        self.angle = torch.tensor([0])

    def forward(self, input, train=True):
        _, _, h, w = input.shape
        if self.angle.shape[0] != input.shape[0]:
            angle = self.angle.repeat(input.shape[0])
        else:
            angle = self.angle

        if self.translations.shape[0] != input.shape[0]:
            translations = self.translations.repeat([input.shape[0], 1])
        else:
            translations = self.translations
        if train:
            translations = torch.sigmoid(translations) * torch.Tensor([h, w])
        else:
            translations = translations * torch.Tensor([h, w])
        # define the rotation center
        center = torch.ones(2)
        center[..., 0] = input.shape[3] / 2  # x
        center[..., 1] = input.shape[2] / 2  # y
        center = center.repeat(input.shape[0], 1)

        # define the scale factor
        scale = torch.ones(input.shape[0])

        # compute the transformation matrix
        M = kornia.get_rotation_matrix2d(center, -angle, scale)

        # Translate
        shape = list(M.shape)
        shape[-1] -= 1
        M = M + torch.cat([torch.zeros(shape), translations.unsqueeze(-1)], 2)

        # apply the transformation to original image
        out = kornia.warp_affine(input, M, dsize=(h, w), padding_mode='zeros')

        return out


tx = torch.tensor([-1], dtype=torch.float32)

ty = torch.tensor([-1], dtype=torch.float32)

translation = DTranslation(x_translation=tx, y_translation=ty)


class Corr(nn.Module):
    def forward(self, x, pred):
        x = x.flatten()
        p = pred.flatten()
        x_n = torch.norm(x)
        p_n = torch.norm(p)

        return torch.dot(x, p) / (x_n * p_n)


criterion = nn.L1Loss()
# criterion = nn.CosineSimilarity(1)
optimizer = torch.optim.Adam(translation.parameters(), lr=0.0001)
img = torch.from_numpy(ex).permute(2, 0, 1)[None, ...]
img.requires_grad_(True)
inst = DTranslation(torch.tensor([0.2], dtype=torch.float32),
                    torch.tensor([0.2], dtype=torch.float32))

with torch.no_grad():
    gt = inst(img)
show(gt)
for i in range(1000):
    optimizer.zero_grad()
    result = translation(img, train=True)
    if i % 10 == 0:
        show(result)
    loss = criterion(gt.view(1, -1), result.view(1, -1))

    loss.backward()
    print(f'Grad: '
          f'{translation.translations.grad}, '
          f'Loss: {loss.item()}, '
          f'Value: {torch.sigmoid(translation.translations.data)}')
    optimizer.step()
</code></pre>
<p>But i have to say i don’t figure out why it doesn’t optimize it.<br>
Gradients are not None, thus it properly prop.<br>
The rot matrix is ok<br>
I’ve tried to maximize corr instead of minimize euclidian but yet the same</p>
<p>I’ll have another look later</p>

username: Niv_Nayman
cooked: <p>Thank you. I wish I knew the answer. To be honest I asked the same question in Kornia’s Github repo and shared the code there:<br>
</p><aside class="onebox githubissue">
  <header class="source">
      <a href="https://github.com/kornia/kornia/issues/682" target="_blank" rel="nofollow noopener">github.com/kornia/kornia</a>
  </header>
  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewBox="0 0 14 16" aria-hidden="true"><path d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/kornia/kornia/issues/682" target="_blank" rel="nofollow noopener">Backpropagation through Translation</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2020-09-15" data-time="07:19:38" data-timezone="UTC">07:19AM - 15 Sep 20 UTC</span>
      </div>


      <div class="user">
        <a href="https://github.com/NivNayman" target="_blank" rel="nofollow noopener">
          <img alt="NivNayman" src="https://avatars0.githubusercontent.com/u/29902065?v=4" class="onebox-avatar-inline" width="20" height="20">
          NivNayman
        </a>
      </div>
    </div>
  </div>
</div>

<div class="github-row">
  <p class="github-content">Hi,
Thank you for the wonderful package !
I have been trying to learn translation (x, y) parameters in the following manner:
`class DTranslation(nn.Module):
def...</p>
</div>

<div class="labels">
</div>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>I hope they find out the reason and correct this.</p>

username: JuanFMontesinos
cooked: <p>I think it’s more theoretical than any other thing.<br>
If you run the example i wrote you will see it kinda get stuck in a local minima.</p>

username: Niv_Nayman
cooked: <p>Thank you for your code. I ran it. Indeed it doesn’t raise that error, yet it doesn’t train.<br>
Actually, Instead of taking all the changes you suggested, I took my code and changed only the line:</p>
<pre><code class="lang-auto">self.translations = torch.stack([x_translation, y_translation], 1)
</code></pre>
<p>To:</p>
<pre><code class="lang-auto">self.translations = nn.Parameter(torch.stack([x_translation, y_translation], 1))
</code></pre>
<p>and indeed the error went away. Yet the gradients w.r.t the original translation tensors remain zero.<br>
I wonder if the constructor nn.Parameter() somehow duplicates the tensor such that the original one is detached from the graph hence no error occur yet neither gradients are back-propagated properly.</p>
<p>I mean, it should be working also without this change, as done in the rotation example in my code.</p>

username: JuanFMontesinos
cooked: <p>Well,<br>
I don’t remember the details of Parameter but my guess is it detach.<br>
It’s not pythonic-&gt;“pytorchnic” to define a parameter outside me module. If that is the case, then you don’t really need to wrap everything in a nn.Module. You can just use a standard function.</p>
<p>When you assign a torch tensor to a nn.Module, it internally calls register buffer. If torch is promoting the stack of to parameters into a tensor, the resulting one is being assigned  as a buffer.</p>
<p>Anyway the input is gonna be a leaf node. If you run a gradient anomaly detector it will rise an error since gradients doesn’t flow to the input due to inplace ops.<br>
Howver you are only using inplace ops to build the rotation matrix. The ops related to the translation vector are not affected and that’s why gradients flow down there.</p>
<p>Besides, the rotation matrix is ok. Ones in the diag + the translation vector .</p>
<p>My point is that:<br>
GT:<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/f/9/f940248eed9066f16db7e907574cdb544e31a9a4.png" alt="image" data-base62-sha1="zyYjyFfT5XJxdwX391aUdcdaHfm" width="640" height="480"><br>
1st pref<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/6/4/64c7133ffec5d27180b79b3331a5b93744f09b65.png" alt="image" data-base62-sha1="enwglBbWvgGAWa8P9VuCcZ4N4C9" width="640" height="480"><br>
Dpending of the learning rate it jumps around that point<br>
or diverge</p>

username: JuanFMontesinos
cooked: <p>With a bit more of analysis.</p>
<pre><code class="lang-auto">from torch import nn
import torch
import kornia
import imageio

import cv2
import matplotlib.pyplot as plt
import numpy as np

ex = imageio.imread(
    'https://nickelquilts.files.wordpress.com/2018/04/1-half-square-triangle-block-with-copyright.jpg') / 255.

ex = cv2.resize(ex, (100, 100)).astype(np.float32)

torch.manual_seed(666)


def show(img):
    if isinstance(img, np.ndarray):
        pass
    else:
        img = img[0].permute(1, 2, 0).detach().cpu().numpy()
    plt.imshow(img)
    plt.show()


show(ex)


class DTranslation(nn.Module):
    def __init__(self, x_translation, y_translation):
        super(DTranslation, self).__init__()
        self.translations = nn.Parameter(torch.stack([x_translation, y_translation], 1))
        self.angle = torch.tensor([0])

    def forward(self, input, train=True):
        _, _, h, w = input.shape
        if self.angle.shape[0] != input.shape[0]:
            angle = self.angle.repeat(input.shape[0])
        else:
            angle = self.angle

        if self.translations.shape[0] != input.shape[0]:
            translations = self.translations.repeat([input.shape[0], 1])
        else:
            translations = self.translations
        if train:
            translations = torch.sigmoid(translations) * torch.Tensor([h, w])
        else:
            translations = translations * torch.Tensor([h, w])
        # define the rotation center
        center = torch.ones(2)
        center[..., 0] = input.shape[3] / 2  # x
        center[..., 1] = input.shape[2] / 2  # y
        center = center.repeat(input.shape[0], 1)

        # define the scale factor
        scale = torch.ones(input.shape[0])

        # compute the transformation matrix
        M = kornia.get_rotation_matrix2d(center, -angle, scale)

        # Translate
        shape = list(M.shape)
        shape[-1] -= 1
        M = M + torch.cat([torch.zeros(shape), translations.unsqueeze(-1)], 2)
        # apply the transformation to original image
        out = kornia.warp_affine(input, M, dsize=(h, w), padding_mode='zeros')

        return out


tx = torch.tensor([-1], dtype=torch.float32)

ty = torch.tensor([-1], dtype=torch.float32)

translation = DTranslation(x_translation=tx, y_translation=ty)


class Corr(nn.Module):
    def forward(self, x, pred):
        x = x.flatten()
        p = pred.flatten()
        x_n = torch.norm(x)
        p_n = torch.norm(p)

        return torch.dot(x, p) / (x_n * p_n)


criterion = nn.MSELoss()
# criterion = nn.CosineSimilarity(1)
optimizer = torch.optim.SGD(translation.parameters(), lr=0.75)
img = torch.from_numpy(ex).permute(2, 0, 1)[None, ...]
img.requires_grad_(True)
inst = DTranslation(torch.tensor([0.2], dtype=torch.float32),
                    torch.tensor([0.2], dtype=torch.float32))

with torch.no_grad():
    gt = inst(img)
show(gt)
loss_h = []
for i in range(100):
    optimizer.zero_grad()
    result = translation(img, train=True)
    if i % 10 == 0:
        show(result)
    loss = criterion(gt.view(1, -1), result.view(1, -1))
    loss_h.append(loss.item())
    loss.backward()
    print(f'Grad: '
          f'{translation.translations.grad}, '
          f'Loss: {loss.item()}, '
          f'Value: {torch.sigmoid(translation.translations.data)}')
    optimizer.step()
plt.plot(loss_h)
plt.show()
</code></pre>
<p>I just went forward SGD to avoid statistical optimizers.</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/2/1/21b45d73cb0ab3266a8e4efffbc8e37e9ca40d49.png" alt="image" data-base62-sha1="4OabAhxATXXoHTG0hdDuQDJj6EF" width="640" height="480"></p>
<p>It get stuck.<br>
GT<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/f/9/f940248eed9066f16db7e907574cdb544e31a9a4.png" alt="image" data-base62-sha1="zyYjyFfT5XJxdwX391aUdcdaHfm" width="640" height="480"><br>
Pred N<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/a/6/a6570c9923826fca453ff020679cd990a8f8231b.png" alt="image" data-base62-sha1="nJvLfMzkMNt6fQtcyUkWXuhQZE7" width="640" height="480"><br>
You can try to reduce LR through time, using an scheduler or manually.<br>
It’s somehow ill posed<br>
For lr=0.5<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/f/1/f1a3a0784fdd68950443c813933cbe8d74a28116.png" alt="image" data-base62-sha1="ytDDDiSqlTve0RZJPJsakzsYTnE" width="640" height="480"><br>
It reaches a good result but diverge in the end.<br>
Maybe an iterative scheme which grabs the best result can help u</p>

username: Niv_Nayman
cooked: <p>Great !</p>
<p>Choosing lr=0.1 does the trick:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/f/1/f1d7e5f08661af5271962978c3091de849ec030a.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/f1d7e5f08661af5271962978c3091de849ec030a" title="Screen Shot 2020-09-15 at 17.05.34"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/f/1/f1d7e5f08661af5271962978c3091de849ec030a_2_155x500.jpeg" alt="Screen Shot 2020-09-15 at 17.05.34" data-base62-sha1="yvrD4gklJ32wiuwDaQAnsqSPOSC" width="155" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/f/1/f1d7e5f08661af5271962978c3091de849ec030a_2_155x500.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/f/1/f1d7e5f08661af5271962978c3091de849ec030a_2_232x750.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/f/1/f1d7e5f08661af5271962978c3091de849ec030a_2_310x1000.jpeg 2x" data-dominant-color="909089"><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">Screen Shot 2020-09-15 at 17.05.34</span><span class="informations">520×1668 178 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> <div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/1/a1c842e872a4e302dd50e28ce2bb105e2f64f137.png" data-download-href="https://discuss.pytorch.org/uploads/default/a1c842e872a4e302dd50e28ce2bb105e2f64f137" title="Screen Shot 2020-09-15 at 17.05.39"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/a/1/a1c842e872a4e302dd50e28ce2bb105e2f64f137_2_355x500.png" alt="Screen Shot 2020-09-15 at 17.05.39" data-base62-sha1="n5bVHI3ujMWuCQYNhN4WzGagK7J" width="355" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/a/1/a1c842e872a4e302dd50e28ce2bb105e2f64f137_2_355x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/a/1/a1c842e872a4e302dd50e28ce2bb105e2f64f137_2_532x750.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/a/1/a1c842e872a4e302dd50e28ce2bb105e2f64f137_2_710x1000.png 2x" data-dominant-color="C6C6C5"><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">Screen Shot 2020-09-15 at 17.05.39</span><span class="informations">772×1086 63.4 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>Still I fail to understand what the problem is. After all I determined the parameter in a nn.Module (LearnableAugs) and then passed it to be processed in an internal nn.Module that doesn’t have those registered as parameters. Nevertheless, I expect the gradients to flow properly as happening for all the other transforms I learn (e.g. rotation, shear, etc.). It’s a mystery to me.</p>

username: JuanFMontesinos
cooked: <p>I think you should predict directly the magnitude value.<br>
Predicting a number between 0-1 makes gradients to be multiplied by 100.</p>
<p>It smoothes the curve.<br>
The best result i got is the following one:<br>
Which goes from the original position to (20,20)<br>
But it’s really dependent on the LR.<br>
LR 50 makes it to diverge and LR 10 makes it to stuck in a minima.</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/7/f/7f34976159b9055e9ea3e4d5b5ac54d1c0e41015.png" alt="image" data-base62-sha1="i9jjxr70fdlbPmi15FGnxB8XCnz" width="640" height="480"></p>
<pre><code class="lang-auto">from torch import nn
import torch
import kornia
import imageio

import cv2
import matplotlib.pyplot as plt
import numpy as np

ex = imageio.imread(
    'https://nickelquilts.files.wordpress.com/2018/04/1-half-square-triangle-block-with-copyright.jpg') / 255.

ex = cv2.resize(ex, (100, 100)).astype(np.float32)

torch.manual_seed(666)


def show(img):
    if isinstance(img, np.ndarray):
        pass
    else:
        img = img[0].permute(1, 2, 0).detach().cpu().numpy()
    plt.imshow(img)
    plt.show()


show(ex)


class DTranslation(nn.Module):
    def __init__(self, x_translation, y_translation):
        super(DTranslation, self).__init__()
        self.translations = nn.Parameter(torch.stack([x_translation, y_translation], 1))
        self.angle = torch.tensor([0])

    def forward(self, input, train=True):
        _, _, h, w = input.shape
        if self.angle.shape[0] != input.shape[0]:
            angle = self.angle.repeat(input.shape[0])
        else:
            angle = self.angle

        if self.translations.shape[0] != input.shape[0]:
            translations = self.translations.repeat([input.shape[0], 1])
        else:
            translations = self.translations
        if train:
            # translations = translations * torch.Tensor([h, w])
            translations = translations
        else:
            # translations = translations * torch.Tensor([h, w])
            translations = translations
        # define the rotation center
        center = torch.ones(2)
        center[..., 0] = input.shape[3] / 2  # x
        center[..., 1] = input.shape[2] / 2  # y
        center = center.repeat(input.shape[0], 1)

        # define the scale factor
        scale = torch.ones(input.shape[0])

        # compute the transformation matrix
        M = kornia.get_rotation_matrix2d(center, -angle, scale)

        # Translate
        shape = list(M.shape)
        shape[-1] -= 1
        M = M + torch.cat([torch.zeros(shape), translations.unsqueeze(-1)], 2)
        # apply the transformation to original image
        out = kornia.warp_affine(input, M, dsize=(h, w), padding_mode='zeros')

        return out


tx = torch.tensor([1.], dtype=torch.float32)

ty = torch.tensor([1.], dtype=torch.float32)

translation = DTranslation(x_translation=tx, y_translation=ty)


class Corr(nn.Module):
    def forward(self, x, pred):
        x = x.flatten()
        p = pred.flatten()
        x_n = torch.norm(x)
        p_n = torch.norm(p)

        return torch.dot(x, p) / (x_n * p_n)


criterion = nn.MSELoss()
# criterion = nn.CosineSimilarity(1)
optimizer = torch.optim.SGD(translation.parameters(), lr=25)
img = torch.from_numpy(ex).permute(2, 0, 1)[None, ...]
img.requires_grad_(True)
inst = DTranslation(torch.tensor([20.], dtype=torch.float32),
                    torch.tensor([20.], dtype=torch.float32))

with torch.no_grad():
    gt = inst(img)
show(gt)
loss_h = []
N=1000
for i in range(N):
    optimizer.zero_grad()
    result = translation(img, train=True)
    if i % (N//100) == 0:
        show(result)
    loss = criterion(gt.view(1, -1), result.view(1, -1))
    loss_h.append(loss.item())
    loss.backward()
    print(f'Grad: '
          f'{translation.translations.grad}, '
          f'Loss: {loss.item()}, '
          f'Value: {translation.translations.data}')
    optimizer.step()
plt.plot(loss_h)
plt.show()
</code></pre>

username: Niv_Nayman
cooked: <p>The issue is solved, with both the kind help from <a class="mention" href="/u/juanfmontesinos">@JuanFMontesinos</a> and at <a href="https://github.com/kornia/kornia/issues/682" class="inline-onebox" rel="noopener nofollow ugc">Backpropagation through Translation · Issue #682 · kornia/kornia · GitHub</a><br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/3/7/370cecde4e836523e4ed55615a76e11d654c8bfc.gif" alt="learn_x_translation_0.30000001192092896_y_translation_-0.20000000298023224_entropy_0.01_step_" data-base62-sha1="7QZX5PZSUDpI5XJ5mQUfeO5lUYI" width="500" height="500" class="animated"></p>
<p>The  <strong>only</strong>  problem was placing  <code>translations = torch.stack([x_translation, y_translation], 1)</code>  at the  <strong>init</strong> () rather than in the forward().</p>
<p>FYI all the inplace operations do not interfere with the differentiation as those are all differentiable, i,e, both:<br>
<code>translations[:, 0] *= h translations[:, 1] *= w</code><br>
And:<br>
<code>M[..., 2] += translations</code></p>
<p>Eventually the working code is:</p>
<pre><code class="lang-auto">class DTranslation(nn.Module):
    def __init__(self, x_translation, y_translation):
        super(DTranslation, self).__init__()
        self.x_translation = x_translation
        self.y_translation = y_translation
        self.angle = torch.tensor([0])
        
        
    def forward(self, input):
        _, _, h, w = input.shape
        if self.angle.shape[0] != input.shape[0]:
            angle = self.angle.repeat(input.shape[0])
        else:
            angle = self.angle
            
        translations = torch.stack([self.x_translation, self.y_translation], 1)
        if translations.shape[0] != input.shape[0]:
            translations = translations.repeat([input.shape[0], 1])
        else:
            translations = translations
            
        translations[:, 0] *= h
        translations[:, 1] *= w
        
        # define the rotation center
        center = torch.ones(2)
        center[..., 0] = input.shape[3] / 2  # x
        center[..., 1] = input.shape[2] / 2  # y
        center = center.repeat(input.shape[0], 1)

        # define the scale factor
        scale = torch.ones(input.shape[0])

        # compute the transformation matrix
        M = kornia.get_rotation_matrix2d(center, -angle, scale)
        
        # Translate
        M[..., 2] += translations  # tx/ty
        
        # apply the transformation to original image
        out = kornia.warp_affine(input, M, dsize=(h, w))
        
        return out
</code></pre>
<p>Closing this issue.</p>