| 22 | |
| 23 | template <typename T, typename Context> |
| 24 | void GridSampleKernel(const Context& dev_ctx, |
| 25 | const DenseTensor& x, |
| 26 | const DenseTensor& grid, |
| 27 | const std::string& mode, |
| 28 | const std::string& padding_mode, |
| 29 | bool align_corners, |
| 30 | DenseTensor* out) { |
| 31 | if (out && out->numel() == 0) { |
| 32 | dev_ctx.template Alloc<T>(out); |
| 33 | return; |
| 34 | } |
| 35 | // attrs |
| 36 | // paddle.nn.functional.grid_sample(x, grid, mode='bilinear', |
| 37 | // padding_mode='zeros', align_corners=True, name=None) |
| 38 | const std::string data_format = DataLayoutToString(x.layout()); |
| 39 | |
| 40 | // attr to real param |
| 41 | bool is_nearest_bool; |
| 42 | if (mode == "bilinear") { |
| 43 | is_nearest_bool = false; |
| 44 | } else if (mode == "nearest") { |
| 45 | is_nearest_bool = true; |
| 46 | } else { |
| 47 | PADDLE_THROW(errors::InvalidArgument( |
| 48 | "should not reach here: mode should be either 'bilinear' or " |
| 49 | "'nearest', bot got %s.", |
| 50 | mode)); |
| 51 | } |
| 52 | |
| 53 | // attention: 0: zeros, 2: reflection, 1: border according to XDNN api. |
| 54 | int padding_mode_int; |
| 55 | if (padding_mode == "zeros") { |
| 56 | padding_mode_int = 0; |
| 57 | } else if (padding_mode == "reflection") { |
| 58 | padding_mode_int = 2; |
| 59 | } else if (padding_mode == "border") { |
| 60 | padding_mode_int = 1; |
| 61 | } else { |
| 62 | PADDLE_THROW(errors::InvalidArgument( |
| 63 | "should not reach here: padding_mode should be either 'zeros' or " |
| 64 | "'reflection' or 'border', bot got %s.", |
| 65 | padding_mode)); |
| 66 | } |
| 67 | |
| 68 | const T* input_data = x.data<T>(); |
| 69 | const T* grid_data = grid.data<T>(); |
| 70 | |
| 71 | int64_t n = x.dims()[0]; |
| 72 | int64_t c = x.dims()[1]; |
| 73 | |
| 74 | if (x.dims().size() == 4) { // 2D grid sample |
| 75 | int64_t h = x.dims()[2]; |
| 76 | int64_t w = x.dims()[3]; |
| 77 | int64_t out_h = grid.dims()[1]; |
| 78 | int64_t out_w = grid.dims()[2]; |
| 79 | |
| 80 | bool is_nchw_bool; |
| 81 | if (data_format == "NCHW") { |