How to remove reduction in a register_hook

username: AlphaBetaGamma96
cooked: <p>Hi All,</p>
<p>I have a bit of a technical question regarding the <code>register_hook</code> function. The question I want to ask is how can I remove the reduction that’s done within <code>register_hook</code> function. I’m not sure if this is the correct terminology so please correct me if I’m wrong.</p>
<p>For example, the code below takes an input, IN, of size 1 with a batch of size BATCH. It passes it through a weight matrix to create an output OUT. The Squared Error (not MSE) is then calculated with Tensor Y and the gradient is calculated for all inputs.</p>
<pre><code class="lang-auto">import torch

BATCH=10
IN=1
NHID=5
OUT=2

X = torch.randn(BATCH, IN)
W = torch.randn(IN, NHID, requires_grad=True)
W.register_hook(lambda x: print("W: ",x))
Y = torch.randn(BATCH, NHID)

OUT=X.mm(W)
loss = (OUT-Y)**2
loss.register_hook(lambda x: print("loss: ",x))
loss.backward(torch.ones_like(loss))
</code></pre>
<p>This prints out the following</p>
<pre><code class="lang-auto">loss:  tensor([[1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.]])
W:  tensor([[  4.9739,   0.1722,   0.9640, -12.8224,   5.5129]])

</code></pre>
<p>From my understanding, the <code>register_hook</code> function calculates the gradient of the loss w.r.t to the Tensor to which it is registered. However, the gradient vector that is shown is of the same dimensionality of <code>W</code> and does not hold the batch dimension. So, it must be summed over or reduced somehow. Is it possible to disable this feature somehow so the hook returns a Tensor of (BATCH, IN, HID) ?</p>
<p>Thanks in advance!</p>

username: albanD
cooked: <p>Hi,</p>
<p>This hook just gives you the gradient that was computed for <code>W</code>. And this gradient has the same size as <code>W</code>.<br>
Note that the hook is not doing any computation, it just shows you something that was computed before.</p>
<p>I am not sure to understand why you expect to see a Batch dimension on the gradient of W here. the <code>torch.mm()</code> operation takes two inputs and compute the gradient for each of them. In particular for W, it will be <code>X.t() * grad_out</code>. As you can see the Batch dimension “naturally” disappears when you compute the gradient.</p>

username: AlphaBetaGamma96
cooked: <p>Hi!</p>
<p>Thanks for the quick response! The reason I’m asking about this is because I want to get access to the gradient of the loss w.r.t W for all inputs in a batch. So, in effective get the gradient -&gt; d(Loss)/d(W) for all inputs in a batch which would have a dimensionality along the lines of something like [hid_features, in_features, batch]</p>
<p>From my understanding, <code>register_hook</code> gives d(Loss)/d(Z) where Z is the output of a module (In my case nn.Linear)  and has the dimensionality [hid_features, in_features]. On the other hand, <code>register_backward_hook</code> gives d(Z)/d(W) which gives <code>grad_output</code> which has the dimensionality  of [batch, hid_features] corresponding to the gradient for each input against weight. My understanding mainly comes from this post here: <a href="https://discuss.pytorch.org/t/exact-meaning-of-grad-input-and-grad-output/14186/9" class="inline-onebox">Exact meaning of grad_input and grad_output</a> (I assume the explanation within this post is correct?). The <code>register_backward_hook</code> has the batch dimension within it, and I was wondering if it were possible to get the gradient of the loss with respect to W for all values in the batch in a similar manner to what the <code>register_backward_hook</code> does, except for d(Loss)/d(W) rather than d(Z)/d(W).</p>
<p>Is this something that is possible with PyTorch?</p>
<p>Thank you for the help and clarification!</p>

username: albanD
cooked: <p>Hi,</p>
<p>The gradient from the output of the module match the output size and so have a batch dimension. But the weight don’t have a batch dimension so you don’t have one for the gradients either.</p>
<p>You can do two things here:</p>
<ul>
<li>Do as many backward as there are sample in the batch. Saving the gradient each time and stack them at the end to get your <code>(batch, in, nhid)</code> gradients.</li>
<li>Expand the weights to `(batch, in, nhid) to have as many weights as you have samples. And in the forward, make sure each element in the weight is used to compute the output for each sample in your batch. Now when you will backprop, you will again get a gradient of the size of the weights that contain all the values you want.</li>
</ul>

username: AlphaBetaGamma96
cooked: <p>Thank you for your help! I shall give it a go!</p>