1 /* Copyright 2018 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 /*
17 See extract_image_patches_op* files and docs for extract_image_patches in
18 ../ops/image_ops.cc.
19
20 Rates are not supported as of now, but the comments hint how to edit the code
21 when rates are to be added.
22 */
23
24 #define USE_EIGEN_TENSOR
25 #define EIGEN_USE_THREADS
26
27 #include "tensorflow/core/kernels/image/extract_volume_patches_op.h"
28
29 #include <vector>
30
31 #include "tensorflow/core/framework/bounds_check.h"
32 #include "tensorflow/core/framework/kernel_shape_util.h"
33 #include "tensorflow/core/framework/numeric_op.h"
34 #include "tensorflow/core/framework/op_kernel.h"
35 #include "tensorflow/core/framework/register_types.h"
36 #include "tensorflow/core/framework/tensor.h"
37 #include "tensorflow/core/kernels/ops_util.h"
38 #include "tensorflow/core/lib/core/errors.h"
39 #include "tensorflow/core/platform/logging.h"
40 #include "tensorflow/core/platform/macros.h"
41 #include "tensorflow/core/util/tensor_format.h"
42
43 namespace tensorflow {
44
45 typedef Eigen::ThreadPoolDevice CPUDevice;
46 typedef Eigen::GpuDevice GPUDevice;
47
ParseAttributeVec5(OpKernelConstruction * context,const string & attr_name,std::vector<int32> * attr)48 static inline void ParseAttributeVec5(OpKernelConstruction* context,
49 const string& attr_name,
50 std::vector<int32>* attr) {
51 OP_REQUIRES_OK(context, context->GetAttr(attr_name, attr));
52 OP_REQUIRES(
53 context, (*attr)[0] == 1 && (*attr)[4] == 1,
54 errors::Unimplemented("Only support ", attr_name, " across space."));
55 OP_REQUIRES(context, (*attr)[1] >= 1 && (*attr)[2] >= 1 && (*attr)[3] >= 1,
56 errors::OutOfRange(attr_name, " is out of range."));
57 }
58
59 template <typename Device, typename T>
60 class ExtractVolumePatchesOp : public UnaryOp<T> {
61 public:
ExtractVolumePatchesOp(OpKernelConstruction * context)62 explicit ExtractVolumePatchesOp(OpKernelConstruction* context)
63 : UnaryOp<T>(context) {
64 ParseAttributeVec5(context, "ksizes", &ksizes_);
65 ParseAttributeVec5(context, "strides", &strides_);
66 // ParseAttributeVec5(context, "rates", &rates_);
67 OP_REQUIRES_OK(context, context->GetAttr("padding", &padding_));
68 }
69
Compute(OpKernelContext * context)70 void Compute(OpKernelContext* context) override {
71 // Input tensor is of the following dimensions:
72 // [ batch, in_planes, in_rows, in_cols, channels ]
73 const Tensor& input = context->input(0);
74 OP_REQUIRES(context, input.dims() == 5,
75 errors::InvalidArgument("input must be 5-dimensional",
76 input.shape().DebugString()));
77
78 const int batch = input.dim_size(0);
79 const int in_planes = input.dim_size(1);
80 const int in_rows = input.dim_size(2);
81 const int in_cols = input.dim_size(3);
82 const int depth = input.dim_size(4);
83
84 const int ksize_planes = ksizes_[1];
85 const int ksize_rows = ksizes_[2];
86 const int ksize_cols = ksizes_[3];
87
88 const int stride_planes = strides_[1];
89 const int stride_rows = strides_[2];
90 const int stride_cols = strides_[3];
91
92 /*
93 // TODO(hsgkim): enable rates
94 // Rates are disabled as of now due to Eigen's definitions of
95 // `extract_volume_patch` functions; none of them accept rates
96 // as its argument and rates are fixed to (1, 1, 1, 1, 1). A
97 // workaround has to be found for this.
98 // In order to enable rates, uncomment the following lines and use
99 // ksize_*_eff instead of ksize_* for the second argument of
100 // GetWindowedOutputSize calls.
101
102 const int rate_planes = rates_[1];
103 const int rate_rows = rates_[2];
104 const int rate_cols = rates_[3];
105
106 const int ksize_planes_eff = ksize_planes +
107 (ksize_planes - 1) * (rate_planes - 1);
108 const int ksize_rows_eff = ksize_rows + (ksize_rows - 1) * (rate_rows - 1);
109 const int ksize_cols_eff = ksize_cols + (ksize_cols - 1) * (rate_cols - 1);
110 */
111
112 int64 out_planes = 0, out_rows = 0, out_cols = 0;
113 int64 pad_planes = 0, pad_rows = 0, pad_cols = 0;
114 OP_REQUIRES_OK(context,
115 GetWindowedOutputSize(in_planes, ksize_planes, stride_planes,
116 padding_, &out_planes, &pad_planes));
117 OP_REQUIRES_OK(context,
118 GetWindowedOutputSize(in_rows, ksize_rows, stride_rows,
119 padding_, &out_rows, &pad_rows));
120 OP_REQUIRES_OK(context,
121 GetWindowedOutputSize(in_cols, ksize_cols, stride_cols,
122 padding_, &out_cols, &pad_cols));
123
124 const std::vector<int64> out_sizes = {
125 batch, out_planes, out_rows, out_cols,
126 ksize_planes * ksize_rows * ksize_cols * depth};
127 TensorShape out_shape(out_sizes);
128
129 Tensor* output = nullptr;
130 OP_REQUIRES_OK(context, context->allocate_output(0, out_shape, &output));
131
132 // If there is nothing to compute, return.
133 if (out_shape.num_elements() == 0) {
134 return;
135 }
136
137 functor::ExtractVolumePatchesForward<Device, T>()(
138 context->eigen_device<Device>(), input.tensor<T, 5>(), ksize_planes,
139 ksize_rows, ksize_cols, stride_planes, stride_rows, stride_cols,
140 /* rate_planes, rate_rows, rate_cols, */
141 BrainPadding2EigenPadding(padding_), output->tensor<T, 5>());
142 }
143
144 private:
145 std::vector<int32> ksizes_;
146 std::vector<int32> strides_;
147 // std::vector<int32> rates_;
148
149 Padding padding_;
150
151 TF_DISALLOW_COPY_AND_ASSIGN(ExtractVolumePatchesOp);
152 };
153
154 // Registration of the CPU implementations.
155 #define REGISTER(T) \
156 REGISTER_KERNEL_BUILDER( \
157 Name("ExtractVolumePatches").Device(DEVICE_CPU).TypeConstraint<T>("T"), \
158 ExtractVolumePatchesOp<CPUDevice, T>);
159
160 TF_CALL_REAL_NUMBER_TYPES(REGISTER);
161
162 #undef REGISTER
163
164 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
165
166 // Forward declarations of the functor specializations for GPU.
167 namespace functor {
168
169 // clang-format off
170 #define DECLARE_GPU_SPEC(T) \
171 template <> \
172 void ExtractVolumePatchesForward<GPUDevice, T>::operator()( \
173 const GPUDevice& d, typename TTypes<T, 5>::ConstTensor input, \
174 int patch_planes, int patch_rows, int patch_cols, \
175 int stride_planes, int stride_rows, int stride_cols, \
176 /* int rate_planes, int rate_rows, int rate_cols, */ \
177 const Eigen::PaddingType& padding, \
178 typename TTypes<T, 5>::Tensor output); \
179 extern template struct ExtractVolumePatchesForward<GPUDevice, T>;
180 // clang-format on
181
182 TF_CALL_GPU_NUMBER_TYPES(DECLARE_GPU_SPEC);
183
184 #undef DECLARE_GPU_SPEC
185
186 } // namespace functor
187
188 // Registration of the GPU implementations.
189 #define REGISTER(T) \
190 REGISTER_KERNEL_BUILDER( \
191 Name("ExtractVolumePatches").Device(DEVICE_GPU).TypeConstraint<T>("T"), \
192 ExtractVolumePatchesOp<GPUDevice, T>);
193
194 TF_CALL_GPU_NUMBER_TYPES(REGISTER);
195
196 #undef REGISTER
197
198 #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
199
200 } // namespace tensorflow
201