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 #ifndef TENSORFLOW_CORE_KERNELS_EIGEN_CUBOID_CONVOLUTION_H_ 17 #define TENSORFLOW_CORE_KERNELS_EIGEN_CUBOID_CONVOLUTION_H_ 18 19 #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" 20 #include "tensorflow/core/kernels/eigen_volume_patch.h" 21 22 namespace Eigen { 23 24 /** CuboidConvolution 25 * \ingroup CXX11_NeuralNetworks_Module 26 * 27 * \brief Applies a 3D convolution over a multichannel input voxel block. 28 * 29 * The input parameter is expected to be a tensor with a rank of 4 or more 30 * (channels, depth, height, width, and optionally others). 31 * The kernel parameter is expected to be a 5D tensor (filters, channels, 32 * kernel_depth, kernel_height, kernel_width). 33 * The result can be assigned to a tensor of rank equal to the rank of the 34 * input. The dimensions of the result will be filters, depth, height, width 35 * (and others if applicable). 36 * 37 * The input and kernel have to be in the same layout, and both row-major and 38 * col-major are supported. The shapes given above are for col-major layout. 39 * For row-major, all dimensions should be reversed. 40 * 41 * It is possible to swap the order of the depth, width, and height dimensions 42 * provided that the same order is used in the input, the kernel, and the 43 * output. 44 */ 45 template <typename Input, typename Kernel> 46 EIGEN_ALWAYS_INLINE static const typename internal::conditional< 47 internal::traits<Input>::Layout == ColMajor, 48 TensorReshapingOp< 49 const DSizes<typename internal::traits<Input>::Index, 50 internal::traits<Input>::NumDimensions>, 51 const TensorContractionOp< 52 const array<IndexPair<typename internal::traits<Input>::Index>, 1>, 53 const TensorReshapingOp< 54 const DSizes<typename internal::traits<Input>::Index, 2>, 55 const Kernel>, 56 const TensorReshapingOp< 57 const DSizes<typename internal::traits<Input>::Index, 2>, 58 const TensorVolumePatchOp<Dynamic, Dynamic, Dynamic, 59 const Input> > > >, 60 TensorReshapingOp< 61 const DSizes<typename internal::traits<Input>::Index, 62 internal::traits<Input>::NumDimensions>, 63 const TensorContractionOp< 64 const array<IndexPair<typename internal::traits<Input>::Index>, 1>, 65 const TensorReshapingOp< 66 const DSizes<typename internal::traits<Input>::Index, 2>, 67 const TensorVolumePatchOp<Dynamic, Dynamic, Dynamic, 68 const Input> >, 69 const TensorReshapingOp< 70 const DSizes<typename internal::traits<Input>::Index, 2>, 71 const Kernel> > > >::type 72 CuboidConvolution(const Input& input, const Kernel& kernel, 73 const DenseIndex stridePlanes = 1, 74 const DenseIndex strideRows = 1, 75 const DenseIndex strideCols = 1, 76 const PaddingType padding_type = PADDING_SAME) { 77 typedef typename internal::traits<Input>::Index TensorIndex; 78 TensorRef<Tensor<typename internal::traits<Input>::Scalar, 79 internal::traits<Input>::NumDimensions, 80 internal::traits<Input>::Layout, TensorIndex> > 81 in(input); 82 TensorRef<Tensor<typename internal::traits<Kernel>::Scalar, 83 internal::traits<Kernel>::NumDimensions, 84 internal::traits<Kernel>::Layout, TensorIndex> > 85 kern(kernel); 86 87 EIGEN_STATIC_ASSERT( 88 internal::traits<Input>::Layout == internal::traits<Kernel>::Layout, 89 YOU_MADE_A_PROGRAMMING_MISTAKE); 90 static const bool isColMajor = (internal::traits<Input>::Layout == ColMajor); 91 static const int NumDims = internal::traits<Input>::NumDimensions; 92 93 // Number of filters to apply. This is the same as the output depth of the 94 // result. 95 const TensorIndex kernelFilters = 96 isColMajor ? kern.dimensions()[0] : kern.dimensions()[4]; 97 const TensorIndex kernelChannels = 98 isColMajor ? kern.dimensions()[1] : kern.dimensions()[3]; 99 100 // Spatial size of the kernel. 101 const TensorIndex kernelDepth = 102 isColMajor ? kern.dimensions()[2] : kern.dimensions()[2]; 103 const TensorIndex kernelRows = 104 isColMajor ? kern.dimensions()[3] : kern.dimensions()[1]; 105 const TensorIndex kernelCols = 106 isColMajor ? kern.dimensions()[4] : kern.dimensions()[0]; 107 108 if (isColMajor) { 109 eigen_assert(kernelChannels == in.dimension(0)); 110 } else { 111 eigen_assert(kernelChannels == in.dimension(NumDims - 1)); 112 } 113 114 const TensorIndex inputPlanes = 115 isColMajor ? in.dimension(1) : in.dimension(NumDims - 2); 116 const TensorIndex inputRows = 117 isColMajor ? in.dimension(2) : in.dimension(NumDims - 3); 118 const TensorIndex inputCols = 119 isColMajor ? in.dimension(3) : in.dimension(NumDims - 4); 120 121 TensorIndex out_depth; 122 TensorIndex out_height; 123 TensorIndex out_width; 124 switch (padding_type) { 125 case PADDING_VALID: 126 out_depth = Eigen::divup(inputPlanes - kernelDepth + 1, 127 static_cast<TensorIndex>(stridePlanes)); 128 out_height = Eigen::divup(inputRows - kernelRows + 1, 129 static_cast<TensorIndex>(strideRows)); 130 out_width = Eigen::divup(inputCols - kernelCols + 1, 131 static_cast<TensorIndex>(strideCols)); 132 break; 133 case PADDING_SAME: 134 out_depth = 135 Eigen::divup(inputPlanes, static_cast<TensorIndex>(stridePlanes)); 136 out_height = 137 Eigen::divup(inputRows, static_cast<TensorIndex>(strideRows)); 138 out_width = Eigen::divup(inputCols, static_cast<TensorIndex>(strideCols)); 139 break; 140 default: 141 out_depth = 0; 142 out_height = 0; 143 out_width = 0; 144 eigen_assert(false && "unexpected padding"); 145 } 146 147 DSizes<TensorIndex, 2> kernel_dims; 148 if (isColMajor) { 149 kernel_dims[0] = kernelFilters; 150 kernel_dims[1] = kernelChannels * kernelDepth * kernelRows * kernelCols; 151 } else { 152 kernel_dims[0] = kernelChannels * kernelDepth * kernelRows * kernelCols; 153 kernel_dims[1] = kernelFilters; 154 } 155 156 // Molds the output of the patch extraction result into a 2D tensor: 157 // - the first dimension (dims[0]): the patch values to be multiplied with the 158 // kernels 159 // - the second dimension (dims[1]): everything else 160 DSizes<TensorIndex, 2> pre_contract_dims; 161 if (isColMajor) { 162 pre_contract_dims[0] = 163 kernelChannels * kernelDepth * kernelRows * kernelCols; 164 pre_contract_dims[1] = out_depth * out_height * out_width; 165 for (int i = 4; i < NumDims; ++i) { 166 pre_contract_dims[1] *= in.dimension(i); 167 } 168 } else { 169 pre_contract_dims[1] = 170 kernelChannels * kernelDepth * kernelRows * kernelCols; 171 pre_contract_dims[0] = out_depth * out_height * out_width; 172 for (int i = 0; i < NumDims - 4; ++i) { 173 pre_contract_dims[0] *= in.dimension(i); 174 } 175 } 176 177 array<IndexPair<TensorIndex>, 1> contract_dims; 178 contract_dims[0] = IndexPair<TensorIndex>(1, 0); 179 180 // Molds the output of the contraction into the shape expected by the user 181 // (assuming ColMajor): 182 // - 1st dim: kernel filters 183 // - 2nd dim: output depth 184 // - 3nd dim: output height 185 // - 4rd dim: output width 186 // - 5th dim and beyond: everything else including batch size 187 DSizes<TensorIndex, NumDims> post_contract_dims; 188 if (isColMajor) { 189 post_contract_dims[0] = kernelFilters; 190 post_contract_dims[1] = out_depth; 191 post_contract_dims[2] = out_height; 192 post_contract_dims[3] = out_width; 193 for (int i = 4; i < NumDims; ++i) { 194 post_contract_dims[i] = in.dimension(i); 195 } 196 } else { 197 post_contract_dims[NumDims - 1] = kernelFilters; 198 post_contract_dims[NumDims - 2] = out_depth; 199 post_contract_dims[NumDims - 3] = out_height; 200 post_contract_dims[NumDims - 4] = out_width; 201 for (int i = 0; i < NumDims - 4; ++i) { 202 post_contract_dims[i] = in.dimension(i); 203 } 204 } 205 206 return choose( 207 Cond<internal::traits<Input>::Layout == ColMajor>(), 208 kernel.reshape(kernel_dims) 209 .contract(input 210 .extract_volume_patches( 211 kernelDepth, kernelRows, kernelCols, stridePlanes, 212 strideRows, strideCols, padding_type) 213 .reshape(pre_contract_dims), 214 contract_dims) 215 .reshape(post_contract_dims), 216 input 217 .extract_volume_patches(kernelDepth, kernelRows, kernelCols, 218 stridePlanes, strideRows, strideCols, 219 padding_type) 220 .reshape(pre_contract_dims) 221 .contract(kernel.reshape(kernel_dims), contract_dims) 222 .reshape(post_contract_dims)); 223 } 224 225 } // end namespace Eigen 226 227 #endif // TENSORFLOW_CORE_KERNELS_EIGEN_CUBOID_CONVOLUTION_H_ 228
