Loss "stuck" while learning affine transform

username: Hmrishav_Bandyopadhy
cooked: <p>I am trying to learn a piecewise affine transform model where the input images are converted via the transform into output data. One example transformation would be (first to second) –<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/2/f/2f77129b949a46b4cb39dbf4707264b6ee309437.jpeg" alt="dst3" data-base62-sha1="6LTyimMniSu2tOwq3OdE10VyYTR" width="128" height="128"> <img src="https://discuss.pytorch.org/uploads/default/original/3X/d/5/d59cd6794a87f18e9b8d90783f883451dc31a12b.jpeg" alt="3_" data-base62-sha1="utHIVqA9nzfFjyQY98h1MiCY3CH" width="128" height="128"><br>
The input images, are however, neither simple down samples with borders, nor distortion free rectangles, which is why I need piecewise affine transform to work.<br>
I use the following code to “learn” the transform but obviously i am going wrong somewhere as the loss is stuck from the very first epoch( i use MSE loss)</p>
<pre><code class="lang-auto">class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()
        # Spatial transformer localization-network
        self.localization = nn.Sequential(
            nn.Conv2d(3, 8, kernel_size=7),
            nn.MaxPool2d(2, stride=2),
            nn.ReLU(True),
            nn.Conv2d(8, 10, kernel_size=5),
            nn.MaxPool2d(2, stride=2),
            nn.MaxPool2d(2, stride=2),
            nn.ReLU(True)
        )

        # Regressor for the 3 * 2 affine matrix
        self.fc_loc = nn.Sequential(
            nn.Linear(10 * 30 * 30, 360),
            nn.ReLU(True),
            nn.Linear(360, 3 * 2)
        )

        # Initialize the weights/bias with identity transformation
        self.fc_loc[2].weight.data.zero_()
        

    # Spatial transformer network forward function
    def stn(self, x):

        xs = self.localization(x)
        xs = xs.view(-1, 30* 10 * 30)
        theta = self.fc_loc(xs)
        theta = theta.view(-1, 2, 3)
        
        grid = F.affine_grid(theta, x.size())
        print(x.shape)
        x = F.grid_sample(x, grid)
        return x

    def forward(self, x):
        # transform the input
        x = self.stn(x)

</code></pre>
<p>Code is heavily borrowed for the pytorch STN example as I first wanted to try out things before stepping it up.<br>
Any help would be highly appreciated <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><strong>Edit:</strong><br>
This is the output from the model –<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/d/b/db9074a42bfa23a2cf4f8b5ba47c2fe471ded929.png" alt="image" data-base62-sha1="vkm3LFzueiFleQQhPJC7ko3CPfX" width="128" height="128"><br>
The loss is stuck as the output is same, no matter what the input <img src="https://discuss.pytorch.org/images/emoji/apple/frowning_face.png?v=12" title=":frowning_face:" class="emoji" alt=":frowning_face:" loading="lazy" width="20" height="20"><br>
Are there any constraints for defining the grid ? Something that maybe I can teach the model–which it would take a long time to learn by itself.</p>

username: ptrblck
cooked: <p>The code seems to look alright, but I’m wondering if the zero initialization is a good idea:</p>
<pre><code class="lang-python">self.fc_loc[2].weight.data.zero_()
</code></pre>
<p>Could you remove it and check, if your model trains better without it?</p>

username: Hmrishav_Bandyopadhy
cooked: <p>Yeah i have tried that but it doesn’t seem to work either <img src="https://discuss.pytorch.org/images/emoji/apple/frowning_face.png?v=9" title=":frowning_face:" class="emoji" alt=":frowning_face:"></p>
<p>I have shuffled with the grid_sample function and only very specific matrices seem to work on it–which led me to believe there might be any constraints to define the matrix</p>