updates doc

ssnl · ssnl · commit 77de15f300f4 · 2018-08-10T12:13:13.000-04:00
diff --git a/torch/nn/functional.py b/torch/nn/functional.py
@@ -2117,38 +2117,53 @@ def upsample_bilinear(input, size=None, scale_factor=None):
 
 def grid_sample(input, grid, mode='bilinear', padding_mode='zeros'):
     r"""Given an :attr:`input` and a flow-field :attr:`grid`, computes the
-    `output` using input pixel locations from the grid.
-
-    Uses bilinear interpolation to sample the input pixels.
-    Currently, only spatial (4 dimensional) and volumetric (5 dimensional)
-    inputs are supported.
-
-    For each output location, :attr:`grid` has `x`, `y`
-    input pixel locations which are used to compute output.
-    In the case of 5D inputs, :attr:`grid` has `x`, `y`, `z` pixel locations.
-
-    .. Note::
-        To avoid confusion in notation, let's note that `x` corresponds to the `width` dimension `IW`,
-        `y` corresponds to the height dimension `IH` and `z` corresponds to the `depth` dimension `ID`.
-
-    :attr:`grid` has values in the range of `[-1, 1]`. This is because the
-    pixel locations are normalized by the input height and width.
-
-    For example, values: x: -1, y: -1 is the left-top pixel of the input, and
-    values: x: 1, y: 1 is the right-bottom pixel of the input.
-
-    If :attr:`grid` has values outside the range of `[-1, 1]`, those locations
-    are handled as defined by `padding_mode`. Options are `zeros` or `border`,
-    defining those locations to use 0 or image border values as contribution
-    to the bilinear interpolation.
-
-    .. Note:: This function is used in building Spatial Transformer Networks
+    ``output`` using :attr:`input` values and pixel locations from :attr:`grid`.
+
+    Currently, only spatial (4-D) and volumetric (5-D) :attr:`input` are
+    supported.
+
+    In the spatial (4-D) case, for :attr:`input` with shape
+    :math:`(N, C, H_\text{in}, W_\text{in})` and :attr:`grid` with shape
+    :math:`(N, H_\text{out}, W_\text{out}, 2)`, the output will have shape
+    :math:`(N, C, H_\text{out}, W_\text{out})`.
+
+    For each output location ``output[n, :, h, w]``, the size-2 vector
+    ``grid[n, h, w]`` specifies :attr:`input` pixel locations ``x`` and ``y``,
+    which are used to interpolate the output value ``output[n, :, h, w]``.
+    In the case of 5D inputs, ``grid[n, d, h, w]`` specifies the
+    ``x``, ``y``, ``z`` pixel locations for interpolating
+    ``output[n, :, d, h, w]``. :attr:`mode` argument specifies ``nearest`` or
+    ``bilinear`` interpolation method to sample the input pixels.
+
+    :attr:`grid` should have most values in the range of ``[-1, 1]``. This is
+    because the pixel locations are normalized by the :attr:`input` spatial
+    dimensions. For example, values ``x = -1, y = -1`` is the left-top pixel of
+    :attr:`input`, and values  ``x = 1, y = 1`` is the right-bottom pixel of
+    :attr:`input`.
+
+    If :attr:`grid` has values outside the range of ``[-1, 1]``, those locations
+    are handled as defined by :attr:`padding_mode`. Options are
+
+        * ``padding_mode="zeros"``: use ``0`` for out-of-bound values,
+        * ``padding_mode="border"``: use border values for out-of-bound values,
+        * ``padding_mode="reflection"``: use values at locations reflected by
+          the border for out-of-bound values. For location far away from the
+          border, it will keep being reflected until becoming in bound, e.g.,
+          (normalized) pixel location ``x = -3.5`` reflects by ``-1`` and
+          becomes ``x' = 2.5``, then reflects by border ``1`` and becomes
+          ``x'' = -0.5``.
+
+    .. Note:: This function is often used in building Spatial Transformer Networks.
 
     Args:
-        input (Tensor): input batch (N x C x IH x IW) or (N x C x ID x IH x IW)
-        grid (Tensor): flow-field of size (N x OH x OW x 2) or (N x OD x OH x OW x 3)
+        input (Tensor): input of shape :math:`(N, C, H_\text{in}, W_\text{in})` (4-D case)
+                        or :math:`(N, C, D_\text{in}, H_\text{in}, W_\text{in})` (5-D case)
+        grid (Tensor): flow-field of shape :math:`(N, H_\text{out}, W_\text{out}, 2)` (4-D case)
+                       or :math:`(N, D_\text{out}, H_\text{out}, W_\text{out}, 3)` (5-D case)
+        mode (str): interpolation mode to calculate output values
+            'bilinear' | 'nearest'. Default: 'bilinear'
         padding_mode (str): padding mode for outside grid values
-            'zeros' | 'border'. Default: 'zeros'
+            'zeros' | 'border' | 'reflection'. Default: 'zeros'
 
     Returns:
         output (Tensor): output Tensor
