Restruct upsample module (#5524)

* fix upsample nearest bug

* fix upsample nearest bug (#5347)

Co-authored-by: oneflow-ci-bot <69100618+oneflow-ci-bot@users.noreply.github.com>

* fix upsample bilinear bug

* restruct upsample module

* auto format by CI

* change docs

* fix bug

* fix bug

Co-authored-by: oneflow-ci-bot <69100618+oneflow-ci-bot@users.noreply.github.com>
Co-authored-by: oneflow-ci-bot <ci-bot@oneflow.org>

oneflow/python/nn/modules/interpolate.py

@@ -83,7 +83,7 @@ class Interpolate(Module):
             raise ValueError("only one of size or scale_factor should be defined")
         elif self.size is not None:
             assert self.scale_factor is None
-            scale_factors = None
+            scale_factors = []
             if isinstance(self.size, (list, tuple)):
                 if len(self.size) != dim:
                     raise ValueError(
@@ -93,6 +93,8 @@ class Interpolate(Module):
                 output_size = self.size
             else:
                 output_size = [self.size for _ in range(dim)]
+            for i in range(dim):
+                scale_factors.append(output_size[i] / x.shape[i + 2])
         elif self.scale_factor is not None:
             assert self.size is None
             output_size = None
@@ -132,13 +134,15 @@ class Interpolate(Module):
         if self.mode == "area" and output_size is None:
             self.recompute_scale_factor = True
 
-        if self.recompute_scale_factor is not None and self.recompute_scale_factor:
+        if self.recompute_scale_factor is True:
             assert scale_factors is not None
             output_size = [
                 int(math.floor(float(input.size(i + 2)) * scale_factors[i]))
                 for i in range(dim)
             ]
-            scale_factors = None
+            scale_factors = []
+            for i in range(dim):
+                scale_factors.append(output_size[i] / x.shape[2 + i])
 
         if len(x.shape) == 3 and self.mode == "nearest":
             return flow.F.upsample_nearest_1d(

oneflow/python/nn/modules/upsampling.py

@@ -23,34 +23,40 @@ from typing import Optional, Union, Tuple
 @oneflow_export("nn.Upsample")
 @experimental_api
 class Upsample(Module):
-    r"""Upsamples a given multi-channel 2D (spatial) data.
+    r"""The interface is consistent with PyTorch.    
+    
+    The documentation is referenced from: https://pytorch.org/docs/1.9.0/_modules/torch/nn/modules/upsampling.html#Upsample
+    
+    Upsamples a given multi-channel 1D (temporal), 2D (spatial) or 3D (volumetric) data.
 
     The input data is assumed to be of the form
-    `minibatch x channels x height x width`.
-    Hence, for spatial inputs, we expect a 4D Tensor.
+    `minibatch x channels x [optional depth] x [optional height] x width`.
+    Hence, for spatial inputs, we expect a 4D Tensor and for volumetric inputs, we expect a 5D Tensor.
 
-    The algorithms available for upsampling are nearest neighbor,
-    bilinear, 4D input Tensor, respectively.
+    The algorithms available for upsampling are nearest neighbor and linear,
+    bilinear, bicubic and trilinear for 3D, 4D and 5D input Tensor,
+    respectively.
 
     One can either give a :attr:`scale_factor` or the target output :attr:`size` to
     calculate the output size. (You cannot give both, as it is ambiguous)
 
     Args:
-        size (int or Tuple[int, int] optional):
+        size (int or Tuple[int] or Tuple[int, int] or Tuple[int, int, int], optional):
             output spatial sizes
-        scale_factor (float or Tuple[float, float], optional):
+        scale_factor (float or Tuple[float] or Tuple[float, float] or Tuple[float, float, float], optional):
             multiplier for spatial size. Has to match input size if it is a tuple.
         mode (str, optional): the upsampling algorithm: one of ``'nearest'``,
-            ``'bilinear'``.
+            ``'linear'``, ``'bilinear'``, ``'bicubic'`` and ``'trilinear'``.
             Default: ``'nearest'``
         align_corners (bool, optional): if ``True``, the corner pixels of the input
             and output tensors are aligned, and thus preserving the values at
-            those pixels. This only has effect when :attr:`mode` is ``'bilinear'``.
-            Default: ``False``
+            those pixels. This only has effect when :attr:`mode` is
+            ``'linear'``, ``'bilinear'``, or ``'trilinear'``. Default: ``False``
 
     Shape:
-        - Input: : :math:`(N, C, H_{in}, W_{in})`
-        - Output: :math:`(N, C, H_{out}, W_{out})` , where
+        - Input: :math:`(N, C, W_{in})`, :math:`(N, C, H_{in}, W_{in})` or :math:`(N, C, D_{in}, H_{in}, W_{in})`
+        - Output: :math:`(N, C, W_{out})`, :math:`(N, C, H_{out}, W_{out})`
+          or :math:`(N, C, D_{out}, H_{out}, W_{out})`, where
 
     .. math::
         D_{out} = \left\lfloor D_{in} \times \text{scale_factor} \right\rfloor
@@ -61,16 +67,25 @@ class Upsample(Module):
     .. math::
         W_{out} = \left\lfloor W_{in} \times \text{scale_factor} \right\rfloor
 
+    .. warning::
+        With ``align_corners = True``, the linearly interpolating modes
+        (`linear`, `bilinear`, `bicubic`, and `trilinear`) don't proportionally
+        align the output and input pixels, and thus the output values can depend
+        on the input size. This was the default behavior for these modes up to
+        version 0.3.1. Since then, the default behavior is
+        ``align_corners = False``. See below for concrete examples on how this
+        affects the outputs.
+
     .. note::
         If you want downsampling/general resizing, you should use :func:`~nn.functional.interpolate`.
 
+
     For example:
 
     .. code-block:: python
 
         >>> import numpy as np
         >>> import oneflow.experimental as flow
-        >>> flow.enable_eager_execution()
 
         >>> input = flow.Tensor(np.arange(1, 5).reshape((1, 1, 2, 2)), dtype=flow.float32)
         >>> input = input.to("cuda")
@@ -92,59 +107,18 @@ class Upsample(Module):
     ):
         super().__init__()
         self.size = size
-        if isinstance(scale_factor, tuple):
-            self.scale_factor = tuple(float(factor) for factor in scale_factor)
-        else:
-            self.scale_factor = float(scale_factor) if scale_factor else None
-
+        self.scale_factor = scale_factor
         self.mode = mode
-        if align_corners == None:
-            align_corners = False
-
         self.align_corners = align_corners
-        self.height_scale = None
-        self.width_scale = None
-
-        if isinstance(self.scale_factor, float):
-            self.height_scale = self.scale_factor
-            self.width_scale = self.scale_factor
-        elif isinstance(self.scale_factor, tuple):
-            self.height_scale = self.scale_factor[0]
-            self.width_scale = self.scale_factor[1]
-        else:
-            pass
-
-        if self.mode != "nearest" and self.mode != "bilinear":
-            raise ValueError('interpolation must be "nearest" or "bilinear".')
-
-        if self.mode == "nearest" and self.align_corners:
-            raise ValueError('interpolation "nearest" does not support align_corners.')
 
     def forward(self, x):
-        assert (
-            self.size != None or self.scale_factor != None
-        ), f"size and scale_factor can not be none at the same time!"
-        h, w = x.shape[2], x.shape[3]
-        if self.height_scale == None:
-            if isinstance(self.size, int):
-                self.height_scale = 1.0 * self.size / h
-            else:
-                self.height_scale = 1.0 * self.size[0] / h
-        if self.width_scale == None:
-            if isinstance(self.size, int):
-                self.width_scale = 1.0 * self.size / w
-            else:
-                self.width_scale = 1.0 * self.size[1] / w
-
-        res = flow.F.upsample(
+        return flow.experimental.nn.functional.interpolate(
             x,
-            height_scale=self.height_scale,
-            width_scale=self.width_scale,
+            size=self.size,
+            scale_factor=self.scale_factor,
+            mode=self.mode,
             align_corners=self.align_corners,
-            interpolation=self.mode,
-            data_format="channels_first",
         )
-        return res
 
     def extra_repr(self) -> str:
         if self.scale_factor is not None: