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