Vmap over autograd.grad of a nn.Module

david-anton: <p>Hi everyone,</p>
<p>I would like to use vmap with autograd and a nn.Module. Here is minimal working example similar to what I would like to implement.</p>
<pre><code class="lang-auto">import torch
from torch import nn, vmap

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self._layers = nn.Sequential(nn.Linear(1,1))

    def forward(self, x):
        return self._layers(x)

net = Model()

batch_size = 5
x = torch.rand(batch_size, 1, requires_grad=True)

def vmap_grad(_x):
    print(_x)
    _y = net(_x)
    return torch.autograd.grad(_y, _x, grad_outputs=torch.ones_like(_x))[0]

y_dx_vmap = vmap(vmap_grad)(x)
</code></pre>
<p>When I run this code snippet, I get the following error message:</p>
<p>*BatchedTensor(lvl=1, bdim=0, value=<br>
tensor([[0.3807],<br>
[0.0681],<br>
[0.3440],<br>
[0.6405],<br>
[0.2308]], requires_grad=True)<br>
)<br>
Traceback (most recent call last):<br>
File “/workspaces/app/test.py”, line 22, in <br>
y_dx_vmap = vmap(vmap_grad)(x)<br>
^^^^^^^^^^^^^^^^^^<br>
File “/usr/local/lib/python3.11/site-packages/torch/_functorch/vmap.py”, line 436, in wrapped<br>
return _flat_vmap(<br>
^^^^^^^^^^^<br>
File “/usr/local/lib/python3.11/site-packages/torch/_functorch/vmap.py”, line 39, in fn<br>
return f(*args, **kwargs)<br>
^^^^^^^^^^^^^^^^^^<br>
File “/usr/local/lib/python3.11/site-packages/torch/_functorch/vmap.py”, line 621, in _flat_vmap<br>
batched_outputs = func(*batched_inputs, *<em>kwargs)<br>
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^<br>
File “/workspaces/app/test.py”, line 20, in vmap_grad<br>
return torch.autograd.grad(_y, _x, grad_outputs=torch.ones_like(_x))[0]<br>
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^<br>
File “/usr/local/lib/python3.11/site-packages/torch/autograd/<strong>init</strong>.py”, line 319, in grad<br>
result = Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass<br>
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^<br>
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn</em></p>
<p>Could you please help me with this and tell me what I am doing wrong? On <a href="https://pytorch.org/docs/stable/func.api.html#utilities-for-working-with-torch-nn-modules" class="inline-onebox" rel="noopener nofollow ugc">torch.func API Reference — PyTorch 2.1 documentation</a> I read that</p>
<blockquote>
<p>In general, you can transform over a function that calls a <code>torch.nn.Module</code> .</p>
</blockquote>
<p>Also, I verified that the <em>requires_grad</em> attribute for the input tensor is <em>True</em> (see print-statement in vmap_grad() and terminal output). Therefore, unfortunately, I can not understand the terminal output.</p>
<p><strong>System info:</strong></p>
<ul>
<li>PyTorch version: 2.1.0.dev20230501+cpu</li>
<li>Installed via pip</li>
<li>Python version: 3.11</li>
</ul>
<p>I would be very grateful if you could help me. Many thanks in advance!</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/david-anton">@david-anton</a>,</p>
<p>You should replace <code>torch.autograd.grad</code> with <code>torch.func.jacrev</code> and you’ll need to use <code>torch.func.functional_call</code> when calling you <code>nn.Module</code> object.</p>
<p>For example,</p>
<pre><code class="lang-auto">import torch
from torch import nn
from torch.func import jacrev, vmap, functional_call

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self._layers = nn.Sequential(nn.Linear(1,1))

    def forward(self, x):
        return self._layers(x)

net = Model()

batch_size = 5
x = torch.rand(batch_size, 1, requires_grad=True)

def vmap_grad(_x):
    print(_x)
    _y = net(_x)
    return torch.autograd.grad(_y, _x, grad_outputs=torch.ones_like(_x))[0]

#torch.func requires a functionalized call of your nn.Module
def fcall(params, x):
  return functional_call(net, params, x)
    
params = dict(net.named_parameters())
    
y_dx_vmap = vmap(jacrev(fcall, argnums=(1)), in_dims=(None, 0))(params, x)

print(x)
print(y_dx_vmap)

print(net._layers[0].weight) #derivative is just the weight of layer (here's a check)
</code></pre>
<p>This script returns,</p>
<pre><code class="lang-auto">tensor([[0.5876],
        [0.9047],
        [0.2497],
        [0.6386],
        [0.3418]], requires_grad=True)
tensor([[[0.0340]],

        [[0.0340]],

        [[0.0340]],

        [[0.0340]],

        [[0.0340]]], grad_fn=&lt;ViewBackward0&gt;)
Parameter containing:
tensor([[0.0340]], requires_grad=True)
</code></pre>

BraveDrXuTF: <p>Why must we use jacrev to replace autograd?</p>

AlphaBetaGamma96: <p>In order to use <code>torch.func.vmap</code> you need to use the respective functions within the <code>torch.func</code> namespace as functorch only works within a functional approach (i.e. params must be an input to the function you’re differentiating), which <code>torch.autograd</code> namespace doesn’t do.</p>

tangzhenjie: <p>I encountered the same issue when using <code>vmap</code>. I found that when applying <code>vmap</code> to a tensor (in the code below, <code>u.T</code>), it causes the tensor’s <code>grad_fn</code> to be lost, and the resulting batched tensor (<code>u_i</code>) has <code>requires_grad=False</code>. Are there any methods to resolve this? Thank you very much.</p>
<pre><code class="lang-auto">def pde_fn_vmap(x, u):
    """
    :param x: Input tensor of shape [num_data, 1, dim], for which gradients are required
    :param u: u(x) has shape [num_data, ens_size] and consists of ens_size neural networks
    :return: The right side of the physical model f(x), shape is [num_data, ens_size]
    """

    def compute_derivatives(u_i):
        # Calculate the first derivative
        u_x_i = torch.autograd.grad(u_i.sum(), x, create_graph=True)[0]  # 形状 [num_data, 1, dim]
        u_x1_i = u_x_i[..., 0]  # shape = [num_data, 1]
        u_x2_i = u_x_i[..., 1]  # shape = [num_data, 1]

        # Calculate the second derivative
        u_x1x1_i = torch.autograd.grad(u_x1_i.sum(), x, create_graph=True)[0][..., 0]  # 形状 [num_data, 1]
        u_x1x2_i = torch.autograd.grad(u_x1_i.sum(), x, create_graph=True)[0][..., 1]  # 形状 [num_data, 1]
        u_x2x1_i = torch.autograd.grad(u_x2_i.sum(), x, create_graph=True)[0][..., 0]  # 形状 [num_data, 1]
        u_x2x2_i = torch.autograd.grad(u_x2_i.sum(), x, create_graph=True)[0][..., 1]  # 形状 [num_data, 1]

        return u_x1_i.detach(), u_x2_i.detach(), u_x1x1_i.detach(), u_x1x2_i.detach(), u_x2x1_i.detach(), u_x2x2_i.detach()

    # Use vmap to batch compute u_i of each ens_size dimension
    results = torch.vmap(compute_derivatives)(u.T)  # u.T shape = [ens_size, num_data]

    u_x1, u_x2, u_x1x1, u_x1x2, u_x2x1, u_x2x2 = results

    u_x1 = torch.stack(u_x1, dim=-1)  # shape = [num_data, ens_size]
    u_x2 = torch.stack(u_x2, dim=-1)  # shape = [num_data, ens_size]
    u_x1x1 = torch.stack(u_x1x1, dim=-1)  # shape = [num_data, ens_size]
    u_x1x2 = torch.stack(u_x1x2, dim=-1)  # shape = [num_data, ens_size]
    u_x2x1 = torch.stack(u_x2x1, dim=-1)  # shape = [num_data, ens_size]
    u_x2x2 = torch.stack(u_x2x2, dim=-1)  # shape = [num_data, ens_size]

    f = 0.01 * (u_x1x1 + u_x2x2) + 1 * u.detach() ** 2  #  shape = [num_data, ens_size]

    return f</code></pre>

AlphaBetaGamma96: <p>I’m pretty sure my previous answer still holds here, the <code>torch.autograd</code> namespace and the <code>torch.func</code> namespace compute gradients in a different manner. I’m not a developer, so I could be wrong, but I have a feeling that you’ll need to re-write this <code>pdf_fn_vmap</code> function with the equivalent <code>torch.func</code> functions.</p>
<p><code>torch.func.vmap</code> works by creating a batched equivalent function which works on a single example, then all samples in a batch are vectorized.</p>
<p>Also, you’re differentiating with respect to <code>x</code> which isn’t even an input to the function <code>compute_derivatives</code>, try re-writing your function within a purely <code>torch.func</code> namespace based approach and see if that resolves the error.</p>

tangzhenjie: <p>Thank you very much for your reply, I will try to modify it using process-based programming.</p>