Per-sample loss and per-sample gradients

jose-solorzano: <p>I’d like to compute per-sample gradients given per-sample loss values and a batch of inputs. For example:</p>
<pre><code class="lang-auto">batch_size = 7
num_features = 12
x = torch.randn((batch_size, num_features)).requires_grad_()
y = x * 3.0
label = torch.randn((batch_size, num_features))
loss = torch.mean((label - y).pow(2), dim=1)
grad1 = ps_grad(loss, x, create_graph=False)
</code></pre>
<p>So I’d like to implement the <em>ps_grad</em> function above. A naive implementation would be:</p>
<pre><code class="lang-auto">def ps_grad(outputs: torch.Tensor, inputs: torch.Tensor, create_graph: bool = False):
    batch_size = inputs.size(0)
    grads = [autograd.grad(outputs[[i]], inputs[[i]], create_graph=create_graph)[0] for i in range(batch_size)]
    return torch.cat(grads, dim=1)
</code></pre>
<p>However, this errors out with:</p>
<blockquote>
<p>RuntimeError: One of the differentiated Tensors appears to not have been used in the graph. Set allow_unused=True if this is the desired behavior.</p>
</blockquote>
<p>This is presumably because individual input tensor elements did not participate in the calculation of the loss.</p>
<p>The <a href="https://pytorch.org/tutorials/intermediate/per_sample_grads.html" rel="noopener nofollow ugc">example here</a> does the equivalent of calculating the loss inside the <em>ps_grad</em> function, but that’s not an option.</p>
<p>I’ve also tried using <em>vmap</em>, but the problem is that it creates a <em>BatchedTensor</em> internally, which loses the <em>requires_grad</em> state of the inputs.</p>
<p>I’m not too concerned with performance at this point. Any ideas?</p>

jose-solorzano: <p>To answer my own question: For most/all practical purposes, this is a valid implementation of <em>ps_grad</em>:</p>
<pre><code class="lang-auto">def ps_grad(outputs: torch.Tensor, inputs: torch.Tensor, create_graph: bool = False):
    return autograd.grad(outputs.sum(), inputs, create_graph=create_graph)[0]
</code></pre>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/jose-solorzano">@jose-solorzano</a>,</p>
<p>So why exactly doesn’t <code>torch.func</code> work for your example? Could you share what you did with <code>torch.func</code>?</p>
<aside class="quote no-group" data-username="jose-solorzano" data-post="1" data-topic="202447">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/jose-solorzano/48/39012_2.png" class="avatar"> jose-solorzano:</div>
<blockquote>
<p>The <a href="https://pytorch.org/tutorials/intermediate/per_sample_grads.html" rel="noopener nofollow ugc">example here</a> does the equivalent of calculating the loss inside the <em>ps_grad</em> function, but that’s not an option.</p>
</blockquote>
</aside>
<p>Also, you ideally shouldn’t mix <code>torch.autograd.grad</code> and <code>torch.func</code> operations as they compute the gradients in different ways. The <code>.grad</code> attribute field accumulates gradients so they have the same shape as the parameter Tensor they’re attributed to.</p>
<p>In the case of <code>torch.func</code>, it creates a functionalized form of your loss and computes per-sample gradients efficiently via removing the batch dimension. In short, it defines the function for 1 sample then vectorizes over all samples. This effectively ignores the <code>.grad</code> attribute entirely as <code>torch.func</code> only sees pytorch primitives. You could potentially attribute the per-sample gradients to the <code>.grad</code> field of your parameter, but you’d get a mis-match error on the shape.</p>

jose-solorzano: <p>I mean something like this:</p>
<pre><code class="lang-auto">def ps_grad(outputs: torch.Tensor, inputs: torch.Tensor, create_graph: bool = False):
    def _singleton_grad(_out: torch.Tensor, _in: torch.Tensor):
        return autograd.grad(_out, _in, create_graph=create_graph)[0]

    return vmap(_singleton_grad)(outputs, inputs)
</code></pre>
<p>Errors out with:</p>
<blockquote>
<p>RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn</p>
</blockquote>
<p>I’m curious if there’s a right way to do that (though it doesn’t matter now, because I have a satisfactory solution for my use case.)</p>

AlphaBetaGamma96: <p>The example you show there wouldn’t work as you’re mixing <code>torch.autograd.grad</code> with <code>torch.func.vmap</code>, the <code>torch.autograd</code> and <code>torch.func</code> namespaces compute gradients in different ways and that leads to the <code>.grad</code> doesn’t exist error you stated above.</p>
<p>What you could do is replace <code>torch.autograd.grad</code> with <code>torch.func.grad</code> and compute per-sample gradients and that would be compatible with the <code>torch.func</code> namespace.</p>