| 32 | |
| 33 | template <typename T> |
| 34 | void DnnPooling3dOp<T>::Compute(OpKernelContext* context, |
| 35 | se::dnn::PoolingMode pooling_mode, |
| 36 | const std::array<int64, 3>& window, |
| 37 | const std::array<int64, 3>& stride, |
| 38 | const std::array<int64, 3>& padding, |
| 39 | TensorFormat data_format, |
| 40 | const Tensor& tensor_in, Tensor* output) { |
| 41 | const auto in_shape = tensor_in.shape(); |
| 42 | const auto out_shape = output->shape(); |
| 43 | |
| 44 | const int64 in_batch = GetTensorDim(tensor_in, data_format, 'N'); |
| 45 | const int64 in_features = GetTensorDim(tensor_in, data_format, 'C'); |
| 46 | |
| 47 | Tensor transformed_input; |
| 48 | if (data_format == FORMAT_NHWC) { |
| 49 | OP_REQUIRES_OK(context, context->allocate_temp( |
| 50 | DataTypeToEnum<T>::value, |
| 51 | ShapeFromFormat(FORMAT_NCHW, tensor_in.shape(), |
| 52 | data_format), |
| 53 | &transformed_input)); |
| 54 | functor::NHWCToNCHW<GPUDevice, T, 5>()(context->eigen_device<GPUDevice>(), |
| 55 | tensor_in.tensor<T, 5>(), |
| 56 | transformed_input.tensor<T, 5>()); |
| 57 | } else { |
| 58 | transformed_input = tensor_in; |
| 59 | } |
| 60 | Tensor transformed_output; |
| 61 | if (data_format == FORMAT_NHWC) { |
| 62 | OP_REQUIRES_OK(context, |
| 63 | context->allocate_temp( |
| 64 | DataTypeToEnum<T>::value, |
| 65 | ShapeFromFormat(FORMAT_NCHW, out_shape, data_format), |
| 66 | &transformed_output)); |
| 67 | } else { |
| 68 | transformed_output = *output; |
| 69 | } |
| 70 | |
| 71 | se::dnn::PoolingDescriptor pooling_desc(3); |
| 72 | pooling_desc.set_pooling_mode(pooling_mode); |
| 73 | se::dnn::BatchDescriptor input_desc(3); |
| 74 | input_desc.set_count(in_batch) |
| 75 | .set_feature_map_count(in_features) |
| 76 | .set_layout(se::dnn::DataLayout::kBatchDepthYX); |
| 77 | se::dnn::BatchDescriptor output_desc(3); |
| 78 | output_desc.set_count(in_batch) |
| 79 | .set_feature_map_count(in_features) |
| 80 | .set_layout(se::dnn::DataLayout::kBatchDepthYX); |
| 81 | for (size_t i = 0; i < window.size(); ++i) { |
| 82 | const auto dim_i = static_cast<se::dnn::DimIndex>(i); |
| 83 | pooling_desc.set_window(dim_i, window[i]); |
| 84 | pooling_desc.set_stride(dim_i, stride[i]); |
| 85 | pooling_desc.set_padding(dim_i, padding[i]); |
| 86 | input_desc.set_spatial_dim(dim_i, |
| 87 | GetTensorDim(tensor_in, data_format, '2' - i)); |
| 88 | output_desc.set_spatial_dim(dim_i, |
| 89 | GetTensorDim(out_shape, data_format, '2' - i)); |
| 90 | } |
| 91 |
nothing calls this directly
no test coverage detected