Efficiently computing the per pixel gradient?

MartensCedric: <p>I have a function that takes a single parameter (alpha) as an input an outputs a N by N image (2048x2048). I want to obtain the gradient of this image with respect to the parameter (alpha). I’m not talking about a sobel filter, I’m looking to see how my image changes as I change alpha. The function takes around 2 seconds to evaluate.</p>
<pre data-code-wrap="python3"><code class="lang-python3">grad_img = torch.autograd.functional.jacobian(render_with_alpha, alpha).squeeze(-1) 
</code></pre>
<p>This works, and does exactly what I want. However, it takes minutes for a 64x64 image, and I terminated the 1024x1024 after 16h before it finished (I’m looking to compute 2048x2048).</p>
<p>Another approach, which is certainly too slow is use <code>backward()</code> for each pixel. To get this working, I had to rerun the forward pass every time, which makes this method impractical (is there a way around this?).</p>
<p>Two alternative methods from the pytorch documentation appear to be jacrev and jacfwd. In my case, since I have a single input and a large output, it appears that <code>jacfwd</code> would be ideal. However, I cannot get it to work with my code. If I understand correctly, it does not use autograd and when I use it, the code has errors about missing storage.</p>
<blockquote>
<p>RuntimeError: Cannot access data pointer of Tensor that doesn’t have storage</p>
</blockquote>
<p>I have this error every time I use a view, or use detach(). I do both of these quite often in my code.</p>
<p>So how can I efficiently compute the per pixel gradient of the image?</p>

KFrank: <p>Hi Martens!</p>
<p>The short story is use <a href="https://pytorch.org/tutorials/intermediate/forward_ad_usage.html" rel="noopener nofollow ugc">forward-mode automatic differentiation.</a></p>
<aside class="quote no-group quote-modified" data-username="MartensCedric" data-post="1" data-topic="217236">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/martenscedric/48/72949_2.png" class="avatar"> MartensCedric:</div>
<blockquote>
<p>I have a function that takes a single parameter (alpha) as an input an outputs a N by N image (2048x2048). I want to obtain the gradient of this image with respect to the parameter (alpha).<br>
…<br>
Another approach, which is certainly too slow is use <code>backward()</code> for each pixel.</p>
</blockquote>
</aside>
<p>Backward-mode (regular) autograd computes the gradient of a single scalar output value<br>
with respect to a multi-dimensional set of parameters with one forward / backward pass.<br>
(More generally, it computes the vector-jacobian product for a single “vector” associated<br>
with a non-scalar output tensor.)  If you wish to use backward-mode autograd to compute<br>
gradients of multiple elements of some output tensor, you have to pay the price of running<br>
multiple forward / backward passes (even if you have pytorch do this for you under the<br>
hood).</p>
<p>On the other hand, <em>forward-mode</em> autograd computes the vector derivative of an output<br>
tensor with respect to a single scalar input (that is, it computes the derivatives of the<br>
elements of the output tenor with respect to the input) with a single (dual-mode) forward<br>
pass.  (More generally, it computes the jacobian-vector product for a single “vector”<br>
associated with a non-scalar input tensor.)  If you wish to use forward-mode autograd<br>
to compute the derivative of the output tensor with respect to multiple elements of some<br>
input tensor, you have to pay the price of running multiple dual-mode forward passes.</p>
<p>(As a consequence of this we have the general rule that if you want to compute the<br>
full jacobian of some tensor output with respect to some tensor input, you would prefer<br>
backward-mode when the number of elements of the input is significantly larger than the<br>
number of elements of the output and prefer forward mode when the size of the output<br>
is larger than the size of the input.)</p>
<p>In your case, you wish to compute derivatives of many output values (your pixels) with<br>
respect to a single “input” value (your <code>alpha</code>), so you are likely to get dramatically<br>
improved performance if you use forward-mode autograd with a single dual-mode<br>
forward pass.</p>
<p>Here is an example script:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

alpha = torch.zeros (1)                        # a single scalar parameter
direc = torch.ones (1)                         # the (trivial) direction for the directional derivative

coeffs = torch.arange (7.)                     # these will be the derivatives with respect to alpha

print ('coeffs = ...')
print (coeffs)

with torch.autograd.forward_ad.dual_level():   # forward-mode context manager
    alpha_dual = torch.autograd.forward_ad.make_dual (alpha, direc)
    output = coeffs * alpha_dual               # (dual version of) vector output
    primal, tangent = torch.autograd.forward_ad.unpack_dual (output)
    print ('tangent = ...')
    print (tangent)                            # vector derivative of output with respect to scalar alpha
</code></pre>
<p>And here is its output:</p>
<pre><code class="lang-plaintext">2.6.0+cu126
coeffs = ...
tensor([0., 1., 2., 3., 4., 5., 6.])
tangent = ...
tensor([0., 1., 2., 3., 4., 5., 6.])
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

MartensCedric: <p>Thanks for the answer. Unfortunately, I’m having problems with running my code with this. The code works fine without.</p>
<pre data-code-wrap="python3"><code class="lang-python3">Cell In[5], line 24
     23 alpha_dual = torch.autograd.forward_ad.make_dual(alpha, torch.tensor(1.0)) 
---&gt; 24 output_dual = render_with_alpha(alpha_dual)
     25 primal, tangent = torch.autograd.forward_ad.unpack_dual(output_dual)

Cell In[5], line 6, in render_with_alpha(a)
      4 modified_vertices = vertices.clone()
      5 modified_vertices[:, 1] += a
----&gt; 6 render, _ = render_result_func(modified_vertices, smoothing=1, temperature=1)
      7 return render

Cell In[3], line 21, in render_result_func(new_vertices, smoothing, temperature, view_n)
     19 visibility_start = pc()
     20 with torch.no_grad():
---&gt; 21     visibility_weights = get_visibility_weights(new_vertices, faces, edges, face_pairs_vec, is_boundary_edge, view_n, temperature=temperature)
     23 edge_weights = ocg_weights * visibility_weights
     24 render_res_start = pc()

File ~/diff-oc/diffoc/geometry/visibility.py:89, in get_visibility_weights(vertices, faces, edges, face_pairs_vec, is_boundary_edge, view_n, temperature)
     87 print(f"edge_midpoints.device : {edge_midpoints.device}")
     88 print(f"ray_origins.device : {ray_origins.device}")
---&gt; 89 faces_z[n_rays_arange, fp_0] = -(edge_midpoints[:,2] - ray_origins[:,2])
     90 ray_directions = query_vertex_rep - ray_origins
     91 mesh = trimesh.Trimesh(vertices=vertices.detach().cpu(), faces=faces.detach().cpu(), process=False, use_embree=True)

File ~/anaconda3/envs/diffoc/lib/python3.12/site-packages/torch/utils/_device.py:104, in DeviceContext.__torch_function__(self, func, types, args, kwargs)
    102 if func in _device_constructors() and kwargs.get('device') is None:
    103     kwargs['device'] = self.device
--&gt; 104 return func(*args, **kwargs)

RuntimeError: CUDA error: unknown error
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
</code></pre>
<p>How would changing from backward differentiation to forward differentiation reveal any CUDA errors? Do I have hopes to debug this without recompiling PyTorch?</p>

KFrank: <p>Hi Martens!</p>
<aside class="quote no-group" data-username="MartensCedric" data-post="3" data-topic="217236">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/martenscedric/48/72949_2.png" class="avatar"> MartensCedric:</div>
<blockquote>
<p><code>RuntimeError: CUDA error: unknown error</code></p>
</blockquote>
</aside>
<p>I don’t have anything to offer about this sort of CUDA error.</p>
<p>You might try running your computation entirely on the cpu.  Maybe that would produce<br>
a more informative error message or perhaps the error would go away.</p>
<p>I would also recommend posting the pytorch and cuda versions that you are using in the<br>
hope that someone may have seen this before.  If you’re not already doing so, you should<br>
try running this on the latest version of pytorch / cuda (assuming that the latest version runs<br>
on your gpu).</p>
<p>Best.</p>
<p>K. Frank</p>