Repository: pytorch
Issue: Recommendations for Grid Sample/Affine Grid/Displacement Fields/Optical Flow
Labels: proposal accepted,triaged,module: interpolation

I compiled a list of recommendations (as well as outstanding issues needing attention) for improving the `grid_sample` and `affine_grid` functions (the functions that are used to form a spatial transformer module), to make them easier to use for a number of computer vision applications.
The list is not exhaustive, so feel free to help out by adding other known issues and/or recommendations.

Some of these are more quick-and-easy fixes that could be easily moved out of the way, while others require some more in-depth discussion and design (marked with  [DISCUSSION]). For the latter category, I would appreciate discussion and feedback from some of the existing stakeholders such as @ssnl, @fmassa, @zou3519, @ailzhang, @soumith, and others.

### Bugs and Issues
- [x] [Fixed in #32829] Fix the border gradient issue, which causes training to occasionally get wild results for grid points at image borders. (originally #23925)
- [ ] [#24823] Fix handling of `NaN` and `inf` grid inputs, which causes segfaults and corrupted CUDA memory (also #19826, [forum/39837](https://discuss.pytorch.org/t/what-may-cause-runtime-error-in-grid-sample/39837)).
- [x] [Fixed in #24929] `affine_grid` does currently support generating 3D affine grids. Update the documentation to support it, and update the tests to test the functionality. (originally #24821, [#24380](https://github.com/pytorch/pytorch/issues/24380#issuecomment-522294803), #24878)
- [x] [Fixed in #24929] `affine_grid` error messages: Produce a more useful error message when the affine matrix has the wrong shape or dtype (#12362), or when it is not a tensor.
- [ ] [#5565] `bilinear` mode for 3D `grid_sample` should really be `trilinear`. Personally, I would prefer to have just a single mode name for `linear`, `bilinear`, and `trilinear`, but this would bring it in line with `interpolate()`.


### Support Optical Flow More Directly
_These changes would greatly simplify the implementation of convolutional networks to do Optical Flow Estimation and Image Registration tasks._
- [x] [Fixed in #24929] Add `align_corners=False` option to make `grid_sample` resolution agnostic, which is important for predicting optical flow at less-than-full resolution (originally #20785, #23923)
- [ ] [#36108] Allow option for sampling using only the residual displacement/flow. Eliminates the need to constantly add an identity grid to the flow/displacement field, which is imprecise, slow, and very [prone to user error](https://discuss.pytorch.org/t/warp-video-frame-from-optical-flow/6013).
- [ ] [pending at #41241] Allow option for the grid/flow input to be in pixel units (absolute coordinates). This is much more natural for some applications, including image registration, where the grid-like input is actually a displacement field, and this eliminates the need for the user to convert to normalized [-1,1] coordinates, which, again, is prone to user error, (especially when users don’t know which setting of `align_corners` to target). (originally #36107)

Implementation: Combinations of these options (residual and/or pixel units) could be added to `grid_sample()`, or as suggested in #4085, put in a separate `flow_sample()` high-level function (which should probably use the same kernels as the existing `grid_sample()`). Either way, the kernel code needs to be changed very little in order to allow them. For pixel units, just skip the call to the `unnormalize` function on the grid. For residual displacement, just add the destination x,y coordinates to the (unnormalized) grid.


### Inconsistencies and Conventions
_These inconsistencies are not necessarily incorrect, so to speak, but some of these do tend to cause a lot of headaches and confusion._
- [ ] Change padding modes names in `grid_sample` to match the names of padding modes in [`torch.nn.functional.pad`](https://pytorch.org/docs/stable/nn.functional.html#torch.nn.functional.pad).
- [ ] [DISCUSSION] The function which on the user-facing python side is named `grid_sample`, is behind the scenes on both the python and the C++ side always named `grid_sampler` (notice the extra ‘r’) and `GridSampler`. Perhaps standardize it by renaming the behind-the-scenes versions to `grid_sample` (respectively `GridSample`), because this direction would not break backwards compatibility. Could still cause BC issue with saved models, though.
- [x] [Fixed in #24929] Add `align_corners` option so that `grid_sample` matches the indexing conventions of `interpolate` (originally #20785, #23923)
- [ ] [DISCUSSION] Switch from a “channels-last” to a “channels-first” convention in order to match the format of all other PyTorch tensors. That is, `[N, C, H, W]` rather than `[N, H, W, C]`, where `C` represents the grid components. Would eliminate the need to constantly call `permute()` to move channels back and forth. This would allow grids/flows/fields to be used normally in most other PyTorch functions, such as `interpolate`, among others. (The only other exceptions to the channels-first rule, as far as I know, are the FFT functions, which also expect channels-last `[N, H, W, C]`.)
Similarly, the components should be expressed as `(y,x)` instead of the current `(x,y)` to match the indexing order. This is currently [very](https://twitter.com/ducha_aiki/status/1024596593302011904) [confusing](https://github.com/pytorch/pytorch/issues/4223) to users and developers alike.
Downside: would be very BC-breaking.
- [ ] [DISCUSSION] Standardize support for half precision grids in CUDA kernels of `grid_sample`. This was partially implemented before, but in a way that caused problems for double precision (by casting everything to single precision floats). Perhaps there should be a conditional cast only if the input grid has less than single precision.


### New Features
- [ ] [#2732] Support broadcasting in `grid_sample` and `affine_grid`.
Specifically, for `affine_grid`, allow it to accept a single 2x3 2D affine matrix, or a single 3x4 3D affine matrix, with no batch dimension, and just produce a single-batch grid (or one with no batch dimension in conjunction with `grid_sample` broadcasting).
For `grid_sample`, apply broadcasting rules to the batch dimension. (Broadcasting doesn’t make sense in the spatial dimensions, of course.)
- [ ] [DISCUSSION] Allow `affine_grid` to accept a python list or a numpy `ndarray`, since when not used as part of a spatial transformer module, users often think of `affine_grid` as a constructor. For example, construct a grid tensor that rotates an image. Can possibly just have it call `torch.tensor()` on the input if not already a tensor.
- [ ] Add a cyclic/circular padding mode, where the image is treated as if repeated/tiled in each dimension
- [ ] [DISCUSSION: #25039] Add new interpolation modes, including `bicubic`, `area`, and possibly [`lanczos`](https://en.wikipedia.org/wiki/Lanczos_resampling).
- [ ] [DISCUSSION: #21457] Improve the ability of `grid_sample` to downsample while warping (that is, when the grid points are spaced out far apart). This would allow a single warp-downsample operation which is more capable than applying the operations in sequence.

### Refactoring
- [x] [Fixed in #24929] There is a ‘dispatcher’ for `affine_grid_generator` that is now left all alone in `nn/_functions/vision.py`, and its sole purpose is to decide if to use cuDNN. I don’t see why we need to keep `vision.py` around just for that. It could probably be moved either up directly into the `affine_grid` function of `nn/functional.py` or down into the C++ implementation.
- [ ] [#24470, #25014] Profile CUDA kernels vs. cuDNN kernels, and consider dropping use of cuDNN grid sampler and affine grid generator if they don’t give much extra value, as [suggested](https://github.com/pytorch/pytorch/issues/20785#issuecomment-496014108) by @soumith
- [ ] Vectorize the 3D CPU kernel of `grid_sample`. See if gives performance gain like it does for the 2D version at #10980.
- [ ] [DISCUSSION] Group together some of the code for `grid_sample` and `interpolate`. They have some shared behavior and kernel code (such as the bilinear/trilinear interpolation) that could potentially be standardized and factored out.

Feel free to open individual issues for any of these points in order to start a discussion (especially for the ones marked with DISCUSSION) and then link back to them here.