ROperator_Pool.hxx

Go to the documentation of this file.

#ifndef TMVA_SOFIE_ROPERATOR_POOL
#define TMVA_SOFIE_ROPERATOR_POOL
 
#include "TMVA/SOFIE_common.hxx"
#include "TMVA/ROperator.hxx"
#include "TMVA/RModel.hxx"
 
#include <memory>
#include <sstream>
#include <algorithm>
#include <stdexcept>
#include <vector>
#include <cassert>
 
namespace TMVA {
namespace Experimental {
namespace SOFIE {
 
struct RAttributes_Pool {
   // structure that contains Pool attribute
   std::string auto_pad = "NOTSET";
   int ceil_mode = 0;
   int count_include_pad = 0;     // not for MaxPool
   int storage_order = 0;         // not for AveragePool
   std::vector<size_t> dilations; // not for AveragePool
   std::vector<size_t> kernel_shape;
   std::vector<size_t> pads;
   std::vector<size_t> strides;
};
 
enum PoolOpMode { InvalidPool, MaxPool, AveragePool, GlobalAveragePool };
 
template<typename T>
class ROperator_Pool final : public ROperator
{
 
private:
 
   PoolOpMode fPoolMode;
 
   size_t fAttrCeilMode;
   size_t fAttrCountIncludePad;
   size_t fAttrStorageOrder;
   std::string fAttrAutopad;
   std::vector<size_t> fAttrDilations;
   std::vector<size_t> fAttrKernelShape;
   std::vector<size_t> fAttrPads;
   std::vector<size_t> fAttrStrides;
 
   std::string fNX;
   std::string fNY;
 
   std::vector<size_t> fShapeX;
   std::vector<size_t> fShapeY;
 
   std::string fType;
 
   size_t fDim;   // dimension of the MaxPool
   bool fUseSession = false;
 
public:
 
   std::string Name() {
      if (fPoolMode == AveragePool)  return "AveragePool";
      if (fPoolMode == MaxPool)  return "MaxPool";
      return "Invalid";
   }
 
   ROperator_Pool() {}
 
   ROperator_Pool(PoolOpMode mode, RAttributes_Pool attr, std::string nameX, std::string nameY)
      : fPoolMode(mode), fAttrCeilMode(attr.ceil_mode), fAttrCountIncludePad(attr.count_include_pad),
        fAttrStorageOrder(attr.storage_order), fAttrAutopad(attr.auto_pad),
        fAttrDilations(attr.dilations), fAttrKernelShape(attr.kernel_shape), fAttrPads(attr.pads), fAttrStrides(attr.strides),
        fNX(UTILITY::Clean_name(nameX)), fNY(UTILITY::Clean_name(nameY))
   {
      if(std::is_same<T, float>::value) {
         fType = "float";
      } else {
         throw
            std::runtime_error("TMVA SOFIE Encountered unsupported type parsing a Pool operator");
      }
      fInputTensorNames = { fNX };
      fOutputTensorNames = { fNY };
   }
 
   // return input type (defined abstract in ROperator class )
   std::vector<ETensorType> TypeInference(std::vector<ETensorType> input) override {
      // only one input in Pool operators
      return input;
   }
 
   // function returning output shape given input
   std::vector<std::vector<size_t>> ShapeInference(std::vector<std::vector<size_t>> input) override {
      // shape of pooling input has to be (according to ONNX): NxCxHxW
      // Where N is batch size, C : input  channels, H : input height, W = input width
      // or it can be [N, C, F1,F2,....FN] . Minimum dimension is 3
      if (input.size() != 1 ) {
         throw std::runtime_error("TMVA SOFIE" + Name() + "Op Shape inference need 1 input tensor");
      }
      if (input[0].size() < 3) {
         throw std::runtime_error("TMVA SOFIE" + Name() + "Op Shape inference only accept tensor with at least 3 dimensions");
      }
      // support only input tensors with dim = 3,4,5
      if (input[0].size() < 3 || input[0].size() >  5) {
         throw std::runtime_error("TMVA SOFIE" + Name() + "Op : tensors with dimension " + std::to_string(input[0].size()) + " are not yet supported");
      }
 
      if (input[0].size() -2 != fDim) {
         throw
            std::runtime_error("TMVA SOFIE Pool Op Shape inference - invalid inputs ");
      }
       // kernel shape
      size_t k1 = ((fAttrKernelShape.empty())? input[0][2] : fAttrKernelShape[0]);
      size_t k2 = (fDim > 1) ? ((fAttrKernelShape.empty()) ? input[0][3] : fAttrKernelShape[1]) : 1;
      size_t k3 = (fDim > 2) ? ((fAttrKernelShape.empty()) ? input[0][4] : fAttrKernelShape[2]) : 1;
 
 
      size_t i1 = (fDim > 1) ? ((fDim > 2) ? 3 : 2) : 1;
      size_t i2 = (fDim > 2) ? 4 : 3;
      size_t i3 = 5;
 
      if (fAttrDilations.empty()) {
         fAttrDilations = {1, 1, 1};
      }
      fAttrDilations.resize(3);
      if (fDim < 3) {
         fAttrDilations.resize(3, 1);
      }
           // Shape of the kernel
      fAttrKernelShape = {k1 + (fAttrDilations[0] - 1) * (k1 - 1),
                          k2 + (fAttrDilations[1] - 1) * (k2 - 1),
                          k3 + (fAttrDilations[2] - 1) * (k3 - 1)};
 
      if (fAttrStrides.empty()) {
         fAttrStrides = {1, 1, 1};
      }
      if (fDim < 3)
         fAttrStrides.resize(3, 1);
 
      if (fAttrAutopad == "NOTSET") {
         // in auto_pad is NOTSET then fAttrPads should have been set or default zero is used
         if (fAttrPads.empty()) {
            fAttrPads = {0, 0, 0, 0, 0, 0};
         }
      } else if (fAttrAutopad == "SAME_UPPER" || fAttrAutopad == "SAME_LOWER") {
         // ONNX SAME padding: total_pad = max(0, (ceil(in/stride)-1)*stride + kernel - in)
         // SAME_UPPER places extra padding at end, SAME_LOWER at beginning
         fAttrPads.assign(6, 0);
         for (size_t d = 0; d < fDim; ++d) {
            size_t inSize = input[0][d + 2];
            size_t stride_d = fAttrStrides[d];
            size_t outSize = (inSize + stride_d - 1) / stride_d;
            int totalPad = std::max(0, (int)((outSize - 1) * stride_d + fAttrKernelShape[d]) - (int)inSize);
            if (fAttrAutopad == "SAME_UPPER") {
               fAttrPads[d] = (size_t)(totalPad / 2);
               fAttrPads[d + fDim] = (size_t)(totalPad - totalPad / 2);
            } else {
               fAttrPads[d] = (size_t)(totalPad - totalPad / 2);
               fAttrPads[d + fDim] = (size_t)(totalPad / 2);
            }
         }
      } else if (fAttrAutopad != "VALID") {
         throw
            std::runtime_error("TMVA SOFIE" + Name() + "Op invalid Autopad value : " + fAttrAutopad);
      }
      // to be sure pad is vector of size 6
      if (fDim < 3) fAttrPads.resize(6, 0);
 
      size_t input1 = input[0][2];
      size_t input2 = (fDim > 1) ? input[0][3] : 1;
      size_t input3 = (fDim > 2) ? input[0][4] : 1;
 
      // use ceiling division when ceil_mode=1, floor otherwise. With ceil_mode the rounding up can add
      // a window that starts past the end of the input (i.e. entirely in the right padding or in the
      // overhang region); ONNX ignores such a window, so clip to the number of valid window starts.
      auto poolOutDim = [this](size_t in, size_t padBegin, size_t padEnd, size_t kern, size_t stride) -> size_t {
         size_t n = in + padBegin + padEnd - kern;
         if (!fAttrCeilMode)
            return n / stride + 1;
         return std::min((n + stride - 1) / stride, (in - 1 + padBegin) / stride) + 1;
      };
 
      size_t output1 = poolOutDim(input1, fAttrPads[0], fAttrPads[i1], fAttrKernelShape[0], fAttrStrides[0]);
 
      size_t batch_size = input[0][0];        // first element in input tensor
      size_t output_channels = input[0][1];   // first element in output tensor
 
      std::vector<std::vector<size_t>> ret({{ batch_size, output_channels, output1 }});
 
      if (fDim == 1)
         return ret;
 
      size_t output2 = poolOutDim(input2, fAttrPads[1], fAttrPads[i2], fAttrKernelShape[1], fAttrStrides[1]);
      // output is N x C x OH x OW
      ret[0].push_back(output2);
      if (fDim == 2)
         return ret;
 
      size_t output3 = poolOutDim(input3, fAttrPads[2], fAttrPads[i3], fAttrKernelShape[2], fAttrStrides[2]);
 
      // output is N x C x OH x OW x OD
      ret[0].push_back(output3);
      return ret;
   }
 
   void Initialize(RModel& model) override {
 
      fUseSession = model.UseSession();
 
      if (!model.CheckIfTensorAlreadyExist(fNX)) {
         throw
            std::runtime_error("TMVA SOFIE Pool op Input Tensor " + fNX + " is not found in model");
      }
      fShapeX = model.GetTensorShape(fNX);
      if (fShapeX.size() < 3 || fShapeX.size()  > 5) {
         std::cout << fNX << " : " << ConvertShapeToString(fShapeX) << std::endl;
         throw
            std::runtime_error("TMVA SOFIE Pool Op input data tensor" + fNX + " is not of 3,4 or 5 dimensions");
      }
       fDim = fShapeX.size() - 2;
      // case of GlobalAveragePool. It is a pool case with kernel shape == image shape
      if (fPoolMode == GlobalAveragePool) {
         fPoolMode = AveragePool;
         fAttrKernelShape.resize(3);
         fAttrKernelShape[0] = fShapeX[2];
         if (fDim > 1)
            fAttrKernelShape[1] = fShapeX[3];
         if (fDim > 2)
            fAttrKernelShape[2] = fShapeX[4];
         fAttrAutopad = "VALID";
         fAttrPads = {0, 0, 0, 0, 0, 0 };
         assert(fAttrStrides.empty());
      }
      // find shape of Y and add it in the list of intermediate tensors
      fShapeY = ShapeInference({fShapeX})[0];
      model.AddIntermediateTensor(fNY, model.GetTensorType(fNX), fShapeY);
 
      // need cmath for INFINITY when using MaxPool
      if (fPoolMode == MaxPool) model.AddNeededStdLib("cmath");
 
   }
 
   std::string GenerateInitCode() override {
      std::stringstream out;
      return out.str();
   }
 
 
   std::string Generate(std::string OpName) override {
      OpName = "op_" + OpName;
 
      if (fShapeX.empty() || fShapeY.empty()) {
         throw std::runtime_error("TMVA SOFIE Pool Op called to Generate without being initialized first");
      }
 
      std::stringstream out;
 
      out << "\n//----  operator " << Name() << "  " << OpName << "\n";
      out << "{\n"; // create a new scope to avoid name clash
 
      assert(fShapeX[0] == fShapeY[0]);
      assert(fShapeX[1] == fShapeY[1]);
      assert(fAttrPads.size() == 6);
      assert(fAttrKernelShape.size() == 3);
      // A window must start inside the input: with ceil_mode the loop bound below can otherwise run one
      // window too far, which ONNX ignores. This mirrors the clipping done in ShapeInference.
      auto clipToInput = [this](int upper, size_t size) {
         return (fAttrCeilMode && upper > (int)size) ? (int)size : upper;
      };
      // find lower bounds of filtered area
      int hmin = - fAttrPads[0];   // minimum lower bound value of filter area
      // use stride instead of 1 when ceil_mode=1, so the loop covers the extra partial window
      int hmax = clipToInput(fShapeX[2] + fAttrPads[fDim] - fAttrKernelShape[0] + (fAttrCeilMode ? (int)fAttrStrides[0] : 1),
                             fShapeX[2]);
      int wmin,wmax,dmin,dmax;
 
      if(fDim >= 2){
         wmin = -fAttrPads[1]; // minimum lower bound value of filter area
         wmax = clipToInput(fShapeX[3] + fAttrPads[fDim + 1] - fAttrKernelShape[1] + (fAttrCeilMode ? (int)fAttrStrides[1] : 1),
                            fShapeX[3]);
      }
      else{
         wmin=1;
         wmax=1;
      }
      if(fDim == 3){
         dmin = -fAttrPads[2]; // minimum lower bound value of filter area
         dmax = clipToInput(fShapeX[4] + fAttrPads[fDim + 2] - fAttrKernelShape[2] + (fAttrCeilMode ? (int)fAttrStrides[2] : 1),
                            fShapeX[4]);
      }
      else{
         dmin=1;
         dmax=1;
      }
      out << SP << "constexpr int hsize = " << fShapeX[2] << ";\n";
      out << SP << "constexpr int hmin = " << hmin << ";\n";
      out << SP << "constexpr int hmax = " << hmax << ";\n";
      out << SP << "constexpr int kh = " << fAttrKernelShape[0] << ";\n";
      if (fDim > 1) {
         size_t wsize = fShapeX[3];
         out << SP << "constexpr int wsize = " << wsize << ";\n";
         out << SP << "constexpr int wmin = " << wmin << ";\n";
         out << SP << "constexpr int wmax = " << wmax << ";\n";
         out << SP << "constexpr int kw = " << fAttrKernelShape[1] << ";\n";
         if (fDim > 2) {
            size_t dsize = fShapeX[4];
            out << SP << "constexpr int dsize = " << dsize << ";\n";
            out << SP << "constexpr int dwsize = " << dsize*wsize << ";\n"; // hstride
            out << SP << "constexpr int dmin = " << dmin << ";\n";
            out << SP << "constexpr int dmax = " << dmax << ";\n";
            out << SP << "constexpr int kd = " << fAttrKernelShape[2] << ";\n";
         }
      }
 
 
      bool doPadding = false;
      for ( auto & e : fAttrPads)
         doPadding |= (e > 0);
 
      // An AveragePool window can cover fewer cells than the kernel area either because it overlaps the
      // padding region or because ceil_mode lets the last window overhang the input. In both cases the
      // divisor has to be computed per window instead of being the constant kernel area.
      bool dynamicDivisor = doPadding || fAttrCeilMode;
      // Number of cells the window starting at "var" covers along one dimension: the kernel extent
      // clipped to [0, size) when count_include_pad = 0, and to the padded input [-padBegin, size +
      // padEnd) otherwise. The lower bound needs no clipping in the latter case, since "var" starts at
      // -padBegin. Note that the cells the window overhangs past the padded input are never counted.
      auto windowExtent = [this](const std::string &var, const std::string &kern, size_t size, size_t padEnd) {
         std::string hi = std::to_string(fAttrCountIncludePad ? size + padEnd : size);
         std::string lo = fAttrCountIncludePad ? var : "(" + var + " > 0 ? " + var + " : 0)";
         return "((" + var + " + " + kern + " < " + hi + " ? " + var + " + " + kern + " : " + hi + ")"
                " - " + lo + ")";
      };
 
 
      if(fDim==1){
         // loop on batches and channels
         out << SP << "size_t outIndex = 0;\n";
         out << SP << "for (size_t n = 0; n < " << fShapeX[0]*fShapeX[1] << "; n++) {\n";
         out << SP << SP << "size_t inputOffset = n*" << fShapeX[2] << ";\n";
         out << SP << SP << "for (int i = hmin; i < hmax; i+=" << fAttrStrides[0] << ") {\n";
         // loop on elements of filter region to compute maximum
         if (fPoolMode == MaxPool)
            out << SP << SP << SP << SP << "float value = -INFINITY;\n";
         else if (fPoolMode == AveragePool) {
            out << SP << SP << SP << SP << "float value = 0;\n";
            if (dynamicDivisor)
               out << SP << SP << SP << SP << "const int nsum = "
                   << windowExtent("i", "kh", fShapeX[2], fAttrPads[fDim]) << ";\n";
            else // every window is full, so the divisor is the kernel area
               out << SP << SP << SP << SP << "constexpr int nsum = kh;\n";
         }
         // loop on rows of filtered region
         out << SP << SP << SP << SP  << "for (int l = i;  l < i + kh; l++) {\n";
         out << SP << SP << SP << SP  << SP << "if (l < 0 || l >= hsize) continue;\n";
         out << SP << SP << SP << SP  << SP << SP <<  "int index = inputOffset + l;\n";
         if (fPoolMode == MaxPool) {
         out << SP << SP << SP << SP << SP << SP << "auto xval = tensor_" << fNX << "[index];\n";
         out << SP << SP << SP << SP << SP << SP << "if (xval > value) value = xval;\n";
         }
         else if (fPoolMode == AveragePool) {
            // compute sum of values
            out << SP << SP << SP << SP << SP << SP << "value += tensor_" << fNX << "[index];\n";
         }
          out << SP << SP << SP << SP << SP << "}\n"; // end loop on region elements
         if (fPoolMode == AveragePool) {
         // compute average
         out << SP << SP << SP << SP << "value /= float(nsum);\n";
         }
 
         out << SP << SP << SP << SP << "tensor_" << fNY << "[outIndex++] = value;\n";
 
         out << SP << SP << "}\n";   // end loop on i (image rows)
         out << SP << "}\n";  // end loop on c*b
      }
      else if(fDim==2){
         // loop on batches and channels
         out << SP << "size_t outIndex = 0;\n";
         out << SP << "for (size_t n = 0; n < " << fShapeX[0]*fShapeX[1] << "; n++) {\n";
         out << SP << SP << "size_t inputOffset = n*" << fShapeX[2]*fShapeX[3] << ";\n";
         out << SP << SP << "for (int i = hmin; i < hmax; i+=" << fAttrStrides[0] << ") {\n";
         out << SP << SP << SP << "for (int j = wmin; j < wmax; j+=" << fAttrStrides[1] << ") {\n";
         // loop on elements of filter region to compute maximum
         if (fPoolMode == MaxPool)
            out << SP << SP << SP << SP << "float value = -INFINITY;\n";
         else if (fPoolMode == AveragePool) {
            out << SP << SP << SP << SP << "float value = 0;\n";
            if (dynamicDivisor)
               out << SP << SP << SP << SP << "const int nsum = "
                   << windowExtent("i", "kh", fShapeX[2], fAttrPads[fDim]) << " * "
                   << windowExtent("j", "kw", fShapeX[3], fAttrPads[fDim + 1]) << ";\n";
            else // every window is full, so the divisor is the kernel area
               out << SP << SP << SP << SP << "constexpr int nsum = kw*kh;\n";
         }
         // loop on rows of filtered region
         out << SP << SP << SP << SP << "for (int l = i;  l < i + kh; l++) {\n";
         out << SP << SP << SP << SP << SP << "if (l < 0 || l >= hsize) continue;\n";
         // loop on columns of filtered region
         out << SP << SP << SP << SP << SP << "for (int m = j; m < j + kw; m++) {\n";
         out << SP << SP << SP << SP << SP << SP << "if (m < 0 || m >= wsize) continue;\n";
         out << SP << SP << SP << SP << SP << SP << SP << "int index = inputOffset + l*wsize + m;\n";
         if (fPoolMode == MaxPool) {
         out << SP << SP << SP << SP << SP << SP << SP << "auto xval = tensor_" << fNX << "[index];\n";
         out << SP << SP << SP << SP << SP << SP << SP << "if (xval > value) value = xval;\n";
         }
         else if (fPoolMode == AveragePool) {
            // compute sum of values
            out << SP << SP << SP << SP << SP << SP << SP << "value += tensor_" << fNX << "[index];\n";
         }
         out << SP << SP << SP << SP << SP << SP << "}\n";
         out << SP << SP << SP << SP << SP << "}\n"; // end loop on region elements
         if (fPoolMode == AveragePool) {
            // compute average
            out << SP << SP << SP << SP << "value /= float(nsum);\n";
         }
         out << SP << SP << SP << SP << "tensor_" << fNY << "[outIndex++] = value;\n";
         out << SP << SP << SP << "}\n";   // end loop on j (columns of image)
         out << SP << SP << "}\n";   // end loop on i (image rows)
         out << SP << "}\n";  // end loop on c*b
      }
      else if(fDim==3){
         // loop on batches and channels
         out << SP << "size_t outIndex = 0;\n";
         out << SP << "for (size_t n = 0; n < " << fShapeX[0]*fShapeX[1] << "; n++) {\n";
         out << SP << SP << "size_t inputOffset = n*" << fShapeX[2]*fShapeX[3]*fShapeX[4] << ";\n";
         out << SP << SP << "for (int i = hmin; i < hmax; i+=" << fAttrStrides[0] << ") {\n";
         out << SP << SP << SP << "for (int j = wmin; j < wmax; j+=" << fAttrStrides[1] << ") {\n";
         out << SP << SP << SP << SP << "for (int k = dmin; k < dmax; k+=" << fAttrStrides[2] << ") {\n";
         // loop on elements of filter region to compute maximum
         if (fPoolMode == MaxPool)
            out << SP << SP << SP << SP << "float value = -INFINITY;\n";
         else if (fPoolMode == AveragePool) {
            out << SP << SP << SP << SP << "float value = 0;\n";
            if (dynamicDivisor)
               out << SP << SP << SP << SP << "const int nsum = "
                   << windowExtent("i", "kh", fShapeX[2], fAttrPads[fDim]) << " * "
                   << windowExtent("j", "kw", fShapeX[3], fAttrPads[fDim + 1]) << " * "
                   << windowExtent("k", "kd", fShapeX[4], fAttrPads[fDim + 2]) << ";\n";
            else // every window is full, so the divisor is the kernel area
               out << SP << SP << SP << SP << "constexpr int nsum = kw*kh*kd;\n";
         }
         // loop on rows of filtered region
         out << SP << SP << SP << SP  << "for (int l = i;  l < i + kh; l++) {\n";
         out << SP << SP << SP << SP  << SP << "if (l < 0 || l >= hsize) continue;\n";
         // loop on columns of filtered region
         out << SP << SP << SP << SP << SP << "for (int m = j; m < j + kw; m++) {\n";
         out << SP << SP << SP << SP << SP << SP << "if (m < 0 || m >= wsize) continue;\n";
         // loop on layers of filtered region
         out << SP << SP << SP << SP << SP << SP << "for (int p = k; p < k + kd; p++) {\n";
         out << SP << SP << SP << SP << SP << SP << SP << "if (p < 0 || p >= dsize) continue;\n";
         out << SP << SP << SP << SP << SP << SP << SP << SP << "int index = inputOffset + l*dwsize + m*dsize + p;\n";
 
         if (fPoolMode == MaxPool) {
            out << SP << SP << SP << SP << SP << SP << SP << SP << "auto xval = tensor_" << fNX << "[index];\n";
            out << SP << SP << SP << SP << SP << SP << SP << SP << "if (xval > value) value = xval;\n";
         }
         else if (fPoolMode == AveragePool) {
            // compute sum of values
            out << SP << SP << SP << SP << SP << SP << SP << SP << "value += tensor_" << fNX << "[index];\n";
         }
         out << SP << SP << SP << SP << SP << SP << "}\n";
         out << SP << SP << SP << SP << SP << "}\n";
         out << SP << SP << SP << SP << "}\n"; // end loop on region elements
         if (fPoolMode == AveragePool) {
            // compute average
            out << SP << SP << SP << SP << "value /= float(nsum);\n";
         }
 
         out << SP << SP << SP << SP << "tensor_" << fNY << "[outIndex++] = value;\n";
         out << SP << SP << SP << SP << "}\n" ; // end loop on k (layers of image)
         out << SP << SP << SP << "}\n";   // end loop on j (columns of image)
         out << SP << SP << "}\n";   // end loop on i (image rows)
         out << SP << "}\n";  // end loop on c*b
      }
      // end scope
      out << SP << "}\n";
 
 
      return out.str();
   }
};
 
} // namespace SOFIE
} // namespace Experimental
} // namespace TMVA
 
 
#endif