[Newbie] PyTorch equivalent of TF tf.gradients(out,in)

andrea.migliorati: <p>Hi, I’m just learning about PyTorch because I need to migrate my TF 1.x code in order to be able to use adv robustness evaluation libraries. I found the process pretty painless so far, but I’m having problems understanding the difference between the TensorFlow function <strong>tf.gradients</strong> which computes the gradient vectors of a variable <strong>out</strong> with respect to a variable <strong>in</strong>, and the <strong>torch.autograd.grad</strong> function, which seems to work on scalars only instead. I’ve read other similar entries on this website but I couldn’t find relevant examples to my problem so here I am.</p>
<p>Going into specifics, I need to migrate the following function written in TF 1.x to PyTorch:</p>
<blockquote>
<p>def cross_lipschitz(f, x, batch_size=100, z_dim=10):<br>
reg = 0<br>
grad_matrix_list = [tf.gradients(f[:, k], x)[0] for k in range(z_dim)]<br>
for l in range(z_dim):<br>
for m in range(l + 1, z_dim):<br>
grad_diff_matrix = grad_matrix_list[l] - grad_matrix_list[m]<br>
norm_for_batch = tf.norm(grad_diff_matrix, ord=2, axis=1)<br>
reg += tf.reduce_sum(tf.square(norm_for_batch))<br>
return reg</p>
</blockquote>
<p>Can you help me out in understanding how to replace <strong>tf.gradients</strong> with <strong>torch.autograd.grad</strong> in this particular situation? I also apologize for the wrong code formatting, it’s not clear to me how to do indendation in the forum text editor.</p>

AlphaBetaGamma96: <aside class="quote no-group" data-username="andrea.migliorati" data-post="1" data-topic="130988">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/77aa72/48.png" class="avatar"> andrea.migliorati:</div>
<blockquote>
<p>and the <strong>torch.autograd.grad</strong> function, which seems to work on scalars only instead</p>
</blockquote>
</aside>
<p>Hi Andrea!</p>
<p><code>torch.autograd.grad</code> isn’t restricted to only scalar outputs, you just need to modify the <code>grad_outputs</code> argument of <code>torch.autograd.grad</code>. So, let’s say you have a function which is N-dimensional input and an M-dimensional output and a batch size of 100. You just need to add <code>torch.ones_like(Y)</code> as your 3rd argument to <code>torch.autograd.grad</code></p>
<pre><code class="lang-auto">B=100
N=4
M=6

X = torch.randn(B, N, requires_grad=True)
fc = torch.nn.Linear(N,M) #our R^N -&gt; R^M function 

Y = fc(X)

dYdX = torch.autograd.grad(Y, X, torch.ones_like(Y)) #gradient of the function
#returns torch.Size([100, 4])
</code></pre>

andrea.migliorati: <p><a class="mention" href="/u/alphabetagamma96">@AlphaBetaGamma96</a> thank you very much for the insight, I couldn’t find this piece of information around. Much appreciated</p>

PTNobel: <p>Hey <a class="mention" href="/u/alphabetagamma96">@AlphaBetaGamma96</a></p>
<p>If I can ask for a bit of clarification, what does the <code>grad_outputs</code> argument actually do?</p>
<p>What does using <code>torch.ones_like(Y)</code> do? Are there other values that change the behavior of the function?</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/ptnobel">@PTNobel</a>,</p>
<p>From my understanding, gradients are calculated via Jacobian-Vector products within PyTorch and so in order to calculate the gradient of all of the outputs we pass a Tensor filled with ones that is the same shape as the output in order for the gradient to be calculated for all components of our output, and for all samples in our batch.</p>
<p>With regards to other values, you can read in more detail in the docs <a href="https://pytorch.org/docs/stable/generated/torch.autograd.grad.html#torch.autograd.grad" class="inline-onebox" rel="noopener nofollow ugc">torch.autograd.grad — PyTorch 1.9.0 documentation</a> but there are <code>retain_graph</code> (used to keep a graph if you’re calculating gradients on the same graph multiple times), <code>create_graph</code> (used if you want to calculate higher-order derivatives), and <code>allow_unused</code> which is for inputs that have no gradient.</p>
<p>I’m not a developer so this is just my understanding (which could be wrong, so if it is, anyone please correct me). One thing that might help you understand the use of <code>grad_outputs</code> is to read up in more detail about Jacobian-Vector products and how they’re used to compute gradients in Automatic Differentiation!</p>