Home | History | Annotate | Download | only in optimized
      1 /* Copyright 2017 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 // Copied from tensorflow/core/kernels/eigen_spatial_convolutions.h.
     17 // TODO(petewarden) - move this to a common location in Eigen itself.
     18 
     19 #ifndef TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_EIGEN_SPATIAL_CONVOLUTIONS_H_
     20 #define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_EIGEN_SPATIAL_CONVOLUTIONS_H_
     21 
     22 #define EIGEN_USE_CUSTOM_THREAD_POOL
     23 #define EIGEN_USE_THREADS
     24 
     25 // NOTE: Eigen is slightly different internally and externally. We need to
     26 // hack the unsupported/Eigen/CXX11/Tensor header instantiation macros at
     27 // specific places, so we need two copies of the hacked file, one for
     28 // internal and one for external.
     29 // If you have trouble simply undef out the reducer macro e.g.
     30 // TFLITE_REDUCE_INSTANTIATIONS_GOOGLE, but be aware this will make
     31 // the binary much bigger!
     32 #define TFLITE_REDUCE_INSTANTIATIONS_OPEN_SOURCE
     33 #define Eigen EigenForTFLite
     34 #if defined(TFLITE_REDUCE_INSTANTIATIONS_GOOGLE)
     35 #include "tensorflow/contrib/lite/kernels/internal/optimized/eigen_tensor_reduced_instantiations_google.h"
     36 #elif defined(TFLITE_REDUCE_INSTANTIATIONS_OPEN_SOURCE)
     37 #include "tensorflow/contrib/lite/kernels/internal/optimized/eigen_tensor_reduced_instantiations_oss.h"
     38 #else
     39 #include "unsupported/Eigen/CXX11/Tensor"
     40 #endif
     41 
     42 namespace Eigen {
     43 
     44 /** SpatialConvolution
     45  * \ingroup CXX11_NeuralNetworks_Module
     46  *
     47  * \brief Applies a 2D convolution over a multichannel input image.
     48  *
     49  * The input parameter is expected to be a tensor with a rank of 3 or more
     50  * (channels, height, width, and optionally others)
     51  * The kernel parameter is expected to be a 4D tensor (filters, channels,
     52  * kernel_height, kernel_width)
     53  * The input and the kernel must both be in col-major layout. The result will
     54  * also be in col-major layout.
     55  *
     56  * If col_in_stride, row_in_stride > 1, then applies convolution with holes
     57  * (aka atrous convolution), sampling every col_in_stride, row_in_stride input
     58  * pixels.
     59  *
     60  * The result can be assigned to a tensor of rank equal to the rank of the
     61  * input. The dimensions of the result will be filters, height, width (and
     62  * others if applicable).
     63  *
     64  * It is possible to swap the order of the width and height dimensions provided
     65  * that the same order is used in the input, the kernel, and the output.
     66  *
     67  */
     68 template <typename Input, typename Kernel>
     69 EIGEN_DEVICE_FUNC
     70     EIGEN_ALWAYS_INLINE static const typename internal::conditional<
     71         internal::traits<Input>::Layout == ColMajor,
     72         TensorReshapingOp<
     73             const DSizes<typename internal::traits<Input>::Index,
     74                          internal::traits<Input>::NumDimensions>,
     75             const TensorContractionOp<
     76                 const array<IndexPair<typename internal::traits<Input>::Index>,
     77                             1>,
     78                 const TensorReshapingOp<
     79                     const DSizes<typename internal::traits<Input>::Index, 2>,
     80                     const Kernel>,
     81                 const TensorReshapingOp<
     82                     const DSizes<typename internal::traits<Input>::Index, 2>,
     83                     const TensorImagePatchOp<Dynamic, Dynamic,
     84                                              const Input> > > >,
     85         TensorReshapingOp<
     86             const DSizes<typename internal::traits<Input>::Index,
     87                          internal::traits<Input>::NumDimensions>,
     88             const TensorContractionOp<
     89                 const array<IndexPair<typename internal::traits<Input>::Index>,
     90                             1>,
     91                 const TensorReshapingOp<
     92                     const DSizes<typename internal::traits<Input>::Index, 2>,
     93                     const TensorImagePatchOp<Dynamic, Dynamic, const Input> >,
     94                 const TensorReshapingOp<
     95                     const DSizes<typename internal::traits<Input>::Index, 2>,
     96                     const Kernel> > > >::type
     97     SpatialConvolution(const Input& input, const Kernel& kernel,
     98                        const DenseIndex row_stride = 1,
     99                        const DenseIndex col_stride = 1,
    100                        const PaddingType padding_type = PADDING_SAME,
    101                        const DenseIndex row_in_stride = 1,
    102                        const DenseIndex col_in_stride = 1) {
    103   typedef typename internal::traits<Input>::Index TensorIndex;
    104   TensorRef<Tensor<typename internal::traits<Input>::Scalar,
    105                    internal::traits<Input>::NumDimensions,
    106                    internal::traits<Input>::Layout, TensorIndex> >
    107       in(input);
    108   TensorRef<Tensor<typename internal::traits<Kernel>::Scalar,
    109                    internal::traits<Kernel>::NumDimensions,
    110                    internal::traits<Kernel>::Layout, TensorIndex> >
    111       kern(kernel);
    112 
    113   EIGEN_STATIC_ASSERT(
    114       internal::traits<Input>::Layout == internal::traits<Kernel>::Layout,
    115       YOU_MADE_A_PROGRAMMING_MISTAKE);
    116   const bool isColMajor = (internal::traits<Input>::Layout == ColMajor);
    117 
    118   const int NumDims = internal::traits<Input>::NumDimensions;
    119 
    120   // Number of filters to apply. This is the same as the output depth of the
    121   // result
    122   const TensorIndex kernelFilters =
    123       isColMajor ? kern.dimensions()[0] : kern.dimensions()[3];
    124   // Number of channels. This is the same as the input depth.
    125   const TensorIndex kernelChannels =
    126       isColMajor ? kern.dimensions()[1] : kern.dimensions()[2];
    127   const TensorIndex kernelRows =
    128       isColMajor ? kern.dimensions()[2] : kern.dimensions()[1];
    129   const TensorIndex kernelCols =
    130       isColMajor ? kern.dimensions()[3] : kern.dimensions()[0];
    131 
    132   const DenseIndex kernelRowsEff =
    133       kernelRows + (kernelRows - 1) * (row_in_stride - 1);
    134   const DenseIndex kernelColsEff =
    135       kernelCols + (kernelCols - 1) * (col_in_stride - 1);
    136 
    137   array<IndexPair<TensorIndex>, 1> contract_dims;
    138   contract_dims[0] = IndexPair<TensorIndex>(1, 0);
    139 
    140   const TensorIndex InputRows =
    141       isColMajor ? in.dimension(1) : in.dimension(NumDims - 2);
    142   const TensorIndex InputCols =
    143       isColMajor ? in.dimension(2) : in.dimension(NumDims - 3);
    144 
    145   TensorIndex out_height;
    146   TensorIndex out_width;
    147   switch (padding_type) {
    148     case PADDING_VALID:
    149       out_height = numext::ceil((InputRows - kernelRowsEff + 1.f) /
    150                                 static_cast<float>(row_stride));
    151       out_width = numext::ceil((InputCols - kernelColsEff + 1.f) /
    152                                static_cast<float>(col_stride));
    153       break;
    154     case PADDING_SAME:
    155       out_height = numext::ceil(InputRows / static_cast<float>(row_stride));
    156       out_width = numext::ceil(InputCols / static_cast<float>(col_stride));
    157       break;
    158     default:
    159       // Initialize unused variables to avoid a compiler warning
    160       out_height = 0;
    161       out_width = 0;
    162       eigen_assert(false && "unexpected padding");
    163   }
    164 
    165   // Molds the output of the patch extraction code into a 2d tensor:
    166   // - the first dimension (dims[0]): the patch values to be multiplied with the
    167   // kernels
    168   // - the second dimension (dims[1]): everything else
    169   DSizes<TensorIndex, 2> pre_contract_dims;
    170   if (isColMajor) {
    171     pre_contract_dims[0] = kernelChannels * kernelRows * kernelCols;
    172     pre_contract_dims[1] = out_height * out_width;
    173     for (int i = 3; i < NumDims; ++i) {
    174       pre_contract_dims[1] *= in.dimension(i);
    175     }
    176   } else {
    177     pre_contract_dims[1] = kernelChannels * kernelRows * kernelCols;
    178     pre_contract_dims[0] = out_height * out_width;
    179     for (int i = 0; i < NumDims - 3; ++i) {
    180       pre_contract_dims[0] *= in.dimension(i);
    181     }
    182   }
    183 
    184   // Molds the output of the contraction into the shape expected by the used
    185   // (assuming this is ColMajor):
    186   // - 1st dim: kernel filters
    187   // - 2nd dim: output height
    188   // - 3rd dim: output width
    189   // - 4th dim and beyond: everything else including batch size
    190   DSizes<TensorIndex, NumDims> post_contract_dims;
    191   if (isColMajor) {
    192     post_contract_dims[0] = kernelFilters;
    193     post_contract_dims[1] = out_height;
    194     post_contract_dims[2] = out_width;
    195     for (int i = 3; i < NumDims; ++i) {
    196       post_contract_dims[i] = in.dimension(i);
    197     }
    198   } else {
    199     post_contract_dims[NumDims - 1] = kernelFilters;
    200     post_contract_dims[NumDims - 2] = out_height;
    201     post_contract_dims[NumDims - 3] = out_width;
    202     for (int i = 0; i < NumDims - 3; ++i) {
    203       post_contract_dims[i] = in.dimension(i);
    204     }
    205   }
    206 
    207   DSizes<TensorIndex, 2> kernel_dims;
    208   if (isColMajor) {
    209     kernel_dims[0] = kernelFilters;
    210     kernel_dims[1] = kernelChannels * kernelRows * kernelCols;
    211   } else {
    212     kernel_dims[0] = kernelChannels * kernelRows * kernelCols;
    213     kernel_dims[1] = kernelFilters;
    214   }
    215   // TODO(yangke): choose() is defined in TensorContraction.h -- consider
    216   // moving it to somewhere more "common".
    217   return input
    218       .extract_image_patches(kernelRows, kernelCols, row_stride, col_stride,
    219                              row_in_stride, col_in_stride, 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 // clang-format on
    228 
    229 #endif  // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_EIGEN_SPATIAL_CONVOLUTIONS_H_
    230