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