1 /* Copyright 2015 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 // See docs in ../ops/nn_ops.cc. 17 18 #define EIGEN_USE_THREADS 19 20 #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" 21 #include "tensorflow/core/framework/numeric_op.h" 22 #include "tensorflow/core/framework/op_kernel.h" 23 #include "tensorflow/core/framework/tensor.h" 24 #include "tensorflow/core/framework/tensor_shape.h" 25 #include "tensorflow/core/kernels/ops_util.h" 26 #include "tensorflow/core/kernels/pooling_ops_common.h" 27 #include "tensorflow/core/lib/core/errors.h" 28 #include "tensorflow/core/platform/logging.h" 29 #include "tensorflow/core/util/padding.h" 30 #include "tensorflow/core/util/tensor_format.h" 31 32 namespace tensorflow { 33 34 typedef Eigen::ThreadPoolDevice CPUDevice; 35 36 template <typename Device, typename T> 37 class QuantizedAvgPoolingOp : public OpKernel { 38 public: QuantizedAvgPoolingOp(OpKernelConstruction * context)39 explicit QuantizedAvgPoolingOp(OpKernelConstruction* context) 40 : OpKernel(context) { 41 OP_REQUIRES_OK(context, context->GetAttr("ksize", &ksize_)); 42 OP_REQUIRES(context, ksize_.size() == 4, 43 errors::InvalidArgument("Sliding window ksize field must " 44 "specify 4 dimensions")); 45 OP_REQUIRES_OK(context, context->GetAttr("strides", &stride_)); 46 OP_REQUIRES(context, stride_.size() == 4, 47 errors::InvalidArgument("Sliding window strides field must " 48 "specify 4 dimensions")); 49 OP_REQUIRES_OK(context, context->GetAttr("padding", &padding_)); 50 OP_REQUIRES(context, ksize_[0] == 1 && stride_[0] == 1, 51 errors::Unimplemented( 52 "Pooling is not yet supported on the batch dimension.")); 53 } 54 Compute(OpKernelContext * context)55 void Compute(OpKernelContext* context) override { 56 const Tensor& tensor_in = context->input(0); 57 PoolParameters params{context, 58 ksize_, 59 stride_, 60 padding_, 61 /*explicit_paddings=*/{}, 62 FORMAT_NHWC, 63 tensor_in.shape()}; 64 if (!context->status().ok()) { 65 return; 66 } 67 68 const float min_input = context->input(1).flat<float>()(0); 69 const float max_input = context->input(2).flat<float>()(0); 70 71 OP_REQUIRES(context, params.depth_window == 1, 72 errors::Unimplemented("Non-spatial pooling is not " 73 "yet supported. Volunteers? :)")); 74 75 OP_REQUIRES(context, tensor_in.dims() == 4, 76 errors::InvalidArgument("tensor_in must be 4-dimensional")); 77 78 Tensor* output = nullptr; 79 OP_REQUIRES_OK(context, context->allocate_output( 80 0, params.forward_output_shape(), &output)); 81 const int32_t highest = static_cast<int32>(Eigen::NumTraits<T>::highest()); 82 const int32_t lowest = static_cast<int32>(Eigen::NumTraits<T>::lowest()); 83 84 // TODO(vrv): Switch this to the Eigen::Tensor version of 85 // SpatialAvgPooling once that version is running quickly. 86 Tensor int32_output(DT_INT32, params.forward_output_shape()); 87 // Cast input to int32 tensor and call SpatialAvgPool. 88 Tensor int32_input(DT_INT32, tensor_in.shape()); 89 int32_input.flat<int32>() = tensor_in.flat<T>().template cast<int32>(); 90 SpatialAvgPool<Device, int32>(context, &int32_output, int32_input, params, 91 padding_); 92 93 // Clamp the int32 output back into quantized space. 94 output->flat<T>() = int32_output.flat<int32>() 95 .cwiseMax(lowest) 96 .cwiseMin(highest) 97 .template cast<T>(); 98 99 Tensor* output_min = nullptr; 100 OP_REQUIRES_OK(context, context->allocate_output(1, {}, &output_min)); 101 output_min->flat<float>()(0) = min_input; 102 Tensor* output_max = nullptr; 103 OP_REQUIRES_OK(context, context->allocate_output(2, {}, &output_max)); 104 output_max->flat<float>()(0) = max_input; 105 } 106 107 private: 108 std::vector<int32> ksize_; 109 std::vector<int32> stride_; 110 Padding padding_; 111 }; 112 113 template <typename Device, typename T> 114 class QuantizedMaxPoolingOp : public MaxPoolingOp<Device, T> { 115 public: QuantizedMaxPoolingOp(OpKernelConstruction * context)116 explicit QuantizedMaxPoolingOp(OpKernelConstruction* context) 117 : MaxPoolingOp<Device, T>(context) {} 118 Compute(OpKernelContext * context)119 void Compute(OpKernelContext* context) override { 120 const float min_input = context->input(1).flat<float>()(0); 121 const float max_input = context->input(2).flat<float>()(0); 122 MaxPoolingOp<Device, T>::Compute(context); 123 Tensor* output_min = nullptr; 124 OP_REQUIRES_OK(context, context->allocate_output(1, {}, &output_min)); 125 output_min->flat<float>()(0) = min_input; 126 Tensor* output_max = nullptr; 127 OP_REQUIRES_OK(context, context->allocate_output(2, {}, &output_max)); 128 output_max->flat<float>()(0) = max_input; 129 } 130 }; 131 132 REGISTER_KERNEL_BUILDER( 133 Name("QuantizedAvgPool").Device(DEVICE_CPU).TypeConstraint<quint8>("T"), 134 QuantizedAvgPoolingOp<CPUDevice, quint8>); 135 136 REGISTER_KERNEL_BUILDER( 137 Name("QuantizedMaxPool").Device(DEVICE_CPU).TypeConstraint<quint8>("T"), 138 QuantizedMaxPoolingOp<CPUDevice, quint8>); 139 140 #ifdef INTEL_MKL 141 REGISTER_KERNEL_BUILDER( 142 Name("QuantizedAvgPool").Device(DEVICE_CPU).TypeConstraint<qint8>("T"), 143 QuantizedAvgPoolingOp<CPUDevice, qint8>); 144 145 REGISTER_KERNEL_BUILDER( 146 Name("QuantizedMaxPool").Device(DEVICE_CPU).TypeConstraint<qint8>("T"), 147 QuantizedMaxPoolingOp<CPUDevice, qint8>); 148 #endif 149 150 } // namespace tensorflow 151