diff --git a/mindspore/ccsrc/backend/kernel_compiler/cpu/grid_sampler_3d_cpu_kernel.cc b/mindspore/ccsrc/backend/kernel_compiler/cpu/grid_sampler_3d_cpu_kernel.cc new file mode 100644 index 00000000000..7b166fa64ac --- /dev/null +++ b/mindspore/ccsrc/backend/kernel_compiler/cpu/grid_sampler_3d_cpu_kernel.cc @@ -0,0 +1,217 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "backend/kernel_compiler/cpu/grid_sampler_3d_cpu_kernel.h" +#include "runtime/device/cpu/cpu_device_address.h" + +namespace { +const size_t kDataSizeThreshold = 64 * 1024; +const size_t kZero = 0; +const size_t kOne = 1; +const size_t kTwo = 2; +const size_t kThree = 3; +const size_t kFour = 4; +const size_t kFive = 5; +} // namespace + +namespace mindspore { +namespace kernel { +void GridSampler3DCPUKernel::InitKernel(const CNodePtr &kernel_node) { + x_shape_ = AnfAlgo::GetInputDeviceShape(kernel_node, kZero); + grid_shape_ = AnfAlgo::GetInputDeviceShape(kernel_node, kOne); + output_shape_ = AnfAlgo::GetOutputDeviceShape(kernel_node, kZero); + output_number_ = + output_shape_[kZero] * output_shape_[kOne] * output_shape_[kTwo] * output_shape_[kThree] * output_shape_[kFour]; + dtype_ = AnfAlgo::GetInputDeviceDataType(kernel_node, kZero); + size_t stride_tmp = kOne; + auto stride_compute = [&](std::vector &stride, std::vector shape) { + for (size_t i = kFour; i > -kOne; i--) { + stride.insert(stride.begin(), stride_tmp); + stride_tmp *= shape[i]; + } + stride_tmp = kOne; + }; + stride_compute(x_stride_, x_shape_); + stride_compute(grid_stride_, grid_shape_); + stride_compute(output_stride_, output_shape_); + interpolation_mode = AnfAlgo::GetNodeAttr(kernel_node, "interpolation_mode"); + padding_mode = AnfAlgo::GetNodeAttr(kernel_node, "padding_mode"); + align_corners = AnfAlgo::GetNodeAttr(kernel_node, "align_corners"); +} + +bool GridSampler3DCPUKernel::Launch(const std::vector &inputs, + const std::vector &, + const std::vector &outputs) { + if (dtype_ == kNumberTypeFloat32) { + LaunchKernel(inputs, outputs); + } else if (dtype_ == kNumberTypeFloat64) { + LaunchKernel(inputs, outputs); + } else { + MS_EXCEPTION(TypeError) << "Input dtype only support float32 and float64!"; + } + return true; +} + +template +void GridSampler3DCPUKernel::ComputeTask(T *x_addr, T *grid_addr, T *output_addr, const size_t &seq) { + size_t out_iter[kFive] = {kZero, seq, kZero, kZero, kZero}; + size_t count = kFour; + while (out_iter[kOne] > kZero) { + if (count == kOne) { + count--; + } + out_iter[count] = out_iter[kOne] % output_shape_[count]; + out_iter[kOne] /= output_shape_[count--]; + } + const size_t out_c = output_shape_[kOne]; + int64_t grid_offset = out_iter[kZero] * grid_stride_[kZero] + out_iter[kTwo] * grid_stride_[kOne] + + out_iter[kThree] * grid_stride_[kTwo] + out_iter[kFour] * grid_stride_[kThree]; + T x = grid_addr[grid_offset]; + T y = grid_addr[grid_offset + grid_stride_[kFour]]; + T z = grid_addr[grid_offset + kTwo * grid_stride_[kFour]]; + x = grid_sampler_compute_source_index(x, x_shape_[kFour], padding_mode, align_corners); + y = grid_sampler_compute_source_index(y, x_shape_[kThree], padding_mode, align_corners); + z = grid_sampler_compute_source_index(z, x_shape_[kTwo], padding_mode, align_corners); + auto x_ptr_NC = out_iter[kZero] * x_stride_[kZero]; + auto output_ptr_NCDHW = out_iter[kZero] * output_stride_[kZero] + out_iter[kTwo] * output_stride_[kTwo] + + out_iter[kThree] * output_stride_[kThree] + out_iter[kFour] * output_stride_[kFour]; + if (interpolation_mode == "bilinear") { + int64_t x_tnw = static_cast(std::floor(x)); + int64_t y_tnw = static_cast(std::floor(y)); + int64_t z_tnw = static_cast(std::floor(z)); + int64_t x_tne = x_tnw + kOne, y_tne = y_tnw, z_tne = z_tnw; + int64_t x_tsw = x_tnw, y_tsw = y_tnw + kOne, z_tsw = z_tnw; + int64_t x_tse = x_tnw + kOne, y_tse = y_tnw + kOne, z_tse = z_tnw; + int64_t x_bnw = x_tnw, y_bnw = y_tnw, z_bnw = z_tnw + kOne; + int64_t x_bne = x_tnw + kOne, y_bne = y_tnw, z_bne = z_tnw + kOne; + int64_t x_bsw = x_tnw, y_bsw = y_tnw + kOne, z_bsw = z_tnw + kOne; + int64_t x_bse = x_tnw + kOne, y_bse = y_tnw + kOne, z_bse = z_tnw + kOne; + T tnw = (x_bse - x) * (y_bse - y) * (z_bse - z), tne = (x - x_bsw) * (y_bsw - y) * (z_bsw - z); + T tsw = (x_bne - x) * (y - y_bne) * (z_bne - z), tse = (x - x_bnw) * (y - y_bnw) * (z_bnw - z); + T bnw = (x_tse - x) * (y_tse - y) * (z - z_tse), bne = (x - x_tsw) * (y_tsw - y) * (z - z_tsw); + T bsw = (x_tne - x) * (y - y_tne) * (z - z_tne), bse = (x - x_tnw) * (y - y_tnw) * (z - z_tnw); + for (size_t c = kZero; c < out_c; c++, x_ptr_NC += x_stride_[kOne], output_ptr_NCDHW += output_stride_[kOne]) { + output_addr[output_ptr_NCDHW] = static_cast(kZero); + if (within_bounds_3d(z_tnw, y_tnw, x_tnw, x_shape_[kTwo], x_shape_[kThree], x_shape_[kFour])) { + auto x_index = x_ptr_NC + z_tnw * x_stride_[kTwo] + y_tnw * x_stride_[kThree] + x_tnw * x_stride_[kFour]; + output_addr[output_ptr_NCDHW] += x_addr[x_index] * tnw; + } + if (within_bounds_3d(z_tne, y_tne, x_tne, x_shape_[kTwo], x_shape_[kThree], x_shape_[kFour])) { + auto x_index = x_ptr_NC + z_tne * x_stride_[kTwo] + y_tne * x_stride_[kThree] + x_tne * x_stride_[kFour]; + output_addr[output_ptr_NCDHW] += x_addr[x_index] * tne; + } + if (within_bounds_3d(z_tsw, y_tsw, x_tsw, x_shape_[kTwo], x_shape_[kThree], x_shape_[kFour])) { + auto x_index = x_ptr_NC + z_tsw * x_stride_[kTwo] + y_tsw * x_stride_[kThree] + x_tsw * x_stride_[kFour]; + output_addr[output_ptr_NCDHW] += x_addr[x_index] * tsw; + } + if (within_bounds_3d(z_tse, y_tse, x_tse, x_shape_[kTwo], x_shape_[kThree], x_shape_[kFour])) { + auto x_index = x_ptr_NC + z_tse * x_stride_[kTwo] + y_tse * x_stride_[kThree] + x_tse * x_stride_[kFour]; + output_addr[output_ptr_NCDHW] += x_addr[x_index] * tse; + } + if (within_bounds_3d(z_bnw, y_bnw, x_bnw, x_shape_[kTwo], x_shape_[kThree], x_shape_[kFour])) { + auto x_index = x_ptr_NC + z_bnw * x_stride_[kTwo] + y_bnw * x_stride_[kThree] + x_bnw * x_stride_[kFour]; + output_addr[output_ptr_NCDHW] += x_addr[x_index] * bnw; + } + if (within_bounds_3d(z_bne, y_bne, x_bne, x_shape_[kTwo], x_shape_[kThree], x_shape_[kFour])) { + auto x_index = x_ptr_NC + z_bne * x_stride_[kTwo] + y_bne * x_stride_[kThree] + x_bne * x_stride_[kFour]; + output_addr[output_ptr_NCDHW] += x_addr[x_index] * bne; + } + if (within_bounds_3d(z_bsw, y_bsw, x_bsw, x_shape_[kTwo], x_shape_[kThree], x_shape_[kFour])) { + auto x_index = x_ptr_NC + z_bsw * x_stride_[kTwo] + y_bsw * x_stride_[kThree] + x_bsw * x_stride_[kFour]; + output_addr[output_ptr_NCDHW] += x_addr[x_index] * bsw; + } + if (within_bounds_3d(z_bse, y_bse, x_bse, x_shape_[kTwo], x_shape_[kThree], x_shape_[kFour])) { + auto x_index = x_ptr_NC + z_bse * x_stride_[kTwo] + y_bse * x_stride_[kThree] + x_bse * x_stride_[kFour]; + output_addr[output_ptr_NCDHW] += x_addr[x_index] * bse; + } + } + } else if (interpolation_mode == "nearest") { + int64_t x_nearest = static_cast(std::round(x)); + int64_t y_nearest = static_cast(std::round(y)); + int64_t z_nearest = static_cast(std::round(z)); + for (size_t c = kZero; c < out_c; c++, x_ptr_NC += x_stride_[kOne], output_ptr_NCDHW += output_stride_[kOne]) { + if (within_bounds_3d(z_nearest, y_nearest, x_nearest, x_shape_[kTwo], x_shape_[kThree], x_shape_[kFour])) { + auto x_index = + x_ptr_NC + z_nearest * x_stride_[kTwo] + y_nearest * x_stride_[kThree] + x_nearest * x_stride_[kFour]; + output_addr[output_ptr_NCDHW] = x_addr[x_index]; + } else { + output_addr[output_ptr_NCDHW] = static_cast(kZero); + } + } + } +} + +template +void GridSampler3DCPUKernel::LaunchKernel(const std::vector &inputs, + const std::vector &outputs) { + auto x_data_addr = reinterpret_cast(inputs[kZero]->addr); + auto grid_data_addr = reinterpret_cast(inputs[kOne]->addr); + auto output_data_addr = reinterpret_cast(outputs[kZero]->addr); + size_t loop_count = output_shape_[kZero] * output_shape_[kTwo] * output_shape_[kThree] * output_shape_[kFour]; + auto task = [&](size_t start, size_t end) { + for (size_t i = start; i < end; i++) { + ComputeTask(x_data_addr, grid_data_addr, output_data_addr, i); + } + }; + if (output_number_ < kDataSizeThreshold) { + task(kZero, loop_count); + } else { + CPUKernelUtils::ParallelFor(task, loop_count); + } +} + +template +T GridSampler3DCPUKernel::grid_sampler_compute_source_index(T coord, int64_t size, std::string padding_mode, + bool align_corners) { + if (align_corners) { + coord = ((coord + 1.f) / kTwo) * (size - kOne); + } else { + coord = ((coord + 1.f) * size - kOne) / kTwo; + } + if (padding_mode == "border") { + coord = std::min(static_cast(size - kOne), std::max(coord, static_cast(kZero))); + } else if (padding_mode == "reflection") { + if (align_corners) { + coord = reflect_coordinates(coord, kZero, kTwo * (size - kOne)); + } else { + coord = reflect_coordinates(coord, -kOne, kTwo * size - kOne); + } + coord = std::min(static_cast(size - kOne), std::max(coord, static_cast(kZero))); + } + return coord; +} + +template +T GridSampler3DCPUKernel::reflect_coordinates(T coord, int64_t twice_low, int64_t twice_high) { + if (twice_low == twice_high) { + return static_cast(kZero); + } + T min = static_cast(twice_low) / kTwo; + T span = static_cast(twice_high - twice_low) / kTwo; + coord = std::fabs(coord - min); + T extra = std::fmod(coord, span); + int64_t flips = static_cast(std::floor(coord / span)); + if (flips % kTwo == kZero) { + return extra + min; + } else { + return span - extra + min; + } +} + +bool GridSampler3DCPUKernel::within_bounds_3d(int64_t d, int64_t h, int64_t w, int64_t D, int64_t H, int64_t W) { + return d >= 0 && d < D && h >= 0 && h < H && w >= 0 && w < W; +} +} // namespace kernel +} // namespace mindspore diff --git a/mindspore/ccsrc/backend/kernel_compiler/cpu/grid_sampler_3d_cpu_kernel.h b/mindspore/ccsrc/backend/kernel_compiler/cpu/grid_sampler_3d_cpu_kernel.h new file mode 100644 index 00000000000..13223670251 --- /dev/null +++ b/mindspore/ccsrc/backend/kernel_compiler/cpu/grid_sampler_3d_cpu_kernel.h @@ -0,0 +1,76 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_GRID_SAMPLER_3D_CPU_KERNEL_H_ +#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_GRID_SAMPLER_3D_CPU_KERNEL_H_ +#include +#include +#include +#include +#include "backend/kernel_compiler/cpu/cpu_kernel.h" +#include "backend/kernel_compiler/cpu/cpu_kernel_factory.h" + +namespace mindspore { +namespace kernel { +class GridSampler3DCPUKernel : public CPUKernel { + public: + GridSampler3DCPUKernel() = default; + ~GridSampler3DCPUKernel() override = default; + + void InitKernel(const CNodePtr &kernel_node) override; + + bool Launch(const std::vector &inputs, const std::vector &workspace, + const std::vector &outputs) override; + + template + void LaunchKernel(const std::vector &inputs, const std::vector &outputs); + + private: + std::vector x_shape_; + std::vector grid_shape_; + std::vector output_shape_; + std::vector x_stride_; + std::vector grid_stride_; + std::vector output_stride_; + std::string interpolation_mode; + std::string padding_mode; + bool align_corners; + size_t output_number_; + TypeId dtype_{kTypeUnknown}; + template + void ComputeTask(T *x_data_addr, T *grid_data_addr, T *output_data_addr, const size_t &seq); + + template + T grid_sampler_compute_source_index(T coord, int64_t size, std::string padding_mode, bool align_corners); + + template + T reflect_coordinates(T coord, int64_t twice_low, int64_t twice_high); + + bool within_bounds_3d(int64_t d, int64_t h, int64_t w, int64_t D, int64_t H, int64_t W); +}; + +MS_REG_CPU_KERNEL( + GridSampler3D, + KernelAttr().AddInputAttr(kNumberTypeFloat32).AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32), + GridSampler3DCPUKernel); + +MS_REG_CPU_KERNEL( + GridSampler3D, + KernelAttr().AddInputAttr(kNumberTypeFloat64).AddInputAttr(kNumberTypeFloat64).AddOutputAttr(kNumberTypeFloat64), + GridSampler3DCPUKernel); +} // namespace kernel +} // namespace mindspore + +#endif // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_GRID_SAMPLER_3D_CPU_KERNEL_H_ diff --git a/mindspore/ccsrc/backend/kernel_compiler/cpu/grid_sampler_3d_grad_cpu_kernel.cc b/mindspore/ccsrc/backend/kernel_compiler/cpu/grid_sampler_3d_grad_cpu_kernel.cc new file mode 100644 index 00000000000..571f697c43c --- /dev/null +++ b/mindspore/ccsrc/backend/kernel_compiler/cpu/grid_sampler_3d_grad_cpu_kernel.cc @@ -0,0 +1,326 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "backend/kernel_compiler/cpu/grid_sampler_3d_grad_cpu_kernel.h" +#include "runtime/device/cpu/cpu_device_address.h" + +namespace { +const size_t kDataSizeThreshold = 64 * 1024; +const size_t kZero = 0; +const size_t kOne = 1; +const size_t kTwo = 2; +const size_t kThree = 3; +const size_t kFour = 4; +} // namespace + +namespace mindspore { +namespace kernel { +void GridSampler3DGradCPUKernel::InitKernel(const CNodePtr &kernel_node) { + grad_shape_ = AnfAlgo::GetInputDeviceShape(kernel_node, kZero); + x_shape_ = AnfAlgo::GetInputDeviceShape(kernel_node, kOne); + grid_shape_ = AnfAlgo::GetInputDeviceShape(kernel_node, kTwo); + dx_shape_ = AnfAlgo::GetOutputDeviceShape(kernel_node, kZero); + dgrid_shape_ = AnfAlgo::GetOutputDeviceShape(kernel_node, kOne); + dx_size_ = dx_shape_[kZero] * dx_shape_[kOne] * dx_shape_[kTwo] * dx_shape_[kThree] * dx_shape_[kFour]; + grid_size_ = grid_shape_[kZero] * grid_shape_[kOne] * grid_shape_[kTwo] * grid_shape_[kThree]; + dtype_ = AnfAlgo::GetInputDeviceDataType(kernel_node, kZero); + size_t stride_tmp = kOne; + auto stride_compute = [&](std::vector &stride, std::vector shape) { + for (size_t i = kFour; i > -kOne; i--) { + stride.insert(stride.begin(), stride_tmp); + stride_tmp *= shape[i]; + } + stride_tmp = kOne; + }; + stride_compute(grad_stride_, grad_shape_); + stride_compute(x_stride_, x_shape_); + stride_compute(grid_stride_, grid_shape_); + stride_compute(dx_stride_, dx_shape_); + stride_compute(dgrid_stride_, dgrid_shape_); + interpolation_mode = AnfAlgo::GetNodeAttr(kernel_node, "interpolation_mode"); + padding_mode = AnfAlgo::GetNodeAttr(kernel_node, "padding_mode"); + align_corners = AnfAlgo::GetNodeAttr(kernel_node, "align_corners"); +} + +bool GridSampler3DGradCPUKernel::Launch(const std::vector &inputs, + const std::vector &, + const std::vector &outputs) { + if (dtype_ == kNumberTypeFloat32) { + LaunchKernel(inputs, outputs); + } else if (dtype_ == kNumberTypeFloat64) { + LaunchKernel(inputs, outputs); + } else { + MS_EXCEPTION(TypeError) << "Input dtype only support float32 and float64!"; + } + return true; +} + +template +void GridSampler3DGradCPUKernel::BilinearKernel(std::vector addr, std::vector location, std::vector mult, + std::vector ptr) { + T x = location[kZero], y = location[kOne], z = location[kTwo]; + int64_t x_tnw = static_cast(std::floor(x)); + int64_t y_tnw = static_cast(std::floor(y)); + int64_t z_tnw = static_cast(std::floor(z)); + int64_t x_tne = x_tnw + kOne, y_tne = y_tnw, z_tne = z_tnw; + int64_t x_tsw = x_tnw, y_tsw = y_tnw + kOne, z_tsw = z_tnw; + int64_t x_tse = x_tnw + kOne, y_tse = y_tnw + kOne, z_tse = z_tnw; + int64_t x_bnw = x_tnw, y_bnw = y_tnw, z_bnw = z_tnw + kOne; + int64_t x_bne = x_tnw + kOne, y_bne = y_tnw, z_bne = z_tnw + kOne; + int64_t x_bsw = x_tnw, y_bsw = y_tnw + kOne, z_bsw = z_tnw + kOne; + int64_t x_bse = x_tnw + kOne, y_bse = y_tnw + kOne, z_bse = z_tnw + kOne; + T tnw = (x_bse - x) * (y_bse - y) * (z_bse - z), tne = (x - x_bsw) * (y_bsw - y) * (z_bsw - z); + T tsw = (x_bne - x) * (y - y_bne) * (z_bne - z), tse = (x - x_bnw) * (y - y_bnw) * (z_bnw - z); + T bnw = (x_tse - x) * (y_tse - y) * (z - z_tse), bne = (x - x_tsw) * (y_tsw - y) * (z - z_tsw); + T bsw = (x_tne - x) * (y - y_tne) * (z - z_tne), bse = (x - x_tnw) * (y - y_tnw) * (z - z_tnw); + T gx = static_cast(kZero), gy = static_cast(kZero), gz = static_cast(kZero); + for (size_t c = kZero; c < x_shape_[kOne]; + c++, ptr[kZero] += grad_stride_[kOne], ptr[kOne] += x_stride_[kOne], ptr[kTwo] += dx_stride_[kOne]) { + T grad_out = addr[kZero][ptr[kZero]]; + safe_add_3d(&addr[kTwo][ptr[kTwo]], z_tnw, y_tnw, x_tnw, dx_stride_[kTwo], dx_stride_[kThree], dx_stride_[kFour], + x_shape_[kTwo], x_shape_[kThree], x_shape_[kFour], tnw * grad_out); + safe_add_3d(&addr[kTwo][ptr[kTwo]], z_tne, y_tne, x_tne, dx_stride_[kTwo], dx_stride_[kThree], dx_stride_[kFour], + x_shape_[kTwo], x_shape_[kThree], x_shape_[kFour], tne * grad_out); + safe_add_3d(&addr[kTwo][ptr[kTwo]], z_tsw, y_tsw, x_tsw, dx_stride_[kTwo], dx_stride_[kThree], dx_stride_[kFour], + x_shape_[kTwo], x_shape_[kThree], x_shape_[kFour], tsw * grad_out); + safe_add_3d(&addr[kTwo][ptr[kTwo]], z_tse, y_tse, x_tse, dx_stride_[kTwo], dx_stride_[kThree], dx_stride_[kFour], + x_shape_[kTwo], x_shape_[kThree], x_shape_[kFour], tse * grad_out); + safe_add_3d(&addr[kTwo][ptr[kTwo]], z_bnw, y_bnw, x_bnw, dx_stride_[kTwo], dx_stride_[kThree], dx_stride_[kFour], + x_shape_[kTwo], x_shape_[kThree], x_shape_[kFour], bnw * grad_out); + safe_add_3d(&addr[kTwo][ptr[kTwo]], z_bne, y_bne, x_bne, dx_stride_[kTwo], dx_stride_[kThree], dx_stride_[kFour], + x_shape_[kTwo], x_shape_[kThree], x_shape_[kFour], bne * grad_out); + safe_add_3d(&addr[kTwo][ptr[kTwo]], z_bsw, y_bsw, x_bsw, dx_stride_[kTwo], dx_stride_[kThree], dx_stride_[kFour], + x_shape_[kTwo], x_shape_[kThree], x_shape_[kFour], bsw * grad_out); + safe_add_3d(&addr[kTwo][ptr[kTwo]], z_bse, y_bse, x_bse, dx_stride_[kTwo], dx_stride_[kThree], dx_stride_[kFour], + x_shape_[kTwo], x_shape_[kThree], x_shape_[kFour], bse * grad_out); + if (within_bounds_3d(z_tnw, y_tnw, x_tnw, x_shape_[kTwo], x_shape_[kThree], x_shape_[kFour])) { + size_t offset = ptr[kOne] + z_tnw * x_stride_[kTwo] + y_tnw * x_stride_[kThree] + x_tnw * x_stride_[kFour]; + T tnw_val = addr[kOne][offset]; + gx -= tnw_val * (y_bse - y) * (z_bse - z) * grad_out; + gy -= tnw_val * (x_bse - x) * (z_bse - z) * grad_out; + gz -= tnw_val * (x_bse - x) * (y_bse - y) * grad_out; + } + if (within_bounds_3d(z_tne, y_tne, x_tne, x_shape_[kTwo], x_shape_[kThree], x_shape_[kFour])) { + size_t offset = ptr[kOne] + z_tne * x_stride_[kTwo] + y_tne * x_stride_[kThree] + x_tne * x_stride_[kFour]; + T tne_val = addr[kOne][offset]; + gx += tne_val * (y_bsw - y) * (z_bsw - z) * grad_out; + gy -= tne_val * (x - x_bsw) * (z_bsw - z) * grad_out; + gz -= tne_val * (x - x_bsw) * (y_bsw - y) * grad_out; + } + if (within_bounds_3d(z_tsw, y_tsw, x_tsw, x_shape_[kTwo], x_shape_[kThree], x_shape_[kFour])) { + size_t offset = ptr[kOne] + z_tsw * x_stride_[kTwo] + y_tsw * x_stride_[kThree] + x_tsw * x_stride_[kFour]; + T tsw_val = addr[kOne][offset]; + gx -= tsw_val * (y - y_bne) * (z_bne - z) * grad_out; + gy += tsw_val * (x_bne - x) * (z_bne - z) * grad_out; + gz -= tsw_val * (x_bne - x) * (y - y_bne) * grad_out; + } + if (within_bounds_3d(z_tse, y_tse, x_tse, x_shape_[kTwo], x_shape_[kThree], x_shape_[kFour])) { + size_t offset = ptr[kOne] + z_tse * x_stride_[kTwo] + y_tse * x_stride_[kThree] + x_tse * x_stride_[kFour]; + T tse_val = addr[kOne][offset]; + gx += tse_val * (y - y_bnw) * (z_bnw - z) * grad_out; + gy += tse_val * (x - x_bnw) * (z_bnw - z) * grad_out; + gz -= tse_val * (x - x_bnw) * (y - y_bnw) * grad_out; + } + if (within_bounds_3d(z_bnw, y_bnw, x_bnw, x_shape_[kTwo], x_shape_[kThree], x_shape_[kFour])) { + size_t offset = ptr[kOne] + z_bnw * x_stride_[kTwo] + y_bnw * x_stride_[kThree] + x_bnw * x_stride_[kFour]; + T bnw_val = addr[kOne][offset]; + gx -= bnw_val * (y_tse - y) * (z - z_tse) * grad_out; + gy -= bnw_val * (x_tse - x) * (z - z_tse) * grad_out; + gz += bnw_val * (x_tse - x) * (y_tse - y) * grad_out; + } + if (within_bounds_3d(z_bne, y_bne, x_bne, x_shape_[kTwo], x_shape_[kThree], x_shape_[kFour])) { + size_t offset = ptr[kOne] + z_bne * x_stride_[kTwo] + y_bne * x_stride_[kThree] + x_bne * x_stride_[kFour]; + T bne_val = addr[kOne][offset]; + gx += bne_val * (y_tsw - y) * (z - z_tsw) * grad_out; + gy -= bne_val * (x - x_tsw) * (z - z_tsw) * grad_out; + gz += bne_val * (x - x_tsw) * (y_tsw - y) * grad_out; + } + if (within_bounds_3d(z_bsw, y_bsw, x_bsw, x_shape_[kTwo], x_shape_[kThree], x_shape_[kFour])) { + size_t offset = ptr[kOne] + z_bsw * x_stride_[kTwo] + y_bsw * x_stride_[kThree] + x_bsw * x_stride_[kFour]; + T bsw_val = addr[kOne][offset]; + gx -= bsw_val * (y - y_tne) * (z - z_tne) * grad_out; + gy += bsw_val * (x_tne - x) * (z - z_tne) * grad_out; + gz += bsw_val * (x_tne - x) * (y - y_tne) * grad_out; + } + if (within_bounds_3d(z_bse, y_bse, x_bse, x_shape_[kTwo], x_shape_[kThree], x_shape_[kFour])) { + size_t offset = ptr[kOne] + z_bse * x_stride_[kTwo] + y_bse * x_stride_[kThree] + x_bse * x_stride_[kFour]; + T bse_val = addr[kOne][offset]; + gx += bse_val * (y - y_tnw) * (z - z_tnw) * grad_out; + gy += bse_val * (x - x_tnw) * (z - z_tnw) * grad_out; + gz += bse_val * (x - x_tnw) * (y - y_tnw) * grad_out; + } + } + addr[kThree][ptr[kThree]] = mult[kZero] * gx; + addr[kThree][ptr[kThree] + kOne] = mult[kOne] * gy; + addr[kThree][ptr[kThree] + kTwo] = mult[kTwo] * gz; +} + +template +void GridSampler3DGradCPUKernel::ComputeTask(T *grad_addr, T *x_addr, T *grid_addr, T *dx_addr, T *dgrid_addr, + const size_t &n) { + size_t grid_ptr_N = n * grid_stride_[kZero]; + size_t dgrid_ptr_NDHW = n * dgrid_stride_[kZero]; + for (size_t d = kZero; d < grid_shape_[kOne]; d++) { + for (size_t h = kZero; h < grid_shape_[kTwo]; h++) { + for (size_t w = kZero; w < grid_shape_[kThree]; w++, dgrid_ptr_NDHW += dgrid_stride_[kThree]) { + size_t grid_ptr_NDHW = grid_ptr_N + d * grid_stride_[kOne] + h * grid_stride_[kTwo] + w * grid_stride_[kThree]; + T x = grid_addr[grid_ptr_NDHW]; + T y = grid_addr[grid_ptr_NDHW + grid_stride_[kFour]]; + T z = grid_addr[grid_ptr_NDHW + kTwo * grid_stride_[kFour]]; + T gx_mult, gy_mult, gz_mult; + x = grid_sampler_compute_source_index_set_grad(x, x_shape_[kFour], padding_mode, align_corners, &gx_mult); + y = grid_sampler_compute_source_index_set_grad(y, x_shape_[kThree], padding_mode, align_corners, &gy_mult); + z = grid_sampler_compute_source_index_set_grad(z, x_shape_[kTwo], padding_mode, align_corners, &gz_mult); + if (interpolation_mode == "bilinear") { + size_t grad_ptr_NCDHW = + n * grad_stride_[kZero] + d * grad_stride_[kTwo] + h * grad_stride_[kThree] + w * grad_stride_[kFour]; + size_t dx_ptr_NC = n * dx_stride_[kZero], x_ptr_NC = n * x_stride_[kZero]; + std::vector addr = {grad_addr, x_addr, dx_addr, dgrid_addr}; + std::vector location = {x, y, z}; + std::vector mult = {gx_mult, gy_mult, gz_mult}; + std::vector ptr = {grad_ptr_NCDHW, x_ptr_NC, dx_ptr_NC, dgrid_ptr_NDHW}; + BilinearKernel(addr, location, mult, ptr); + } else if (interpolation_mode == "nearest") { + int64_t x_nearest = static_cast(std::round(x)); + int64_t y_nearest = static_cast(std::round(y)); + int64_t z_nearest = static_cast(std::round(z)); + size_t grad_ptr_NCDHW = + n * grad_stride_[kZero] + d * grad_stride_[kTwo] + h * grad_stride_[kThree] + w * grad_stride_[kFour]; + size_t dx_ptr_NC = n * dx_stride_[kZero]; + for (size_t c = kZero; c < x_shape_[kOne]; + c++, grad_ptr_NCDHW += grad_stride_[kOne], dx_ptr_NC += dx_stride_[kOne]) { + safe_add_3d(&dx_addr[dx_ptr_NC], z_nearest, y_nearest, x_nearest, dx_stride_[kTwo], dx_stride_[kThree], + dx_stride_[kFour], x_shape_[kTwo], x_shape_[kThree], x_shape_[kFour], + grad_addr[grad_ptr_NCDHW]); + } + dgrid_addr[dgrid_ptr_NDHW] = static_cast(kZero); + dgrid_addr[dgrid_ptr_NDHW + kOne] = static_cast(kZero); + dgrid_addr[dgrid_ptr_NDHW + kTwo] = static_cast(kZero); + } + } + } + } +} + +template +void GridSampler3DGradCPUKernel::LaunchKernel(const std::vector &inputs, + const std::vector &outputs) { + auto grad_data_addr = reinterpret_cast(inputs[kZero]->addr); + auto x_data_addr = reinterpret_cast(inputs[kOne]->addr); + auto grid_data_addr = reinterpret_cast(inputs[kTwo]->addr); + auto dx_data_addr = reinterpret_cast(outputs[kZero]->addr); + auto dgrid_data_addr = reinterpret_cast(outputs[kOne]->addr); + size_t loop_count = x_shape_[kZero]; + for (size_t i = kZero; i < dx_size_; i++) { + dx_data_addr[i] = static_cast(kZero); + } + auto task = [&](size_t start, size_t end) { + for (size_t n = start; n < end; n++) { + ComputeTask(grad_data_addr, x_data_addr, grid_data_addr, dx_data_addr, dgrid_data_addr, n); + } + }; + if (grid_size_ * sizeof(T) < kDataSizeThreshold) { + task(kZero, loop_count); + } else { + CPUKernelUtils::ParallelFor(task, loop_count); + } +} + +template +T GridSampler3DGradCPUKernel::grid_sampler_compute_source_index_set_grad(T coord, size_t size, std::string padding_mode, + bool align_corners, T *grad_x) { + T grad_clip, grad_refl; + if (align_corners) { + *grad_x = static_cast(size - kOne) / kTwo; + coord = ((coord + kOne) / kTwo) * (size - kOne); + } else { + *grad_x = static_cast(size) / kTwo; + coord = ((coord + kOne) * size - kOne) / kTwo; + } + if (padding_mode == "border") { + coord = clip_coordinates_set_grad(coord, size, &grad_clip); + *grad_x = (*grad_x) * grad_clip; + } else if (padding_mode == "reflection") { + if (align_corners) { + coord = reflect_coordinates_set_grad(coord, kZero, kTwo * (size - kOne), &grad_refl); + } else { + coord = reflect_coordinates_set_grad(coord, -kOne, kTwo * size - kOne, &grad_refl); + } + coord = clip_coordinates_set_grad(coord, size, &grad_clip); + *grad_x = (*grad_x) * grad_refl * grad_clip; + } + return coord; +} + +template +T GridSampler3DGradCPUKernel::clip_coordinates_set_grad(T x, int64_t clip_limit, T *grad_x) { + if (x <= static_cast(kZero)) { + *grad_x = static_cast(kZero); + return static_cast(kZero); + } else { + T max = static_cast(clip_limit - kOne); + if (x >= max) { + *grad_x = static_cast(kZero); + return max; + } else { + *grad_x = static_cast(kOne); + return x; + } + } +} + +template +T GridSampler3DGradCPUKernel::reflect_coordinates_set_grad(T x, int64_t twice_low, int64_t twice_high, T *grad_x) { + if (twice_low == twice_high) { + *grad_x = static_cast(kZero); + return static_cast(kZero); + } + int64_t grad_x_mult_; + T min = static_cast(twice_low) / kTwo; + T span = static_cast(twice_high - twice_low) / kTwo; + x = x - min; + if (x < static_cast(kZero)) { + grad_x_mult_ = -kOne; + x = -x; + } else { + grad_x_mult_ = kOne; + } + T extra = std::fmod(x, span); + int64_t flips = static_cast(std::floor(x / span)); + if (flips % kTwo == kZero) { + *grad_x = static_cast(grad_x_mult_); + return extra + min; + } else { + *grad_x = static_cast(-grad_x_mult_); + return span - extra + min; + } +} + +template +void GridSampler3DGradCPUKernel::safe_add_3d(T *data, int64_t d, int64_t h, int64_t w, size_t sD, size_t sH, size_t sW, + size_t D, size_t H, size_t W, T delta) { + if (within_bounds_3d(d, h, w, D, H, W)) { + data[d * sD + h * sH + w * sW] += delta; + } +} + +bool GridSampler3DGradCPUKernel::within_bounds_3d(int64_t d, int64_t h, int64_t w, size_t D, size_t H, size_t W) { + int64_t iD = static_cast(D); + int64_t iH = static_cast(H); + int64_t iW = static_cast(W); + return d >= 0 && d < iD && h >= 0 && h < iH && w >= 0 && w < iW; +} +} // namespace kernel +} // namespace mindspore diff --git a/mindspore/ccsrc/backend/kernel_compiler/cpu/grid_sampler_3d_grad_cpu_kernel.h b/mindspore/ccsrc/backend/kernel_compiler/cpu/grid_sampler_3d_grad_cpu_kernel.h new file mode 100644 index 00000000000..0ece5b77409 --- /dev/null +++ b/mindspore/ccsrc/backend/kernel_compiler/cpu/grid_sampler_3d_grad_cpu_kernel.h @@ -0,0 +1,99 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_GRID_SAMPLER_3D_GRAD_CPU_KERNEL_H_ +#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_GRID_SAMPLER_3D_GRAD_CPU_KERNEL_H_ +#include +#include +#include +#include "backend/kernel_compiler/cpu/cpu_kernel.h" +#include "backend/kernel_compiler/cpu/cpu_kernel_factory.h" + +namespace mindspore { +namespace kernel { +class GridSampler3DGradCPUKernel : public CPUKernel { + public: + GridSampler3DGradCPUKernel() = default; + ~GridSampler3DGradCPUKernel() override = default; + + void InitKernel(const CNodePtr &kernel_node) override; + + bool Launch(const std::vector &inputs, const std::vector &workspace, + const std::vector &outputs) override; + + template + void LaunchKernel(const std::vector &inputs, const std::vector &outputs); + + private: + std::vector grad_shape_; + std::vector x_shape_; + std::vector grid_shape_; + std::vector dx_shape_; + std::vector dgrid_shape_; + std::vector grad_stride_; + std::vector x_stride_; + std::vector grid_stride_; + std::vector dx_stride_; + std::vector dgrid_stride_; + std::string interpolation_mode; + std::string padding_mode; + bool align_corners; + size_t dx_size_; + size_t grid_size_; + TypeId dtype_{kTypeUnknown}; + template + void BilinearKernel(std::vector addr, std::vector location, std::vector mult, std::vector ptr); + + template + void ComputeTask(T *grad_addr, T *x_addr, T *grid_addr, T *dx_addr, T *dgrid_addr, const size_t &n); + + template + T grid_sampler_compute_source_index_set_grad(T coord, size_t size, std::string padding_mode, bool align_corners, + T *grad_x); + + template + T reflect_coordinates_set_grad(T x, int64_t twice_low, int64_t twice_high, T *grad_x); + + template + T clip_coordinates_set_grad(T x, int64_t clip_limit, T *grad_x); + + template + void safe_add_3d(T *data, int64_t d, int64_t h, int64_t w, size_t sD, size_t sH, size_t sW, size_t D, size_t H, + size_t W, T delta); + + bool within_bounds_3d(int64_t d, int64_t h, int64_t w, size_t D, size_t H, size_t W); +}; + +MS_REG_CPU_KERNEL(GridSampler3DGrad, + KernelAttr() + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddOutputAttr(kNumberTypeFloat32) + .AddOutputAttr(kNumberTypeFloat32), + GridSampler3DGradCPUKernel); + +MS_REG_CPU_KERNEL(GridSampler3DGrad, + KernelAttr() + .AddInputAttr(kNumberTypeFloat64) + .AddInputAttr(kNumberTypeFloat64) + .AddInputAttr(kNumberTypeFloat64) + .AddOutputAttr(kNumberTypeFloat64) + .AddOutputAttr(kNumberTypeFloat64), + GridSampler3DGradCPUKernel); +} // namespace kernel +} // namespace mindspore + +#endif // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_GRID_SAMPLER_3D_GRAD_CPU_KERNEL_H_ diff --git a/mindspore/core/base/core_ops.h b/mindspore/core/base/core_ops.h index ef09f3151ba..b21ab7dbf8e 100644 --- a/mindspore/core/base/core_ops.h +++ b/mindspore/core/base/core_ops.h @@ -128,6 +128,8 @@ constexpr auto kConv2DTranspose = "Conv2DTranspose"; constexpr auto kSparseApplyAdadelta = "SparseApplyAdadelta"; constexpr auto kRoll = "Roll"; constexpr auto kTanh = "Tanh"; +constexpr auto kGridSampler3D = "GridSampler3D"; +constexpr auto kGridSampler3DGrad = "GridSampler3DGrad"; // Others constexpr auto kEnvironCreate = "EnvironCreate"; @@ -485,6 +487,8 @@ inline const PrimitivePtr kPrimApplyAddSign = std::make_shared("Apply inline const PrimitivePtr kPrimApplyAdagrad = std::make_shared("ApplyAdagrad"); inline const PrimitivePtr kPrimApplyAdadelta = std::make_shared("ApplyAdadelta"); inline const PrimitivePtr kPrimApplyAdamWithAmsgrad = std::make_shared("ApplyAdamWithAmsgrad"); +inline const PrimitivePtr kPrimGridSampler3D = std::make_shared(kGridSampler3D); +inline const PrimitivePtr kPrimGridSampler3DGrad = std::make_shared(kGridSampler3DGrad); // Comm ops inline const PrimitivePtr kPrimMirror = std::make_shared("_MirrorOperator"); diff --git a/mindspore/core/ops/grad/grid_sampler_3d_grad.cc b/mindspore/core/ops/grad/grid_sampler_3d_grad.cc new file mode 100644 index 00000000000..5475dfd63ad --- /dev/null +++ b/mindspore/core/ops/grad/grid_sampler_3d_grad.cc @@ -0,0 +1,102 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include "ops/grad/grid_sampler_3d_grad.h" + +namespace mindspore { +namespace ops { +namespace { +const size_t kZero = 0; +const size_t kOne = 1; +const size_t kTwo = 2; +const size_t kThree = 3; +const size_t kFour = 4; +const size_t kFive = 5; + +abstract::TupleShapePtr GridSampler3DGradInferShape(const PrimitivePtr &primitive, + const std::vector &input_args) { + auto grad_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[kZero]->BuildShape())[kShape]; + auto input_x_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[kOne]->BuildShape())[kShape]; + auto grid_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[kTwo]->BuildShape())[kShape]; + if (grad_shape.size() != kFive) { + MS_EXCEPTION(ValueError) << "Grad must be a 5-dimensional tensor, but got " << std::to_string(grad_shape.size()) + << "-dimensional tensor."; + } + if (input_x_shape.size() != kFive) { + MS_EXCEPTION(ValueError) << "Input_x must be a 5-dimensional tensor, but got " + << std::to_string(input_x_shape.size()) << "-dimensional tensor."; + } + if (grid_shape.size() != kFive) { + MS_EXCEPTION(ValueError) << "Grid must be a 5-dimensional tensor, but got " << std::to_string(grid_shape.size()) + << "-dimensional tensor."; + } + if (input_x_shape[kZero] != grid_shape[kZero]) { + MS_EXCEPTION(ValueError) << "The shape of grid is " << input_args[kTwo]->BuildShape()->ToString() + << " , but the shape of input_x is " << input_args[kOne]->BuildShape()->ToString() + << " . The first dimension of grid and input_x must be equal."; + } + if (grid_shape[kFour] != kThree) { + MS_EXCEPTION(ValueError) << "The last dimension of grid must be 3, but got " << std::to_string(grid_shape[kFour]); + } + std::vector out_shape = {input_x_shape[kZero], input_x_shape[kOne], grid_shape[kOne], grid_shape[kTwo], + grid_shape[kThree]}; + bool shape_error = false; + for (size_t i = kZero; i < kFive; i++) { + if (out_shape[i] != grad_shape[i]) { + shape_error = true; + break; + } + } + if (shape_error) { + MS_EXCEPTION(ValueError) << "The shape of grad, which is the same as that of output, is " + << input_args[kZero]->BuildShape()->ToString() << ", but the shape of output is (" + << std::to_string(out_shape[kZero]) << ", " << std::to_string(out_shape[kOne]) << ", " + << std::to_string(out_shape[kTwo]) << ", " << std::to_string(out_shape[kThree]) << ", " + << std::to_string(out_shape[kFour]) << ")."; + } + abstract::ShapePtr dx_shape = std::make_shared(input_x_shape); + abstract::ShapePtr dgrid_shape = std::make_shared(grid_shape); + return std::make_shared(std::vector{dx_shape, dgrid_shape}); +} + +TuplePtr GridSampler3DGradInferType(const PrimitivePtr &primitive, const std::vector &input_args) { + std::map types; + std::set valid_types = {kFloat32, kFloat64}; + TypePtr grad_type = input_args[kZero]->BuildType(); + TypePtr input_x_type = input_args[kOne]->BuildType(); + TypePtr grid_type = input_args[kTwo]->BuildType(); + (void)types.emplace("grad", grad_type); + (void)types.emplace("input_x", input_x_type); + (void)types.emplace("grid", grid_type); + (void)CheckAndConvertUtils::CheckTensorTypeSame(types, valid_types, primitive->name()); + return std::make_shared(std::vector{input_x_type, grid_type}); +} +} // namespace + +AbstractBasePtr GridSampler3DGradInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive, + const std::vector &input_args) { + MS_EXCEPTION_IF_NULL(primitive); + const int64_t input_num = kThree; + CheckAndConvertUtils::CheckInputArgs(input_args, kEqual, input_num, primitive->name()); + auto infer_types = GridSampler3DGradInferType(primitive, input_args); + auto infer_shapes = GridSampler3DGradInferShape(primitive, input_args); + return abstract::MakeAbstract(infer_shapes, infer_types); +} + +REGISTER_PRIMITIVE_EVAL_IMPL(GridSampler3DGrad, prim::kPrimGridSampler3DGrad, GridSampler3DGradInfer, nullptr, true); +} // namespace ops +} // namespace mindspore diff --git a/mindspore/core/ops/grad/grid_sampler_3d_grad.h b/mindspore/core/ops/grad/grid_sampler_3d_grad.h new file mode 100644 index 00000000000..c39515810ff --- /dev/null +++ b/mindspore/core/ops/grad/grid_sampler_3d_grad.h @@ -0,0 +1,44 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef MINDSPORE_CORE_OPS_GRID_SAMPLER_3D_H_ +#define MINDSPORE_CORE_OPS_GRID_SAMPLER_3D_H_ +#include +#include +#include +#include + +#include "abstract/abstract_value.h" +#include "ops/op_utils.h" +#include "ops/primitive_c.h" +#include "utils/check_convert_utils.h" + +namespace mindspore { +namespace ops { +constexpr auto kNameGridSampler3DGrad = "GridSampler3DGrad"; +class GridSampler3DGrad : public PrimitiveC { + public: + GridSampler3DGrad() : PrimitiveC(kNameGridSampler3DGrad) { InitIOName({"grad", "input_x", "grid"}, {"dx", "dgrid"}); } + ~GridSampler3DGrad() = default; + MS_DECLARE_PARENT(GridSampler3DGrad, PrimitiveC); +}; +AbstractBasePtr GridSampler3DGradInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive, + const std::vector &input_args); +using PrimGridSampler3DGrad = std::shared_ptr; +} // namespace ops +} // namespace mindspore + +#endif // MINDSPORE_CORE_OPS_GRID_SAMPLER_3D_H_ diff --git a/mindspore/core/ops/grid_sampler_3d.cc b/mindspore/core/ops/grid_sampler_3d.cc new file mode 100644 index 00000000000..5055ab48976 --- /dev/null +++ b/mindspore/core/ops/grid_sampler_3d.cc @@ -0,0 +1,79 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include "ops/grid_sampler_3d.h" + +namespace mindspore { +namespace ops { +namespace { +const size_t kZero = 0; +const size_t kOne = 1; +const size_t kTwo = 2; +const size_t kThree = 3; +const size_t kFour = 4; +const size_t kFive = 5; + +abstract::ShapePtr GridSampler3DInferShape(const PrimitivePtr &primitive, + const std::vector &input_args) { + auto input_x_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[kZero]->BuildShape())[kShape]; + auto grid_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[kOne]->BuildShape())[kShape]; + if (input_x_shape.size() != kFive) { + MS_EXCEPTION(ValueError) << "Input_x must be a 5-dimensional tensor, but got " + << std::to_string(input_x_shape.size()) << "-dimensional tensor."; + } + if (grid_shape.size() != kFive) { + MS_EXCEPTION(ValueError) << "Grid must be a 5-dimensional tensor, but got " << std::to_string(grid_shape.size()) + << "-dimensional tensor."; + } + if (input_x_shape[kZero] != grid_shape[kZero]) { + MS_EXCEPTION(ValueError) << "The shape of grid is " << input_args[kOne]->BuildShape()->ToString() + << " , but the shape of input_x is " << input_args[kZero]->BuildShape()->ToString() + << " . The first dimension of grid and input_x must be equal."; + } + if (grid_shape[kFour] != kThree) { + MS_EXCEPTION(ValueError) << "The last dimension of grid must be 3, but got " << std::to_string(grid_shape[kFour]); + } + std::vector output_shape = {input_x_shape[kZero], input_x_shape[kOne], grid_shape[kOne], grid_shape[kTwo], + grid_shape[kThree]}; + return std::make_shared(output_shape); +} + +TypePtr GridSampler3DInferType(const PrimitivePtr &primitive, const std::vector &input_args) { + std::map types; + std::set valid_types = {kFloat32, kFloat64}; + TypePtr input_x_type = input_args[kZero]->BuildType(); + TypePtr grid_type = input_args[kOne]->BuildType(); + (void)types.emplace("input_x", input_x_type); + (void)types.emplace("grid", grid_type); + (void)CheckAndConvertUtils::CheckTensorTypeSame(types, valid_types, primitive->name()); + return input_x_type; +} +} // namespace + +AbstractBasePtr GridSampler3DInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive, + const std::vector &input_args) { + MS_EXCEPTION_IF_NULL(primitive); + const int64_t input_num = kTwo; + CheckAndConvertUtils::CheckInputArgs(input_args, kEqual, input_num, primitive->name()); + auto infer_type = GridSampler3DInferType(primitive, input_args); + auto infer_shape = GridSampler3DInferShape(primitive, input_args); + return abstract::MakeAbstract(infer_shape, infer_type); +} + +REGISTER_PRIMITIVE_EVAL_IMPL(GridSampler3D, prim::kPrimGridSampler3D, GridSampler3DInfer, nullptr, true); +} // namespace ops +} // namespace mindspore diff --git a/mindspore/core/ops/grid_sampler_3d.h b/mindspore/core/ops/grid_sampler_3d.h new file mode 100644 index 00000000000..316fac206c5 --- /dev/null +++ b/mindspore/core/ops/grid_sampler_3d.h @@ -0,0 +1,44 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef MINDSPORE_CORE_OPS_GRID_SAMPLER_3D_H_ +#define MINDSPORE_CORE_OPS_GRID_SAMPLER_3D_H_ +#include +#include +#include +#include + +#include "abstract/abstract_value.h" +#include "ops/op_utils.h" +#include "ops/primitive_c.h" +#include "utils/check_convert_utils.h" + +namespace mindspore { +namespace ops { +constexpr auto kNameGridSampler3D = "GridSampler3D"; +class GridSampler3D : public PrimitiveC { + public: + GridSampler3D() : PrimitiveC(kNameGridSampler3D) { InitIOName({"input_x", "grid"}, {"output"}); } + ~GridSampler3D() = default; + MS_DECLARE_PARENT(GridSampler3D, PrimitiveC); +}; +AbstractBasePtr GridSampler3DInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive, + const std::vector &input_args); +using PrimGridSampler3D = std::shared_ptr; +} // namespace ops +} // namespace mindspore + +#endif // MINDSPORE_CORE_OPS_GRID_SAMPLER_3D_H_ diff --git a/mindspore/python/mindspore/ops/_grad_experimental/grad_nn_ops.py b/mindspore/python/mindspore/ops/_grad_experimental/grad_nn_ops.py index 46f08ad2bfd..e9da366ff2a 100644 --- a/mindspore/python/mindspore/ops/_grad_experimental/grad_nn_ops.py +++ b/mindspore/python/mindspore/ops/_grad_experimental/grad_nn_ops.py @@ -84,3 +84,15 @@ def get_bprop_celu(self): return (dx,) return bprop + + +@bprop_getters.register(P.GridSampler3D) +def get_bprop_grid_sampler_3d(self): + """Grad definition for `GridSampler3D` operation.""" + grad = G.GridSampler3DGrad(self.interpolation_mode, self.padding_mode, self.align_corners) + + def bprop(input_x, grid, out, dout): + dx, dgrid = grad(dout, input_x, grid) + return dx, dgrid + + return bprop diff --git a/mindspore/python/mindspore/ops/_op_impl/aicpu/__init__.py b/mindspore/python/mindspore/ops/_op_impl/aicpu/__init__.py index f56c27120cb..718b7840342 100644 --- a/mindspore/python/mindspore/ops/_op_impl/aicpu/__init__.py +++ b/mindspore/python/mindspore/ops/_op_impl/aicpu/__init__.py @@ -95,3 +95,5 @@ from .scatter_elements import _scatter_elements_aicpu from .non_max_suppression import _non_max_suppression_aicpu from .square import _square_aicpu from .lower_bound import _lower_bound_aicpu +from .grid_sampler_3d import _grid_sampler_3d_aicpu +from .grid_sampler_3d_grad import _grid_sampler_3d_grad_aicpu diff --git a/mindspore/python/mindspore/ops/_op_impl/aicpu/grid_sampler_3d.py b/mindspore/python/mindspore/ops/_op_impl/aicpu/grid_sampler_3d.py new file mode 100644 index 00000000000..cbea7edf530 --- /dev/null +++ b/mindspore/python/mindspore/ops/_op_impl/aicpu/grid_sampler_3d.py @@ -0,0 +1,34 @@ +# Copyright 2021 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ + +"""GridSampler3D op""" +from mindspore.ops.op_info_register import op_info_register, AiCPURegOp, DataType + +grid_sampler_3d_op_info = AiCPURegOp("GridSampler3D") \ + .fusion_type("OPAQUE") \ + .attr("interpolation_mode", "str") \ + .attr("padding_mode", "str") \ + .attr("align_corners", "bool") \ + .input(0, "input_x", "required") \ + .input(1, "grid", "required") \ + .output(0, "output", "required") \ + .dtype_format(DataType.F32_Default, DataType.F32_Default, DataType.F32_Default) \ + .dtype_format(DataType.F64_Default, DataType.F64_Default, DataType.F64_Default) \ + .get_op_info() + +@op_info_register(grid_sampler_3d_op_info) +def _grid_sampler_3d_aicpu(): + """GridSampler3D aicpu register""" + return diff --git a/mindspore/python/mindspore/ops/_op_impl/aicpu/grid_sampler_3d_grad.py b/mindspore/python/mindspore/ops/_op_impl/aicpu/grid_sampler_3d_grad.py new file mode 100644 index 00000000000..303ec036ac2 --- /dev/null +++ b/mindspore/python/mindspore/ops/_op_impl/aicpu/grid_sampler_3d_grad.py @@ -0,0 +1,38 @@ +# Copyright 2021 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ + +"""GridSampler3DGrad op""" +from mindspore.ops.op_info_register import op_info_register, AiCPURegOp, DataType + +grid_sampler_3d_grad_op_info = AiCPURegOp("GridSampler3DGrad") \ + .fusion_type("OPAQUE") \ + .attr("interpolation_mode", "str") \ + .attr("padding_mode", "str") \ + .attr("align_corners", "bool") \ + .input(0, "grad", "required") \ + .input(1, "input_x", "required") \ + .input(2, "grid", "required") \ + .output(0, "dx", "required") \ + .output(1, "dgrid", "required") \ + .dtype_format(DataType.F32_Default, DataType.F32_Default, DataType.F32_Default, \ + DataType.F32_Default, DataType.F32_Default) \ + .dtype_format(DataType.F64_Default, DataType.F64_Default, DataType.F64_Default, \ + DataType.F64_Default, DataType.F64_Default) \ + .get_op_info() + +@op_info_register(grid_sampler_3d_grad_op_info) +def _grid_sampler_3d_grad_aicpu(): + """GridSampler3DGrad aicpu register""" + return diff --git a/mindspore/python/mindspore/ops/operations/__init__.py b/mindspore/python/mindspore/ops/operations/__init__.py index 53cc9fe44e4..8f2d8b96097 100644 --- a/mindspore/python/mindspore/ops/operations/__init__.py +++ b/mindspore/python/mindspore/ops/operations/__init__.py @@ -90,7 +90,7 @@ from .nn_ops import (LSTM, SGD, Adam, AdamWeightDecay, FusedSparseAdam, FusedSpa ApplyAdaMax, ApplyAdadelta, ApplyAdagrad, ApplyAdagradV2, ApplyAdagradDA, ApplyAddSign, ApplyPowerSign, ApplyGradientDescent, ApplyProximalGradientDescent, ApplyRMSProp, ApplyCenteredRMSProp, BasicLSTMCell, InTopK, AdaptiveAvgPool2D, SoftShrink, - ApplyAdamWithAmsgrad) + ApplyAdamWithAmsgrad, GridSampler3D) from . import _quant_ops from ._quant_ops import * from .other_ops import (Assign, IOU, BoundingBoxDecode, BoundingBoxEncode, @@ -129,6 +129,7 @@ from .custom_ops import (Custom) __all__ = [ 'CeLU', 'Ger', + 'GridSampler3D', 'Unique', 'ReverseSequence', 'Sort', diff --git a/mindspore/python/mindspore/ops/operations/_grad_ops.py b/mindspore/python/mindspore/ops/operations/_grad_ops.py index bf59d7f5b62..4a5bd2e6f79 100644 --- a/mindspore/python/mindspore/ops/operations/_grad_ops.py +++ b/mindspore/python/mindspore/ops/operations/_grad_ops.py @@ -2346,3 +2346,52 @@ class ParallelResizeBilinearGrad(PrimitiveWithInfer): return {'shape': output_shape, 'dtype': x_dtype, 'value': None} + + +class GridSampler3DGrad(Primitive): + """ + Computes gradients for GridSampler3D operation. + + Args: + interpolation_mode (str): An optional string specifying the interpolation method. The optional values are + "bilinear" or "nearest". Default: "bilinear". + padding_mode (str): An optional string specifying the pad method. The optional values are "zeros", "border" or + "reflection". Default: "zeros". + align_corners (bool): An optional bool. If "true", the centers of the corner pixels of the input and output + tensors are aligned. Defaults to "false". + + Inputs: + - **grad** (Tensor) - A 5-D tensor whose dtype is float32 or float64 and whose shape is :math:`(N, C, D_{out}, + H_{out}, W_{out})`. The shape is inconsistent with the shape of the output result of forward calculation. + - **input_x** (Tensor) - A 5-D tensor whose dtype is the same as `grad` and whose shape is :math:`(N, C, + D_{in}, H_{in}, W_{in})`. + - **grid** (Tensor) - A 5-D tensor whose dtype is the same as `grad` and whose shape is :math:`(N, D_{out}, + H_{out}, W_{out}, 3)`. + + Outputs: + - **dx** (Tensor) - A 5-D tensor whose dtype and shape are the same as `input_x`. + - **dgrid** (Tensor) - A 5-D tensor whose dtype and shape are the same as `grid`. + + Raises: + TypeError: If `grad`, `input_x` or `grid` is not a Tensor. + TypeError: If the dtypes of `grad`, `input_x` and `grid` are inconsistent. + TypeError: If the dtype of `grad`, `input_x` or `grid` is not a valid type. + TypeError: If `align_corners` is not a boolean value. + ValueError: If the rank of `grad`, `input_x` or `grid` is not equal to 5. + ValueError: If the first dimension of `grad`, `input_x` and `grid` are inconsistent. + ValueError: If the last dimension of `grid` is not equal to 3. + ValueError: If `interpolation_mode` is not "bilinear", "nearest" or a string value. + ValueError: If `padding_mode` is not "zeros", "border", "reflection" or a string value. + ValueError: If the shape of `grad` is inconsistent with the shape of the output result of forward calculation. + + Supported Platforms: + ``CPU`` + """ + + @prim_attr_register + def __init__(self, interpolation_mode='bilinear', padding_mode='zeros', align_corners=False): + """Initialize GridSampler3DGrad.""" + validator.check_string(interpolation_mode, ['bilinear', 'nearest'], 'interpolation_mode', self.name) + validator.check_string(padding_mode, ['zeros', 'border', 'reflection'], 'padding_mode', self.name) + validator.check_bool(align_corners, 'align_corners', self.name) + self.init_prim_io_names(inputs=['grad', 'input_x', 'grid'], outputs=['dx', 'dgrid']) diff --git a/mindspore/python/mindspore/ops/operations/nn_ops.py b/mindspore/python/mindspore/ops/operations/nn_ops.py index 99f23fde512..28e11ff9e6e 100644 --- a/mindspore/python/mindspore/ops/operations/nn_ops.py +++ b/mindspore/python/mindspore/ops/operations/nn_ops.py @@ -8982,3 +8982,82 @@ class ApplyAdamWithAmsgrad(Primitive): validator.check_value_type("beta2", beta2, [float], self.name) validator.check_value_type("epsilon", epsilon, [float], self.name) validator.check_value_type("use_locking", use_locking, [bool], self.name) + + +class GridSampler3D(Primitive): + """ + Given an `input_x` and a flow-field `grid`, computes the `output` using `input_x` values and pixel locations from + `grid`. Only volumetric (5-D) `input_x` is supported. + + For `input_x` with shape :math:`(N, C, D_{in}, H_{in}, W_{in})` and `grid` with shape :math:`(N, D_{out}, H_{out}, + W_{out}, 3)`, the `output` will have shape :math:`(N, C, D_{out}, H_{out}, W_{out})`. + + For each output location `output[n, :, d, h, w]`, the size-3 vector `grid[n, d, h, w]` specifies `input_x` pixel + locations x, y, z, which are used to interpolate the output value `output[n, :, d, h, w]`. And `interpolation_mode` + argument specifies "nearest" or "bilinear" interpolation method to sample the input pixels. + + `grid` specifies the sampling pixel locations normalized by the `input_x` spatial dimensions. Therefore, it should + have most values in the range of :math:`[-1, 1]`. + + If `grid` has values outside the range of :math:`[-1, 1]`, the corresponding outputs are handled as defined by + `padding_mode`. If `padding_mode` is set to be "zeros", use :math:`0` for out-of-bound grid locations. If + `padding_mode` is set to be "border", use border values for out-of-bound grid locations. If `padding_mode` is set + to be "reflection", use values at locations reflected by the border for out-of-bound grid locations. For location + far away from the border, it will keep being reflected until becoming in bound. + + Args: + interpolation_mode (str): An optional string specifying the interpolation method. The optional values are + "bilinear" or "nearest". Default: "bilinear". + padding_mode (str): An optional string specifying the pad method. The optional values are "zeros", "border" or + "reflection". Default: "zeros". + align_corners (bool): An optional bool. If set to `True`, the extrema (-1 and 1) are considered as referring to + the center points of the input’s corner pixels. If set to `False`, they are instead considered as referring + to the corner points of the input’s corner pixels, making the sampling more resolution agnostic. Default: + `False`. + + Inputs: + - **input_x** (Tensor) - A 5-D tensor with dtype of float32 or float64 and shape of :math:`(N, C, D_{in}, + H_{in}, W_{in})`. + - **grid** (Tensor) - A 5-D tensor whose dtype is the same as `input_x` and whose shape is :math:`(N, D_{out}, + H_{out}, W_{out}, 3)`. + + Outputs: + A 5-D Tensor whose dtype is the same as `input_x` and whose shape is :math:`(N, C, D_{out}, H_{out}, W_{out})`. + + Raises: + TypeError: If `input_x` or `grid` is not a Tensor. + TypeError: If the dtypes of `input_x` and `grid` are inconsistent. + TypeError: If the dtype of `input_x` or `grid` is not a valid type. + TypeError: If `align_corners` is not a boolean value. + ValueError: If the rank of `input_x` or `grid` is not equal to 5. + ValueError: If the first dimension of `input_x` is not equal to that of `grid`. + ValueError: If the last dimension of `grid` is not equal to 3. + ValueError: If `interpolation_mode` is not "bilinear", "nearest" or a string value. + ValueError: If `padding_mode` is not "zeros", "border", "reflection" or a string value. + + Supported Platforms: + ``CPU`` + + Examples: + >>> gridsampler = ops.GridSampler3D(interpolation_mode='bilinear', padding_mode='zeros', align_corners=True) + >>> input_x = Tensor(np.arange(32).reshape((2, 2, 2, 2, 2)).astype(np.float32)) + >>> grid = Tensor(np.arange(-0.2, 1, 0.1).reshape((2, 2, 1, 1, 3)).astype(np.float32)) + >>> output = gridsampler(input_x, grid) + >>> print(output) + [[[[[ 3.3 ]] + [[ 4.35 ]]] + [[[11.300001]] + [[12.349999]]]] + [[[[21.4 ]] + [[22.449999]]] + [[[29.4 ]] + [[30.449999]]]]] + """ + + @prim_attr_register + def __init__(self, interpolation_mode='bilinear', padding_mode='zeros', align_corners=False): + """Initialize GridSampler3D.""" + validator.check_string(interpolation_mode, ['bilinear', 'nearest'], 'interpolation_mode', self.name) + validator.check_string(padding_mode, ['zeros', 'border', 'reflection'], 'padding_mode', self.name) + validator.check_bool(align_corners, 'align_corners', self.name) + self.init_prim_io_names(inputs=['input_x', 'grid'], outputs=['output']) diff --git a/tests/ut/python/ops/test_ops.py b/tests/ut/python/ops/test_ops.py index a5ad80641f5..a3becb31856 100755 --- a/tests/ut/python/ops/test_ops.py +++ b/tests/ut/python/ops/test_ops.py @@ -2448,6 +2448,11 @@ test_case_nn_ops = [ Tensor(np.array([[0.3, 0.2], [0.4, 0.1]]).astype(np.float32)), Tensor(0.99, mstype.float32)], 'skip': ['backward']}), + ('GridSampler3D', { + 'block': P.GridSampler3D(interpolation_mode='bilinear', padding_mode='zeros', align_corners=False), + 'desc_inputs': [Tensor(np.arange(32).reshape((2, 2, 2, 2, 2)).astype(np.float32)), + Tensor(np.arange(-0.2, 1, 0.1).reshape((2, 2, 1, 1, 3)).astype(np.float32))], + 'desc_bprop': [Tensor(np.arange(-1, 1, 0.25).reshape((2, 2, 2, 1, 1)).astype(np.float32))]}), ] test_case_array_ops = [