Calculating the divergence

username: AlexiaJM
cooked: <p>How to calculate the divergence efficiently? I’m not talking about a GAN divergence, but the actual divergence which is the sum of the partial derivative of all elements of a vector (<a href="https://en.wikipedia.org/wiki/Divergence" class="inline-onebox" rel="noopener nofollow ugc">Divergence - Wikipedia</a>).</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/2X/2/2e619e562f08aa8e240a7787c42754a3ecd91757.png" alt="image" data-base62-sha1="6Cj6Q3P6P6WfJLNd1oxTsxngmMf" width="553" height="51"></p>
<p>Assume f(x): R^d-&gt; R^d. I could use autograd to get the derivative matrix (of size d x d) and then simply take the sum of the diagonals. But this is seems terribly inefficient and wasteful. There has to be a better way!</p>

username: KFrank
cooked: <p>Hello Alexia!</p>
<aside class="quote no-group quote-modified" data-username="AlexiaJM" data-post="1" data-topic="53409" data-full="true">
<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/alexiajm/48/14877_2.png" class="avatar"> AlexiaJM:</div>
<blockquote>
<p>How to calculate the divergence efficiently?<br>
…<br>
Assume f(x): R^d-&gt; R^d. I could use autograd to get the derivative matrix (of size d x d) and then simply take the sum of the diagonals. But this is seems terribly inefficient and wasteful. There has to be a better way!</p>
</blockquote>
</aside>
<p>The short answer is to use <code>requires_grad = True</code> on<br>
each of your <code>d</code> input variables one at at time, and calculate<br>
one derivative at a time for each of your <code>d</code> output variables.<br>
Then you sum the derivatives together.</p>
<p>But …</p>
<p>I’m not aware of any easy way to do precisely what you want.<br>
The reason is that there is no automatic way to disentangle<br>
calculation that is shared by your <code>d</code> output variables.</p>
<p>In three-dimensional language, if your three output variables<br>
as functions of your three input variables are <code>F_u (x, y, z)</code>,<br>
<code>F_v (x, y, z)</code>, and <code>F_w (x, y, z)</code>, and those three functions<br>
are completely unrelated to one another, then no efficiency is<br>
is lost by calculating them separately, and using three separate<br>
runs of autograd’s <code>backward()</code> to calculate  <code>d F_u / d x</code>,<br>
<code>d F_v / d u</code>, and <code>d F_w_ / d z</code> separately.</p>
<p>If <code>F_u</code>, <code>F_v</code>, and <code>F_w</code> have almost all of their calculational<br>
work shared, very little efficiency will be lost by using one run<br>
of autograd’s <code>backward()</code> to calculate all three components<br>
(<code>x</code>, <code>y</code>, and <code>z</code>) of the gradient at the same time and discarding<br>
the off-diagonal elements.  That is, you already had to do<br>
almost all of the calculational work needed for the off-diagonal<br>
elements, so the <em>incremental</em> cost of calculating them (and<br>
then not using them) is small.</p>
<p>If <code>F_x</code>, <code>F_y</code>, and <code>F_z</code> have significant shared computation,<br>
but also significant independent computation, you will have to<br>
disentangle the shared computation by hand if you want to<br>
calculate your divergence with maximum efficiency.</p>
<p>If you calculate the three needed partial derivatives independently<br>
(using <code>requires_grad = True</code> one at a time on each variable),<br>
you will be needlessly repeating the shared computation.  But if<br>
you calculate the full gradient all at once (performing one<br>
<code>backward()</code> run with <code>requires_grad = True</code> set on all<br>
the input variables at the same time), you will be needlessly<br>
performing the off-diagonal pieces of the “independent”<br>
computation.</p>
<p>Note, in a typical multi-layer neural network, most of the<br>
computation that leads to the value of your output values<br>
will be shared, so, as a practical matter, I would use autograd<br>
to calculate the full gradient all at once, sum the diagonal<br>
elements to get the divergence, and discard the off-diagonal<br>
elements.  I would be wasting a little bit of effort to calculate the<br>
off-diagonal elements, but, because most of the calculational<br>
cost was shared in the earlier, upstream layers of the neural<br>
network, the wasted work would likely be small.</p>
<p>Good luck.</p>
<p>K. Frank</p>

username: AlexiaJM
cooked: <p>Thanks <a class="mention" href="/u/kfrank">@KFrank</a> , I appreciate the well thought response. I will try both approaches.</p>

username: abdulfatir
cooked: <p>Hi <a class="mention" href="/u/alexiajm">@AlexiaJM</a>!</p>
<p>Did you figure out a way to do perform this computation efficiently?</p>

username: Sentient07
cooked: <p>Hope I’m not late. I think there are two ways, 1) Divergence = Trace(Jacobian). This is a trivial extension of the definition, however, computation of Jacobian using <code>torch.autograd.functional.jacobian</code> might be memory intensive. 2) You can use the <code>torch.autograd.grad</code> module for this.</p>
<p>Example,</p>
<pre><code class="lang-auto">from torch.autograd import grad
# Assuming your_inp is of size (B,N,3)
g = grad(your_network_callable, your_inp, grad_outputs=torch.ones_like(your_inp), create_graph=True)[0]
</code></pre>
<p><a class="mention" href="/u/ptrblck">@ptrblck</a> can you please confirm if my soln is correct? Thank you!</p>

username: pascal_notsawo
cooked: <p>Maybe this will help you :</p><aside class="quote quote-modified" data-post="14" data-topic="14968">
  <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/pascal_notsawo/48/22829_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/how-to-compute-jacobian-matrix-in-pytorch/14968/14">How to compute Jacobian matrix in PyTorch?</a> <a class="badge-category__wrapper " href="/c/autograd/7"><span data-category-id="7" data-drop-close="true" class="badge-category " title="A category of posts relating to the autograd engine itself."><span class="badge-category__name">autograd</span></span></a>
  </div>
  <blockquote>
    Here are some functions that can help you. 
import torch

def gradient(y, x, grad_outputs=None):
    """Compute dy/dx @ grad_outputs"""
    if grad_outputs is None:
        grad_outputs = torch.ones_like(y)
    grad = torch.autograd.grad(y, [x], grad_outputs = grad_outputs, create_graph=True)[0]
    return grad

def jacobian(y, x):
    """Compute dy/dx = dy/dx @ grad_outputs; 
    for grad_outputs in [1, 0, ..., 0], [0, 1, 0, ..., 0], ...., [0, ..., 0, 1]"""
    jac = torch.zeros(y.shape[0], x.s…
  </blockquote>
</aside>