Batch gradient descent with custom model (tensor operations)

Luigi: <p>Hi everyone!</p>
<p><strong>Model description:</strong><br>
My custom model can be simplified and represented with:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/4/64776202e3d93788856fc1cb2a1d306a01793e32.png" data-download-href="https://discuss.pytorch.org/uploads/default/64776202e3d93788856fc1cb2a1d306a01793e32" title="tensors"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/6/4/64776202e3d93788856fc1cb2a1d306a01793e32_2_690x366.png" alt="tensors" data-base62-sha1="ekLwsR0EJ6CthasQI8Ka45LRhCi" width="690" height="366" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/6/4/64776202e3d93788856fc1cb2a1d306a01793e32_2_690x366.png, https://discuss.pytorch.org/uploads/default/original/3X/6/4/64776202e3d93788856fc1cb2a1d306a01793e32.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/6/4/64776202e3d93788856fc1cb2a1d306a01793e32.png 2x" data-dominant-color="F4F2F2"><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">tensors</span><span class="informations">735×390 12.1 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>These are tensor operations between some constant tensor -in grey- and a vector of  “torch.nn.Parameter(…)”, which contains the terms I want to train by gradient descent.</p>
<p>I created a “nn.Module” child class that implements these tensor operations in the forward method, while the constructor takes care of initialising the constant tensors and the vector of parameters.</p>
<p><strong>Issue:</strong><br>
How can I perform/implement (mini-)batch gradient descent using the model above, considering that the forward method carries out the computations for only <strong>one</strong> input at the time?</p>
<p>Thank you in advance for helping!</p>

albanD: <p>Hi,</p>
<p>The simplest think to do would be to make your nn.Module support batch dimension (using bmm and matmul) to replace your current version.</p>
<p>If you are using nightly builds, you might be able to get that automatically by using the new vmap Module. <a class="mention" href="/u/richard">@richard</a> will be able to help you with that if you want to go down that path <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

Luigi: <p>Hey alban, thank you for replying! <img src="https://discuss.pytorch.org/images/emoji/apple/slightly_smiling_face.png?v=9" title=":slightly_smiling_face:" class="emoji" alt=":slightly_smiling_face:"></p>
<p>I may have solved the problem, torch seems already able to process mini-batches for this kind of computation without any further effort! <img src="https://discuss.pytorch.org/images/emoji/apple/grin.png?v=9" title=":grin:" class="emoji" alt=":grin:"><br>
The only thing is to make sure you have in input a tensor of size [mini_batch, variable_to_optimise.shape].</p>

Luigi: <p>I forgot to post the (very simple) class I was using to test how the batch computation worked, in case anyone wants to try it here it is:</p>
<pre><code class="lang-auto">class linearRegression(nn.Module):

    def __init__(self):
        nn.Module.__init__(self)
        self.params = torch.nn.Parameter(torch.nn.init.normal_(torch.empty(2, 1).float(), mean=0.0, std=0.1))

        self.A = torch.tensor([1, 0], requires_grad=False, dtype=torch.float32).view(1, 2)
        self.B = torch.tensor([0, 1], requires_grad=False, dtype=torch.float32).view(1, 2)

    def forward(self, obs):
        y = (self.A @ self.params) * obs + self.B @ self.params

        return y
</code></pre>
<p>It implements the line equation: y = a*x + b, where I want to train a and b given input x and y_ground_truth.</p>
<pre><code class="lang-auto">criterion = torch.nn.MSELoss(reduction='mean')
</code></pre>
<p>Hope it hepls! <img src="https://discuss.pytorch.org/images/emoji/apple/grin.png?v=9" title=":grin:" class="emoji" alt=":grin:"></p>