Parameter trained only partly

Nalaxon: <p>Hi,<br>
I have a stack of images (B x 1 x W x H) where B ist the number of image.<br>
I want to optimize translation pararameter <code>tx</code>and <code>ty</code> for each image.<br>
For some reason only the last element of <code>tx</code> and <code>ty</code> converge in an expected result.<br>
The Gradients are converge faster to 0 for the other elements. And resulting values are not as expected.</p>
<p>At the end I attached an pytrochvz visualisation of the gradient flow.</p>
<p><strong>Can anyone tell me or point me in the direction what I miss here?</strong></p>
<p>Do I have to split gradient calculation? And if so, how would I do that?</p>
<p>Thanks a lot in advance,<br>
N.</p>
<pre><code class="lang-auto">#!/usr/bin/env python
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
from torchviz import make_dot, make_dot_from_trace
torch.autograd.set_detect_anomaly(True)

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 = (torch.arange(w).repeat((b,1)) + self.translation_x)/(w-1)
        y = (torch.arange(h).repeat((b,1)) + self.translation_y)/(h-1)
        ddf = [torch.dstack(torch.meshgrid(x[i], y[i], indexing='xy'))*2-1 for i in range(b)]
        out = TF.grid_sample(input, torch.stack(ddf)) #needs input=b,c,d,h,w, grid=b,h,w,2 #2d
        return out

#constants for testing
out_path = f'./pytorch/'
b0, w0, h0 = 3, 100, 100
#get and prepare image
ex = imageio.imread(
    'https://nickelquilts.files.wordpress.com/2018/04/1-half-square-triangle-block-with-copyright.jpg') / 255.
sa = cv2.resize(ex, (w0, h0)).astype(np.float32)
sa = cv2.cvtColor(sa, cv2.COLOR_BGR2GRAY)
sa = torch.tensor(sa).unsqueeze(0).repeat((b0,1,1,1))
sa_in = torch.empty(tuple(sa.shape))


#augment test input
sa_in[0] = TTF.crop(TF.pad(sa[0], [10, 0, 0, 0], 'constant', 0), 0, 0, w0, h0).clone()
sa_in[1] = TTF.crop(TF.pad(sa[1], [0, 10, 0, 0], 'constant', 0), 0, 10, w0, h0).clone()
sa_in[2] = TTF.crop(TF.pad(sa[2], [5, 0, 0, 0], 'constant', 0), 0, 0, w0, h0).clone()
b, c, w, h = sa_in.shape

tx = torch.tensor([0], dtype=torch.float32).repeat(b).reshape((b,1))
ty = torch.tensor([0], dtype=torch.float32).repeat(b).reshape((b,1))
    
# instantiate model
model = Translation2D(translation_x=tx, translation_y=ty)
# Instantiate optimizer
config = {'lr': 6.12693}
optimizer = torch.optim.SGD(model.parameters(), **config)

max_epochs = 50
losses = []
# dot = make_dot(model(sa_in), params=dict(model.named_parameters()), show_attrs=True, show_saved=True)
# dot.view()
for i in range(max_epochs):
    optimizer.zero_grad()

    pred = model(sa_in)

    loss = TF.mse_loss(sa, pred).sqrt()
    loss.backward()
    losses.append(loss.item())

    print(f'train Values x: {model.translation_x.data.T}\ny: {model.translation_y.data.T}')
    print(f'train Grad x: {model.translation_x.grad.T}\ny: {model.translation_y.grad.T}')
    print('---------------')
    optimizer.step()


</code></pre>
<p>Gradient flow visualization<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/9/4/9453ac8f78bbdcf7520384d813c5e2b19e052cfe.png" data-download-href="https://discuss.pytorch.org/uploads/default/9453ac8f78bbdcf7520384d813c5e2b19e052cfe" title="grad_graph"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/9/4/9453ac8f78bbdcf7520384d813c5e2b19e052cfe_2_554x500.png" alt="grad_graph" data-base62-sha1="la9W3PW2u6siNiiv0VWldYdAIdU" width="554" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/9/4/9453ac8f78bbdcf7520384d813c5e2b19e052cfe_2_554x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/9/4/9453ac8f78bbdcf7520384d813c5e2b19e052cfe_2_831x750.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/9/4/9453ac8f78bbdcf7520384d813c5e2b19e052cfe_2_1108x1000.png 2x" data-dominant-color="9FA29E"><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">grad_graph</span><span class="informations">2219×2000 303 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>