Image Sampler 3D

magz: <p>Hey everyone,<br>
since I’m dealing with 3D (depth, height, width) data, I’m wondering if someone has implemented a fast image warping function in PyTorch. I’m aware of the grid_sample function, but so far, it only supports two dimensional data. If anyone has some hints where to look, please let me know! Or if anyone from the core team could provide some information on wether the expansion to 3D is planned in the near future, that would be great!<br>
Thank you all in advance.</p>

smth: <p>The function is implemented at it’s core in these two files:</p>
<p><aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/e33df2b88a31c08567c81f49c45f1eb530cd7ef4/aten/src/THNN/generic/SpatialGridSamplerBilinear.c" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/e33df2b88a31c08567c81f49c45f1eb530cd7ef4/aten/src/THNN/generic/SpatialGridSamplerBilinear.c" target="_blank" rel="nofollow noopener">pytorch/pytorch/blob/e33df2b88a31c08567c81f49c45f1eb530cd7ef4/aten/src/THNN/generic/SpatialGridSamplerBilinear.c</a></h4>
<pre><code class="lang-c">#ifndef TH_GENERIC_FILE
#define TH_GENERIC_FILE "generic/SpatialGridSamplerBilinear.c"
#else

#undef MIN
#define MIN(a,b) ( ((a)&lt;(b)) ? (a) : (b) )
#undef MAX
#define MAX(a,b) ( ((a)&gt;(b)) ? (a) : (b) )

#undef MODE_BORDER
#define MODE_BORDER 1

static inline void THNN_(SpatialGridSamplerBilinear_shapeCheck)
     (THTensor *input, THTensor *grid, THTensor *gradOutput) {
  THNN_ARGCHECK(input-&gt;nDimension == 4, 2, input,
    "4D input tensor expected but got: %s");
  THNN_ARGCHECK(grid-&gt;nDimension == 4, 2, grid,
    "4D grid tensor expected but got: %s");

  int nbatch   = THTensor_(size)(input, 0);
</code></pre>

  This file has been truncated. <a href="https://github.com/pytorch/pytorch/blob/e33df2b88a31c08567c81f49c45f1eb530cd7ef4/aten/src/THNN/generic/SpatialGridSamplerBilinear.c" target="_blank" rel="nofollow noopener">show original</a>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
<br>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/e33df2b88a31c08567c81f49c45f1eb530cd7ef4/aten/src/THCUNN/SpatialGridSamplerBilinear.cu" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/e33df2b88a31c08567c81f49c45f1eb530cd7ef4/aten/src/THCUNN/SpatialGridSamplerBilinear.cu" target="_blank" rel="nofollow noopener">pytorch/pytorch/blob/e33df2b88a31c08567c81f49c45f1eb530cd7ef4/aten/src/THCUNN/SpatialGridSamplerBilinear.cu</a></h4>
<pre><code class="lang-cu">#include "THCUNN.h"
#include "common.h"
#include "THCDeviceTensor.cuh"
#include "THCDeviceTensorUtils.cuh"
#include "THCDeviceUtils.cuh"
#include "THCHalf.h"
#include "THCHalfAutoNumerics.cuh"
#include "THCAtomics.cuh"

#define WITHIN_BOUNDS(x, y, H, W) (x &gt;= 0 &amp;&amp; x &lt; W &amp;&amp; y &gt;= 0 &amp;&amp; y &lt; H)
#define SAFE_ADD(input, x, y, n, c, H, W, value)    \
  do {    \
    if (WITHIN_BOUNDS(x, y, H, W)) {    \
      atomicAdd(&amp;input[n][c][y][x], value);   \
    }   \
  } while(0)

#undef MIN
#define MIN(a,b) ( ((a)&lt;(b)) ? (a) : (b) )
#undef MAX
</code></pre>

  This file has been truncated. <a href="https://github.com/pytorch/pytorch/blob/e33df2b88a31c08567c81f49c45f1eb530cd7ef4/aten/src/THCUNN/SpatialGridSamplerBilinear.cu" target="_blank" rel="nofollow noopener">show original</a>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
</p>
<p>It is a relatively small task to extend this for 3D inputs. One has to make VolumetricGridSamplerBilinear.c / cu and call them from here: <a href="https://github.com/pytorch/pytorch/blob/1c0fbd27a10b7036b1db643014e65cae7f6d0266/torch/nn/_functions/vision.py#L52">https://github.com/pytorch/pytorch/blob/1c0fbd27a10b7036b1db643014e65cae7f6d0266/torch/nn/_functions/vision.py#L52</a></p>

magz: <p>Thank you for this fast reply!<br>
Since I’m not familiar with writing code in cuda, this will surely be a hard task for me, but I’ll try my best <img src="//discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=5" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

smth: <p>The way the code is written, you only have to add an additional for loop everywhere (even in CUDA case), so it shouldn’t need any working knowledge of cuda <img src="//discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=5" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

magz: <p>Awesome. I’ll give it a shot <img src="//discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=5" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

schurinkje: <p>Is there any chance that this will be implemented in the soon future?<br>
If it’s a relatively small task, then why not implement it in the core? <img src="https://discuss.pytorch.org/images/emoji/apple/slightly_smiling_face.png?v=6" title=":slightly_smiling_face:" class="emoji" alt=":slightly_smiling_face:"></p>

smth: <p><code>grid_sample</code> supports 3D since February, and release <code>0.4.1</code> supports this feature:</p>
<p>See latest doc: <a href="https://pytorch.org/docs/stable/nn.html#torch.nn.functional.grid_sample">https://pytorch.org/docs/stable/nn.html#torch.nn.functional.grid_sample</a></p>

Mohamed_Hassan: <p>Hi <a class="mention" href="/u/smth">@smth</a>  I am trying to use <code>grid_sample</code> to implement stn for 3d data. Here is a simple test code</p>
<pre><code class="lang-auto">import torch

def create_grid(height=32, width=32, depth=32):
    # create normalized 2D grid
    x = torch.linspace(-1, 1, width)
    y = torch.linspace(-1, 1, height)
    z = torch.linspace(-1, 1, depth)
    x_t, y_t, z_t = torch.meshgrid([x, y, z])

    # flatten
    x_t_flat = x_t.contiguous().view(-1)
    y_t_flat = y_t.contiguous().view(-1)
    z_t_flat = z_t.contiguous().view(-1)

    # reshape to [x_t, y_t , 1] - (homogeneous form)
    ones = torch.ones_like(x_t_flat)
    sampling_grid = torch.stack((x_t_flat, y_t_flat, z_t_flat, ones), 0)

    # repeat grid num_batch times
    print(sampling_grid.shape)
    sampling_grid = sampling_grid.unsqueeze(0)
    print(sampling_grid.shape)
    return sampling_grid

def stn(grid, sampling_grid, theta, grid_size):
    # grab batch size
    num_batch = theta.shape[0]
    sampling_grid = sampling_grid.repeat(num_batch, 1, 1)
    sampling_grid_t = torch.matmul(theta, sampling_grid) # B*3*HWD
    sampling_grid_t = sampling_grid_t.permute([0,2,1]) # B*HWD*3

    # sampling_grid_t = torch.tensor(sampling_grid_t.numpy().reshape(num_batch, grid_size, grid_size, grid_size, 3))
    sampling_grid_t = sampling_grid_t.view(num_batch, grid_size, grid_size, grid_size, 3)
    out = torch.nn.functional.grid_sample(grid, sampling_grid_t)
    return out
</code></pre>
<p>But when i pass the <code>stn</code> a 3d grid with identity transformation</p>
<pre><code class="lang-auto">theta = torch.cat((torch.eye(3),torch.zeros(3,1)), 1)
theta = theta.unsqueeze(0)
</code></pre>
<p>I got a rotated grid for some reason as shown in the picture below. I suspect it’s a problem with the way the sampling grid is defined. but there are no details in the documentation about how to do this in 3d. Can you please help?<br>
<img src="https://discuss.pytorch.org/uploads/default/original/2X/b/bac429841930883537189ec8d0515ef75365a7b8.png" alt="Screenshot%20from%202018-09-20%2015-11-47" data-base62-sha1="qEd5gOiJ6G9yVggxYqEQLo6t0qQ" width="400" height="378"></p>

soulslicer: <p>were you able to figure this out?</p>

tyoc213: <p>How did you plot this? is your data or what the grid_sampler is doing?</p>