forked from huawei/mindspore2022
Add GridSampler3D and GridSampler3DGrad
This commit is contained in:
bubtltk
parent
60c431410e
commit
f5ffb282ee
|
|
@ -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<size_t> &stride, std::vector<size_t> 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<std::string>(kernel_node, "interpolation_mode");
|
||||
padding_mode = AnfAlgo::GetNodeAttr<std::string>(kernel_node, "padding_mode");
|
||||
align_corners = AnfAlgo::GetNodeAttr<bool>(kernel_node, "align_corners");
|
||||
}
|
||||
|
||||
bool GridSampler3DCPUKernel::Launch(const std::vector<kernel::AddressPtr> &inputs,
|
||||
const std::vector<kernel::AddressPtr> &,
|
||||
const std::vector<kernel::AddressPtr> &outputs) {
|
||||
if (dtype_ == kNumberTypeFloat32) {
|
||||
LaunchKernel<float>(inputs, outputs);
|
||||
} else if (dtype_ == kNumberTypeFloat64) {
|
||||
LaunchKernel<double>(inputs, outputs);
|
||||
} else {
|
||||
MS_EXCEPTION(TypeError) << "Input dtype only support float32 and float64!";
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
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<int64_t>(std::floor(x));
|
||||
int64_t y_tnw = static_cast<int64_t>(std::floor(y));
|
||||
int64_t z_tnw = static_cast<int64_t>(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<T>(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<int64_t>(std::round(x));
|
||||
int64_t y_nearest = static_cast<int64_t>(std::round(y));
|
||||
int64_t z_nearest = static_cast<int64_t>(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<T>(kZero);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void GridSampler3DCPUKernel::LaunchKernel(const std::vector<AddressPtr> &inputs,
|
||||
const std::vector<AddressPtr> &outputs) {
|
||||
auto x_data_addr = reinterpret_cast<T *>(inputs[kZero]->addr);
|
||||
auto grid_data_addr = reinterpret_cast<T *>(inputs[kOne]->addr);
|
||||
auto output_data_addr = reinterpret_cast<T *>(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<T>(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 <typename T>
|
||||
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<T>(size - kOne), std::max(coord, static_cast<T>(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<T>(size - kOne), std::max(coord, static_cast<T>(kZero)));
|
||||
}
|
||||
return coord;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
T GridSampler3DCPUKernel::reflect_coordinates(T coord, int64_t twice_low, int64_t twice_high) {
|
||||
if (twice_low == twice_high) {
|
||||
return static_cast<T>(kZero);
|
||||
}
|
||||
T min = static_cast<T>(twice_low) / kTwo;
|
||||
T span = static_cast<T>(twice_high - twice_low) / kTwo;
|
||||
coord = std::fabs(coord - min);
|
||||
T extra = std::fmod(coord, span);
|
||||
int64_t flips = static_cast<int64_t>(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
|
||||
|
|
@ -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 <algorithm>
|
||||
#include <vector>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#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<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
|
||||
const std::vector<AddressPtr> &outputs) override;
|
||||
|
||||
template <typename T>
|
||||
void LaunchKernel(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &outputs);
|
||||
|
||||
private:
|
||||
std::vector<size_t> x_shape_;
|
||||
std::vector<size_t> grid_shape_;
|
||||
std::vector<size_t> output_shape_;
|
||||
std::vector<size_t> x_stride_;
|
||||
std::vector<size_t> grid_stride_;
|
||||
std::vector<size_t> output_stride_;
|
||||
std::string interpolation_mode;
|
||||
std::string padding_mode;
|
||||
bool align_corners;
|
||||
size_t output_number_;
|
||||
TypeId dtype_{kTypeUnknown};
|
||||
template <typename T>
|
||||
void ComputeTask(T *x_data_addr, T *grid_data_addr, T *output_data_addr, const size_t &seq);
|
||||
|
||||
template <typename T>
|
||||
T grid_sampler_compute_source_index(T coord, int64_t size, std::string padding_mode, bool align_corners);
|
||||
|
||||
template <typename T>
|
||||
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_
|
||||
|
|
@ -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<size_t> &stride, std::vector<size_t> 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<std::string>(kernel_node, "interpolation_mode");
|
||||
padding_mode = AnfAlgo::GetNodeAttr<std::string>(kernel_node, "padding_mode");
|
||||
align_corners = AnfAlgo::GetNodeAttr<bool>(kernel_node, "align_corners");
|
||||
}
|
||||
|
||||
bool GridSampler3DGradCPUKernel::Launch(const std::vector<kernel::AddressPtr> &inputs,
|
||||
const std::vector<kernel::AddressPtr> &,
|
||||
const std::vector<kernel::AddressPtr> &outputs) {
|
||||
if (dtype_ == kNumberTypeFloat32) {
|
||||
LaunchKernel<float>(inputs, outputs);
|
||||
} else if (dtype_ == kNumberTypeFloat64) {
|
||||
LaunchKernel<double>(inputs, outputs);
|
||||
} else {
|
||||
MS_EXCEPTION(TypeError) << "Input dtype only support float32 and float64!";
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void GridSampler3DGradCPUKernel::BilinearKernel(std::vector<T *> addr, std::vector<T> location, std::vector<T> mult,
|
||||
std::vector<size_t> ptr) {
|
||||
T x = location[kZero], y = location[kOne], z = location[kTwo];
|
||||
int64_t x_tnw = static_cast<int64_t>(std::floor(x));
|
||||
int64_t y_tnw = static_cast<int64_t>(std::floor(y));
|
||||
int64_t z_tnw = static_cast<int64_t>(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<T>(kZero), gy = static_cast<T>(kZero), gz = static_cast<T>(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 <typename T>
|
||||
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<T *> addr = {grad_addr, x_addr, dx_addr, dgrid_addr};
|
||||
std::vector<T> location = {x, y, z};
|
||||
std::vector<T> mult = {gx_mult, gy_mult, gz_mult};
|
||||
std::vector<size_t> ptr = {grad_ptr_NCDHW, x_ptr_NC, dx_ptr_NC, dgrid_ptr_NDHW};
|
||||
BilinearKernel<T>(addr, location, mult, ptr);
|
||||
} else if (interpolation_mode == "nearest") {
|
||||
int64_t x_nearest = static_cast<int64_t>(std::round(x));
|
||||
int64_t y_nearest = static_cast<int64_t>(std::round(y));
|
||||
int64_t z_nearest = static_cast<int64_t>(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<T>(kZero);
|
||||
dgrid_addr[dgrid_ptr_NDHW + kOne] = static_cast<T>(kZero);
|
||||
dgrid_addr[dgrid_ptr_NDHW + kTwo] = static_cast<T>(kZero);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void GridSampler3DGradCPUKernel::LaunchKernel(const std::vector<AddressPtr> &inputs,
|
||||
const std::vector<AddressPtr> &outputs) {
|
||||
auto grad_data_addr = reinterpret_cast<T *>(inputs[kZero]->addr);
|
||||
auto x_data_addr = reinterpret_cast<T *>(inputs[kOne]->addr);
|
||||
auto grid_data_addr = reinterpret_cast<T *>(inputs[kTwo]->addr);
|
||||
auto dx_data_addr = reinterpret_cast<T *>(outputs[kZero]->addr);
|
||||
auto dgrid_data_addr = reinterpret_cast<T *>(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<T>(kZero);
|
||||
}
|
||||
auto task = [&](size_t start, size_t end) {
|
||||
for (size_t n = start; n < end; n++) {
|
||||
ComputeTask<T>(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 <typename T>
|
||||
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<T>(size - kOne) / kTwo;
|
||||
coord = ((coord + kOne) / kTwo) * (size - kOne);
|
||||
} else {
|
||||
*grad_x = static_cast<T>(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 <typename T>
|
||||
T GridSampler3DGradCPUKernel::clip_coordinates_set_grad(T x, int64_t clip_limit, T *grad_x) {
|
||||
if (x <= static_cast<T>(kZero)) {
|
||||
*grad_x = static_cast<T>(kZero);
|
||||
return static_cast<T>(kZero);
|
||||
} else {
|
||||
T max = static_cast<T>(clip_limit - kOne);
|
||||
if (x >= max) {
|
||||
*grad_x = static_cast<T>(kZero);
|
||||
return max;
|
||||
} else {
|
||||
*grad_x = static_cast<T>(kOne);
|
||||
return x;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
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<T>(kZero);
|
||||
return static_cast<T>(kZero);
|
||||
}
|
||||
int64_t grad_x_mult_;
|
||||
T min = static_cast<T>(twice_low) / kTwo;
|
||||
T span = static_cast<T>(twice_high - twice_low) / kTwo;
|
||||
x = x - min;
|
||||
if (x < static_cast<T>(kZero)) {
|
||||
grad_x_mult_ = -kOne;
|
||||
x = -x;
|
||||
} else {
|
||||
grad_x_mult_ = kOne;
|
||||
}
|
||||
T extra = std::fmod(x, span);
|
||||
int64_t flips = static_cast<int64_t>(std::floor(x / span));
|
||||
if (flips % kTwo == kZero) {
|
||||
*grad_x = static_cast<T>(grad_x_mult_);
|
||||
return extra + min;
|
||||
} else {
|
||||
*grad_x = static_cast<T>(-grad_x_mult_);
|
||||
return span - extra + min;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
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<int64_t>(D);
|
||||
int64_t iH = static_cast<int64_t>(H);
|
||||
int64_t iW = static_cast<int64_t>(W);
|
||||
return d >= 0 && d < iD && h >= 0 && h < iH && w >= 0 && w < iW;
|
||||
}
|
||||
} // namespace kernel
|
||||
} // namespace mindspore
|
||||
|
|
@ -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 <vector>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#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<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
|
||||
const std::vector<AddressPtr> &outputs) override;
|
||||
|
||||
template <typename T>
|
||||
void LaunchKernel(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &outputs);
|
||||
|
||||
private:
|
||||
std::vector<size_t> grad_shape_;
|
||||
std::vector<size_t> x_shape_;
|
||||
std::vector<size_t> grid_shape_;
|
||||
std::vector<size_t> dx_shape_;
|
||||
std::vector<size_t> dgrid_shape_;
|
||||
std::vector<size_t> grad_stride_;
|
||||
std::vector<size_t> x_stride_;
|
||||
std::vector<size_t> grid_stride_;
|
||||
std::vector<size_t> dx_stride_;
|
||||
std::vector<size_t> 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 <typename T>
|
||||
void BilinearKernel(std::vector<T *> addr, std::vector<T> location, std::vector<T> mult, std::vector<size_t> ptr);
|
||||
|
||||
template <typename T>
|
||||
void ComputeTask(T *grad_addr, T *x_addr, T *grid_addr, T *dx_addr, T *dgrid_addr, const size_t &n);
|
||||
|
||||
template <typename T>
|
||||
T grid_sampler_compute_source_index_set_grad(T coord, size_t size, std::string padding_mode, bool align_corners,
|
||||
T *grad_x);
|
||||
|
||||
template <typename T>
|
||||
T reflect_coordinates_set_grad(T x, int64_t twice_low, int64_t twice_high, T *grad_x);
|
||||
|
||||
template <typename T>
|
||||
T clip_coordinates_set_grad(T x, int64_t clip_limit, T *grad_x);
|
||||
|
||||
template <typename T>
|
||||
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_
|
||||
|
|
@ -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<Primitive>("Apply
|
|||
inline const PrimitivePtr kPrimApplyAdagrad = std::make_shared<Primitive>("ApplyAdagrad");
|
||||
inline const PrimitivePtr kPrimApplyAdadelta = std::make_shared<Primitive>("ApplyAdadelta");
|
||||
inline const PrimitivePtr kPrimApplyAdamWithAmsgrad = std::make_shared<Primitive>("ApplyAdamWithAmsgrad");
|
||||
inline const PrimitivePtr kPrimGridSampler3D = std::make_shared<Primitive>(kGridSampler3D);
|
||||
inline const PrimitivePtr kPrimGridSampler3DGrad = std::make_shared<Primitive>(kGridSampler3DGrad);
|
||||
|
||||
// Comm ops
|
||||
inline const PrimitivePtr kPrimMirror = std::make_shared<Primitive>("_MirrorOperator");
|
||||
|
|
|
|||
|
|
@ -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 <set>
|
||||
#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<AbstractBasePtr> &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<int64_t> 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<abstract::Shape>(input_x_shape);
|
||||
abstract::ShapePtr dgrid_shape = std::make_shared<abstract::Shape>(grid_shape);
|
||||
return std::make_shared<abstract::TupleShape>(std::vector<abstract::BaseShapePtr>{dx_shape, dgrid_shape});
|
||||
}
|
||||
|
||||
TuplePtr GridSampler3DGradInferType(const PrimitivePtr &primitive, const std::vector<AbstractBasePtr> &input_args) {
|
||||
std::map<std::string, TypePtr> types;
|
||||
std::set<TypePtr> 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<Tuple>(std::vector<TypePtr>{input_x_type, grid_type});
|
||||
}
|
||||
} // namespace
|
||||
|
||||
AbstractBasePtr GridSampler3DGradInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive,
|
||||
const std::vector<AbstractBasePtr> &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
|
||||
|
|
@ -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 <map>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#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<AbstractBasePtr> &input_args);
|
||||
using PrimGridSampler3DGrad = std::shared_ptr<GridSampler3DGrad>;
|
||||
} // namespace ops
|
||||
} // namespace mindspore
|
||||
|
||||
#endif // MINDSPORE_CORE_OPS_GRID_SAMPLER_3D_H_
|
||||
|
|
@ -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 <set>
|
||||
#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<AbstractBasePtr> &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<int64_t> output_shape = {input_x_shape[kZero], input_x_shape[kOne], grid_shape[kOne], grid_shape[kTwo],
|
||||
grid_shape[kThree]};
|
||||
return std::make_shared<abstract::Shape>(output_shape);
|
||||
}
|
||||
|
||||
TypePtr GridSampler3DInferType(const PrimitivePtr &primitive, const std::vector<AbstractBasePtr> &input_args) {
|
||||
std::map<std::string, TypePtr> types;
|
||||
std::set<TypePtr> 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<AbstractBasePtr> &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
|
||||
|
|
@ -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 <map>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#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<AbstractBasePtr> &input_args);
|
||||
using PrimGridSampler3D = std::shared_ptr<GridSampler3D>;
|
||||
} // namespace ops
|
||||
} // namespace mindspore
|
||||
|
||||
#endif // MINDSPORE_CORE_OPS_GRID_SAMPLER_3D_H_
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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',
|
||||
|
|
|
|||
|
|
@ -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'])
|
||||
|
|
|
|||
|
|
@ -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'])
|
||||
|
|
|
|||
|
|
@ -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 = [
|
||||
|
|
|
|||
Loading…
Reference in New Issue