masahi commented on a change in pull request #4584: [TOPI] add 3D upsampling Op.
URL: https://github.com/apache/incubator-tvm/pull/4584#discussion_r361580306
##########
File path: topi/python/topi/image/resize.py
##########
@@ -210,3 +210,167 @@ def _bicubic(*indices):
raise ValueError('%s method is not supported.' % method)
return tvm.compute(output_shape, compute_func, name='resize',
tag=tag.INJECTIVE)
+
+def resize3d(data, size, layout="NCDHW", method="nearest_neighbor",
+ coordinate_transformation_mode="align_corners", out_dtype=None):
+ """Perform resize operation on the data.
+
+ Parameters
+ ----------
+ inputs: tvm.Tensor
+ inputs is a 5-D tensor with shape
+ [batch, channel, in_depth, in_height, in_width]
+ or [batch, in_depth, in_height, in_width, channel]
+
+ size: Tuple
+ Output resolution scale to
+
+ layout: string, optional
+ "NCDHW", "NDHWC", or "NCDHWc".
+
+ coordinate_transformation_mode: string, optional
+ Describes how to transform the coordinate in the resized tensor
+ to the coordinate in the original tensor.
+ Refer to the ONNX Resize operator specification for details.
+ Available options are "half_pixel", "align_corners" and "asymmetric".
+
+ method: {"trilinear", "nearest_neighbor"}
+ Method to be used for resizing.
+
+ out_dtype: string, optional
+ Type to return. If left None will be same as input type.
+
+ Returns
+ -------
+ output : tvm.Tensor
+ 5-D with shape [batch, channel, in_depth*scale, in_height*scale,
in_width*scale]
+ or [batch, in_depth*scale, in_height*scale, in_width*scale, channel]
+ or 5-D with shape [batch, channel-major, in_depth*scale,
in_height*scale, in_width*scale,
+ channel-minor]
+ """
+ method = method.lower()
+
+ if layout == 'NDHWC':
+ in_n, in_d, in_h, in_w, in_c = data.shape
+ output_shape = [in_n, size[0], size[1], size[2], in_c]
+ elif layout == 'NCDHW':
+ in_n, in_c, in_d, in_h, in_w = data.shape
+ output_shape = [in_n, in_c, size[0], size[1], size[2]]
+ # Otherwise layout must be NCHWxc
+ else:
+ in_n, in_c, in_d, in_h, in_w, in_cc = data.shape
+ output_shape = [in_n, in_c, size[0], size[1], size[2], in_cc]
+
+ if coordinate_transformation_mode == "align_corners":
+ z_ratio = (in_d - 1).astype('float') / (size[0] - 1)
+ y_ratio = (in_h - 1).astype('float') / (size[1] - 1)
+ x_ratio = (in_w - 1).astype('float') / (size[2] - 1)
+ elif coordinate_transformation_mode in ["asymmetric", "half_pixel"]:
+ z_ratio = (in_d).astype('float') / (size[0])
+ y_ratio = (in_h).astype('float') / (size[1])
+ x_ratio = (in_w).astype('float') / (size[2])
+ else:
+ raise ValueError("Unsupported coordinate_transformation_mode:
{}".format(
+ coordinate_transformation_mode))
+
+ def _get_pixel(n, c, z, y, x, cc):
+ z = tvm.max(tvm.min(z, in_d - 1), 0)
+ y = tvm.max(tvm.min(y, in_h - 1), 0)
+ x = tvm.max(tvm.min(x, in_w - 1), 0)
+ if layout == 'NDHWC':
+ return data(n, z, y, x, c).astype('float')
+ if layout == 'NCDHW':
+ return data(n, c, z, y, x).astype('float')
+ # else must be NCDHWxc
+ return data(n, c, z, y, x, cc).astype('float')
+
+ def _get_indices(*indices):
+ if layout == 'NDHWC':
+ n, z, y, x, c = indices
+ cc = None
+ elif layout == 'NCDHW':
+ n, c, z, y, x = indices
+ cc = None
+ else:
+ n, c, z, y, x, cc = indices
+
+ return n, c, z, y, x, cc
+
+ def _cast_output(value):
+ if out_dtype:
+ dtype = out_dtype
+ else:
+ dtype = data.dtype
+ return value.astype(dtype)
+
+ # Nearest neighbor computation
+ def _nearest_neighbor(*indices):
+ n, c, z, y, x, cc = _get_indices(*indices)
+
+ in_z = z_ratio * z
+ in_y = y_ratio * y
+ in_x = x_ratio * x
+
+ if coordinate_transformation_mode == "align_corners":
+ zint = tvm.round(in_z).astype('int32')
+ yint = tvm.round(in_y).astype('int32')
+ xint = tvm.round(in_x).astype('int32')
+ elif coordinate_transformation_mode in ["asymmetric", "half_pixel"]:
+ # Add epsilon to floor to prevent gpu rounding errors.
+ epsilon = 1e-5
+ zint = tvm.floor(in_z + epsilon).astype('int32')
+ yint = tvm.floor(in_y + epsilon).astype('int32')
+ xint = tvm.floor(in_x + epsilon).astype('int32')
+ else:
+ raise ValueError("Unsupported coordinate_transformation_mode:
{}".format(
+ coordinate_transformation_mode))
+
+ return _cast_output(_get_pixel(n, c, zint, yint, xint, cc))
+
+ # Trilinear helper functions and computation.
+ def _lerp(A, B, t):
+ return A * (1.0 - t) + B * t
+
+ def _trilinear(*indices):
+ n, c, z, y, x, cc = _get_indices(*indices)
+
+ in_z = z_ratio * z
+ in_y = y_ratio * y
+ in_x = x_ratio * x
+
Review comment:
bilinear_resize_python will be updated with my PR, you don't need to touch
it. Update trilinear_resize3d_python.py and add test cases for half_pixel
(again, see my PR for 2D case)
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
