`grid_sample` in Swift

jeremy · March 23, 2019, 1:42pm

We’re going to need something like affine_grid and grid_sample from PyTorch. (Which AFAICT is much the same as remap in OpenCV). Whilst we can simply use remap on CPU, that won’t help us on GPU. (Also remap doesn’t support batch transforms.)

It appears that tf doesn’t support this. Although, in classic tf fashion, some half-hearted implementations are dumped in various other places - most notably buried in Spatial Transformer in the official tf repo:

github.com

tensorflow/models/blob/master/research/transformer/spatial_transformer.py#L63


"""


def _repeat(x, n_repeats):
    with tf.variable_scope('_repeat'):
        rep = tf.transpose(
            tf.expand_dims(tf.ones(shape=tf.stack([n_repeats, ])), 1), [1, 0])
        rep = tf.cast(rep, 'int32')
        x = tf.matmul(tf.reshape(x, (-1, 1)), rep)
        return tf.reshape(x, [-1])


def _interpolate(im, x, y, out_size):
    with tf.variable_scope('_interpolate'):
        # constants
        num_batch = tf.shape(im)[0]
        height = tf.shape(im)[1]
        width = tf.shape(im)[2]
        channels = tf.shape(im)[3]


        x = tf.cast(x, 'float32')
        y = tf.cast(y, 'float32')
        height_f = tf.cast(height, 'float32')

And something in tf.contrib:

https://www.tensorflow.org/api_docs/python/tf/contrib/resampler/resampler

And something in an external lib:

github.com

adalca/neuron/blob/master/neuron/utils.py#L40


# local imports
import pynd.ndutils as nd


# often changed file
from imp import reload
import keras
import keras.backend as K
import tensorflow as tf
reload(pl)


def interpn(vol, loc, interp_method='linear'):
"""
N-D gridded interpolation in tensorflow


vol can have more dimensions than loc[i], in which case loc[i] acts as a slice 
for the first dimensions


Parameters:
    vol: volume with size vol_shape or [*vol_shape, nb_features]
    loc: a N-long list of N-D Tensors (the interpolation locations) for the new grid
        each tensor has to have the same size (but not nec. same size as vol)

@saeta @clattner is this something that might be doable in Swift (and be performant on GPU)? Perhaps the s4tf team could see what the best approach might be?

jeremy · March 25, 2019, 2:19pm

Just discovered that cudnn already supports both these operations

https://docs.nvidia.com/deeplearning/sdk/cudnn-developer-guide/index.html#cudnnSpatialTfSamplerForward

So how do we wrap a cudnn thing in s4tf most conveniently?

vova · March 25, 2019, 3:52pm

It seems the latest OpenCV also has cuda implementation of remap (in extra modules list)
https://docs.opencv.org/master/db/d29/group__cudawarping.html