RuntimeError: derivative for aten::grid_sampler_3d_backward is not implemented

qiminchen: <p>Hi, I was trying <a href="https://github.com/threestudio-project/threestudio" rel="noopener nofollow ugc">threestudio</a> with 3D volume grid and I ran into this issue, I found a similar issue (<a href="https://github.com/pytorch/pytorch/issues/34704" class="inline-onebox" rel="noopener nofollow ugc">RuntimeError: derivative for grid_sampler_2d_backward is not implemented · Issue #34704 · pytorch/pytorch · GitHub</a>) for grid_sampler_2d_backward but no solutions yet. I wonder if anyone have any ideas?</p>
<p>Pytorch version: 2.1.0+cu121</p>
<p>The original error</p>
<pre><code class="lang-auto">File "/mnt/localssd/miniconda3/envs/threestudio/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 543, in fit                                      
    call._call_and_handle_interrupt(                                                                                                                                       
  File "/mnt/localssd/miniconda3/envs/threestudio/lib/python3.9/site-packages/pytorch_lightning/trainer/call.py", line 44, in _call_and_handle_interrupt                   
    return trainer_fn(*args, **kwargs)                                                                                                                                     
  File "/mnt/localssd/miniconda3/envs/threestudio/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 579, in _fit_impl                                
    self._run(model, ckpt_path=ckpt_path)                                                                                                                                  
  File "/mnt/localssd/miniconda3/envs/threestudio/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 986, in _run                                     
    results = self._run_stage()                                                                                                                                            
  File "/mnt/localssd/miniconda3/envs/threestudio/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1030, in _run_stage                              
    self.fit_loop.run()                   
  File "/mnt/localssd/miniconda3/envs/threestudio/lib/python3.9/site-packages/pytorch_lightning/loops/fit_loop.py", line 205, in run                                       
    self.advance()                        
  File "/mnt/localssd/miniconda3/envs/threestudio/lib/python3.9/site-packages/pytorch_lightning/loops/fit_loop.py", line 363, in advance                                   
    self.epoch_loop.run(self._data_fetcher)                                          
  File "/mnt/localssd/miniconda3/envs/threestudio/lib/python3.9/site-packages/pytorch_lightning/loops/training_epoch_loop.py", line 140, in run                            
    self.advance(data_fetcher)            
  File "/mnt/localssd/miniconda3/envs/threestudio/lib/python3.9/site-packages/pytorch_lightning/loops/training_epoch_loop.py", line 250, in advance                        
    batch_output = self.automatic_optimization.run(trainer.optimizers[0], batch_idx, kwargs)                                                                               
  File "/mnt/localssd/miniconda3/envs/threestudio/lib/python3.9/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 190, in run                         
    self._optimizer_step(batch_idx, closure)                                         
  File "/mnt/localssd/miniconda3/envs/threestudio/lib/python3.9/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 268, in _optimizer_step
    call._call_lightning_module_hook(                                                
  File "/mnt/localssd/miniconda3/envs/threestudio/lib/python3.9/site-packages/pytorch_lightning/trainer/call.py", line 159, in _call_lightning_module_hook
    output = fn(*args, **kwargs)          
  File "/mnt/localssd/miniconda3/envs/threestudio/lib/python3.9/site-packages/pytorch_lightning/core/module.py", line 1308, in optimizer_step                              
    optimizer.step(closure=optimizer_closure)                                        
  File "/mnt/localssd/miniconda3/envs/threestudio/lib/python3.9/site-packages/pytorch_lightning/core/optimizer.py", line 153, in step                                      
    step_output = self._strategy.optimizer_step(self._optimizer, closure, **kwargs)                                                                                        
  File "/mnt/localssd/miniconda3/envs/threestudio/lib/python3.9/site-packages/pytorch_lightning/strategies/strategy.py", line 238, in optimizer_step                       
    return self.precision_plugin.optimizer_step(optimizer, model=model, closure=closure, **kwargs)                                                                         
  File "/mnt/localssd/miniconda3/envs/threestudio/lib/python3.9/site-packages/pytorch_lightning/plugins/precision/amp.py", line 77, in optimizer_step                      
    closure_result = closure()            
  File "/mnt/localssd/miniconda3/envs/threestudio/lib/python3.9/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 144, in __call__
    self._result = self.closure(*args, **kwargs)                                     
  File "/mnt/localssd/miniconda3/envs/threestudio/lib/python3.9/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context                                   
    return func(*args, **kwargs)          
  File "/mnt/localssd/miniconda3/envs/threestudio/lib/python3.9/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 138, in closure
    self._backward_fn(step_output.closure_loss) 
File "/mnt/localssd/miniconda3/envs/threestudio/lib/python3.9/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 239, in backward_fn
    call._call_strategy_hook(self.trainer, "backward", loss, optimizer)
  File "/mnt/localssd/miniconda3/envs/threestudio/lib/python3.9/site-packages/pytorch_lightning/trainer/call.py", line 311, in _call_strategy_hook
    output = fn(*args, **kwargs)
  File "/mnt/localssd/miniconda3/envs/threestudio/lib/python3.9/site-packages/pytorch_lightning/strategies/strategy.py", line 212, in backward
    self.precision_plugin.backward(closure_loss, self.lightning_module, optimizer, *args, **kwargs)
  File "/mnt/localssd/miniconda3/envs/threestudio/lib/python3.9/site-packages/pytorch_lightning/plugins/precision/precision.py", line 72, in backward
    model.backward(tensor, *args, **kwargs)
  File "/mnt/localssd/miniconda3/envs/threestudio/lib/python3.9/site-packages/pytorch_lightning/core/module.py", line 1103, in backward
    loss.backward(*args, **kwargs)
  File "/mnt/localssd/miniconda3/envs/threestudio/lib/python3.9/site-packages/torch/_tensor.py", line 492, in backward
    torch.autograd.backward(
  File "/mnt/localssd/miniconda3/envs/threestudio/lib/python3.9/site-packages/torch/autograd/__init__.py", line 251, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: derivative for aten::grid_sampler_3d_backward is not implemented                                     
</code></pre>

ptrblck: <p>It seems to work for me using the example from another user:</p>
<pre data-code-wrap="python"><code class="lang-python">input = torch.tensor([[
    [[1., 2., 3.],
     [4., 5., 6.]]]],
    dtype=torch.float64,
    requires_grad=True)

# 2x3x2
grid = torch.tensor([[  # x,y
    [[ 1.,  1.],  # 1
     [-1., -1.],  # 6
     [ 0.,  1.]], # 5

    [[ 0.,  0.],     # between 2 and 5   == (2 + 5) / 2 == 3.5
     [ 0.,  0.5],    # between 3.5 and 5 == (3.5 + 5) / 2 == 4.25
     [-1., -1.]]]],  # etc
     dtype=torch.float64,
     requires_grad=True)

interpolation_mode = 'bilinear'
padding_mode = 'zeros'
align_corners = True

res = torch.nn.functional.grid_sample(input, grid, interpolation_mode, padding_mode, align_corners)
res.mean().backward()
print(input.grad)
# tensor([[[[0.3333, 0.1250, 0.0000],
#           [0.0000, 0.3750, 0.1667]]]], dtype=torch.float64)
</code></pre>
<p>Could you post a minimal and executable code snippet reproducing the error?</p>

yunxiao: <p>Similar issue with minimal example code attached:</p>
<aside class="quote quote-modified" data-post="1" data-topic="209214">
  <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/y/dbc845/48.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/runtimeerror-derivative-for-aten-grid-sampler-3d-backward-is-not-implemented-minimal-reproducible-example-attached/209214">RuntimeError: derivative for aten::grid_sampler_3d_backward is not implemented (minimal reproducible example attached)</a> <a class="badge-category__wrapper " href="/c/vision/5"><span data-category-id="5" style="--category-badge-color: #AB9364; --category-badge-text-color: #FFFFFF;" data-drop-close="true" class="badge-category " title="Topics related to either pytorch/vision or vision research related topics"><span class="badge-category__name">vision</span></span></a>
  </div>
  <blockquote>
    Hello! So I’m trying to render a normal map from a density field. One of the ways to do this to grab the gradients of the density field w.r.t. to the point locations, and later the gradients participates in loss calculation. An example 
import torch
import torch.nn.functional as F

feat = torch.randn(1, 1, 200, 200, 16) # features
inds = torch.rand(1, 1, 1, 420000, 3) # point locations

feat.requires_grad_(True)
inds.requires_grad_(True)

feat_ = F.grid_sample(feat, inds, 'bilinear', 'zeros', al…
  </blockquote>
</aside>


Zhang_michael: <p>I have the exactly same issue, when I tried to compute normal map from sdf grid and backpropagate. Did you solve that? I tried finite difference, but always OOM.</p>