root/TMVA/ROperator_Reshape.hxx

0001 #ifndef TMVA_SOFIE_ROPERATOR_RESHAPE
0002 #define TMVA_SOFIE_ROPERATOR_RESHAPE
0003
0004 #include "TMVA/SOFIE_common.hxx"
0005 #include "TMVA/ROperator.hxx"
0006 #include "TMVA/RModel.hxx"
0007
0008 #include <cassert>
0009 #include <cctype>
0010 #include <sstream>
0011 #include <algorithm>
0012
0013 namespace TMVA{
0014 namespace Experimental{
0015 namespace SOFIE{
0016
0017 enum ReshapeOpMode { Reshape, Flatten, Squeeze, Unsqueeze };
0018
0019
0020 class ROperator_Reshape final : public ROperator
0021 {
0022
0023 private:
0024
0025    bool fVerbose = false;
0026    bool fDimInput = false;
0027    bool fDynamicShape = false;
0028    ReshapeOpMode fOpMode = Reshape;   // type of Reshape operator
0029
0030    int fAllowZero = 0; // (for Reshape) zero in tensor shape makes output shape equal to input tensor shape
0031    int fAxis = 1;      // (for Flatten)
0032
0033    std::string fNData;        // input data tensor name
0034    std::string fNInput2;       // reshape or axes tensor name depending on operator
0035    std::string fNOutput;               // output tensor name
0036    std::vector<Dim> fShapeInput;     // input shape data
0037    std::vector<Dim> fShapeOutput;   // output shape data
0038    std::vector<int64_t> fAttrAxes;         // axes attributes (provided for all version of Squeeze/Unsqueeze)
0039    std::vector<int64_t> fShape;     // shape tensor values provided for Reshape
0040
0041 public:
0042
0043    std::string Name() const {
0044       if (fOpMode == Reshape) return "Reshape";
0045       if (fOpMode == Flatten) return "Flatten";
0046       if (fOpMode == Squeeze) return "Squeeze";
0047       if (fOpMode == Unsqueeze) return "Unsqueeze";
0048       return "";
0049    }
0050
0051    ROperator_Reshape(){}
0052    ROperator_Reshape(ReshapeOpMode opMode, int attr_value, std::string nameData, std::string nameInput2, std::string nameOutput)
0053       : fOpMode(opMode), fNData(UTILITY::Clean_name(nameData)), fNInput2(UTILITY::Clean_name(nameInput2)),
0054          fNOutput(UTILITY::Clean_name(nameOutput))
0055    {
0056       if (opMode == Reshape) fAllowZero = attr_value;
0057       if (opMode == Flatten) fAxis = attr_value;
0058
0059       fInputTensorNames = { fNData };
0060       if(!fNInput2.empty()){
0061          fInputTensorNames.emplace_back(fNInput2);
0062       }
0063       fOutputTensorNames = { fNOutput };
0064    }
0065
0066    // for squeeze/unsqueezed operators following old ONNX version (< 10)
0067    // In this cases axes are passed as attribute values
0068    ROperator_Reshape(ReshapeOpMode opMode, std::vector<int64_t> attrAxes, std::string nameData, std::string nameOutput)
0069       : fOpMode(opMode), fNData(UTILITY::Clean_name(nameData)), fNOutput(UTILITY::Clean_name(nameOutput)),
0070         fAttrAxes(attrAxes)
0071    {
0072       assert(fOpMode == Squeeze || fOpMode == Unsqueeze);
0073       fInputTensorNames = { fNData };
0074       fOutputTensorNames = { fNOutput };
0075    }
0076
0077    // output type is same as input
0078    std::vector<ETensorType> TypeInference(std::vector<ETensorType> input) override {
0079       auto ret = std::vector<ETensorType>(1, input[0]);
0080       return ret;
0081    }
0082    std::vector<std::vector<size_t>> ShapeInference(std::vector<std::vector<size_t>> input) override {
0083       return input;
0084    }
0085
0086    // output shape
0087    std::vector<std::vector<Dim>> ShapeInference(const std::vector<std::vector<Dim>> & input)  {
0088       std::vector<std::vector<Dim>> ret;
0089       auto & input_shape = input[0];
0090       if (fOpMode == Reshape) {
0091          // correct the provided shape (here we have the value) for 0 or -1
0092          std::vector<Dim> output_shape(fShape.size());
0093          assert(!fShape.empty() && !fDynamicShape);
0094          for (size_t i = 0; i < output_shape.size(); i++) {
0095             if (fShape[i] > 0 || (fAllowZero && fShape[i] >= 0))
0096                output_shape[i] = Dim{ static_cast<size_t>(fShape[i]) };
0097             else if (!fAllowZero && fShape[i] == 0)
0098                output_shape[i] = input_shape[i];
0099          }
0100          // now case of -1 in shape
0101          for (size_t i = 0; i < output_shape.size(); i++) {
0102             if (fShape[i] == -1) {
0103                auto tmp = output_shape;
0104                tmp.erase(tmp.begin() + i);
0105                auto tmp_length = ConvertDimShapeToLength(tmp);
0106                auto input_length = ConvertDimShapeToLength(input_shape);
0107                if (fVerbose)
0108                   std::cout << "reshape- try simplifying " << ConvertDimShapeToString(input_shape) << " with length "
0109                             << input_length << " to " << tmp_length << std::endl;
0110
0111                if (IsInteger(tmp_length) && IsInteger(input_length))
0112                   output_shape[i] = Dim{static_cast<size_t>(std::stoi(input_length) / std::stoi(tmp_length))};
0113                else if (IsInteger(tmp_length) && std::stoi(tmp_length) == 1) {
0114                   output_shape[i] = Dim{input_length, static_cast<size_t>(-1)};
0115                }
0116                else {
0117                   //we can try simplifying expression if tmp_length is integer and part of input_length
0118                   // contains tmp_length
0119                   bool canSimplify = false;
0120                   std::vector <Dim> reduced_input;
0121                   if (IsInteger(tmp_length)) {
0122
0123                      // try to tokenize with * the input length
0124
0125                      std::stringstream ss(input_length);
0126
0127                      std::string token;
0128
0129                      // Tokenizing w.r.t. space '*'
0130                      while(getline(ss, token, '*'))
0131                      {
0132                         // remove any whitespace
0133                         token.erase(std::remove_if(token.begin(), token.end(),
0134                                                    [](unsigned char x) { return std::isspace(x); }), token.end());
0135                         if (token != tmp_length) {
0136                            if (IsInteger(token)) {
0137                                  size_t il = static_cast<size_t>(std::stoi(input_length));
0138                                  size_t tl = static_cast<size_t>(std::stoi(tmp_length));
0139                                  if ((il % tl) == 0) {
0140                                  canSimplify = true;
0141                                  reduced_input.push_back(Dim{il / tl});
0142                                  }
0143                            } else {
0144                               reduced_input.push_back(Dim{token});
0145                            }
0146                         } else {
0147                            // token is equal to tmp_length, can be not considered and is simplified
0148                            canSimplify = true;
0149                         }
0150                      }
0151                   }
0152                   if (canSimplify) {
0153                      // if length contains * we need to add some brackets
0154                      std::string res_shape = ConvertDimShapeToLength(reduced_input);
0155                      if (res_shape.find('*') != std::string::npos)
0156                         output_shape[i] = Dim{std::string("(") + res_shape + ")", static_cast<size_t>(-1)};
0157                      else
0158                         output_shape[i] = Dim{res_shape};
0159                   }
0160                   if (!canSimplify)
0161                      output_shape[i] = Dim{std::string("(") + input_length + " / (" + tmp_length + "))", static_cast<size_t>(-1)};
0162                }
0163
0164                break; // cannot have more than -1
0165             }
0166             //  throw std::runtime_error(
0167             //                   "TMVA Reshape Op : output shape has multiple negative or zero values");
0168          }
0169
0170          if (fVerbose)
0171             std::cout << "Reshape: correct output shape  to " << ConvertDimShapeToString(output_shape) << std::endl;
0172
0173          if (!fDimInput && ConvertDimShapeToLength(output_shape) != ConvertDimShapeToLength(input_shape)) {
0174             throw std::runtime_error("TMVA Reshape Op : Invalid  shapes : " + ConvertDimShapeToString(input_shape) +
0175                                      ConvertDimShapeToString(output_shape));
0176          }
0177          ret.push_back(output_shape);
0178
0179       } else if (fOpMode == Flatten) {
0180          // flatten case
0181          if (fAxis < 0)
0182             fAxis += input_shape.size();
0183          auto s1 = std::vector<Dim>(input_shape.begin(), input_shape.begin() + fAxis);
0184          auto s2 = std::vector<Dim>(input_shape.begin() + fAxis, input_shape.end());
0185          auto l1 = ConvertDimShapeToLength(s1);
0186          auto l2 = ConvertDimShapeToLength(s2);
0187          std::vector<Dim> newShape = {Dim{l1}, Dim{l2}};
0188          ret.push_back(newShape);
0189       } else if (fOpMode == Squeeze) {
0190          // squeeze
0191          // assume no axis is provided - remove all axes with value equal to 1
0192          auto output_shape = input_shape;
0193          if (fAttrAxes.empty()) {
0194             size_t i = 0;
0195             while (i < output_shape.size()) {
0196                if (output_shape[i] == Dim{1}) {
0197                   output_shape.erase(output_shape.begin() + i);
0198                } else {
0199                   i++;
0200                }
0201             }
0202          } else {
0203             std::cout << "getting shape for Squeeze...from attribute\n";
0204             auto axes = fAttrAxes;
0205             for (size_t i = 0; i < axes.size(); i++) {
0206                std::cout << i << "  " << axes[i] << std::endl;
0207                if (axes[i] < 0)
0208                   axes[i] += input_shape.size();
0209                if (!(output_shape[axes[i]] == Dim{1}))
0210                   throw std::runtime_error("TMVA Squeeze Op : Invalid  axis value " + std::to_string(axes[i]) +
0211                                            " for " + ConvertDimShapeToString(output_shape));
0212             }
0213             // for calling vector::erase we must sort axes in decreasing order to avoid
0214             std::sort(axes.begin(), axes.end(), std::greater<int>());
0215             for (auto & axis : axes) {
0216                std::cout << "erase give axis " << axis << "  -> ";
0217                for (auto & o : output_shape) std::cout << o << " , ";
0218                std::cout << std::endl;
0219                output_shape.erase(output_shape.begin() + axis);
0220             }
0221          }
0222          ret.push_back(output_shape);
0223       }
0224       else if (fOpMode == Unsqueeze) {
0225          // unsqueeze
0226          assert(!fAttrAxes.empty());
0227          auto output_shape = input_shape;
0228          auto &axes = fAttrAxes;
0229          // output rank
0230          int64_t r = input[0].size() + axes.size();
0231          for (auto &a : axes) {
0232             int64_t i = static_cast<int64_t>(a);
0233             if (i < -r || i > r - 1)
0234                throw std::runtime_error("TMVA Unsqueeze Op - axes input is not in correct range");
0235             if (i >= 0)
0236                output_shape.insert(output_shape.begin() + i, Dim{1});
0237             else
0238                // negative axes
0239                output_shape.insert(output_shape.end() + i + 1, Dim{1});
0240          }
0241          ret.push_back(output_shape);
0242       }
0243       return ret;
0244    }
0245
0246    void Initialize(RModel& model) override {
0247
0248       fVerbose = model.Verbose();
0249       if (fVerbose)
0250          std::cout << "initialize reshape op type " << fOpMode << " - " << fNInput2 << " " << fNData << std::endl;
0251
0252       if (model.CheckIfTensorAlreadyExist(fNData) == false) {
0253           // input must be a graph input, or already initialized intermediate tensor
0254          throw std::runtime_error("TMVA Reshape Op Input Tensor " + fNData + "  is not found in model");
0255       }
0256       fShapeInput = model.GetDimTensorShape(fNData);
0257       fDimInput = model.IsDynamicTensor(fNData);
0258       // check if optional tensor exists defining shape or axes
0259       if (!fNInput2.empty()) {
0260          if (model.CheckIfTensorAlreadyExist(fNInput2)) {
0261             if (model.IsInitializedTensor(fNInput2)) {
0262                // assume input shape is an initialized tensor
0263                auto dptr = model.GetInitializedTensorData(fNInput2);
0264                auto values = static_cast<int64_t *>(dptr.get());
0265                auto vec = model.GetTensorShape(fNInput2);
0266                size_t n = 1;
0267                if (vec.size() > 0)
0268                   n = vec[0]; // size of shape input tensor
0269                // copy values in fShape vector or fAttrAxes
0270                if (fOpMode == Reshape)
0271                   fShape = std::vector<int64_t>(values, values + n);
0272                else
0273                   fAttrAxes = std::vector<int64_t>(values, values + n);
0274
0275                fShapeOutput = ShapeInference({fShapeInput})[0];
0276                // set flag to not write tensor in weight file. Its data will be hard-coded in way model is constructed
0277                model.SetNotWritableInitializedTensor(fNInput2);
0278             } else if (model.IsShapeTensor(fNInput2)) {
0279                auto shapeData = model.GetShapeTensorValues(fNInput2);
0280                fShapeOutput = shapeData;
0281             } else {
0282                // we cannot get shape at initialization time but at run-time
0283                fDynamicShape = true;
0284                // size of shape output us given by size of shape input tensor
0285                auto shapeInput2 = model.GetTensorShape(fNInput2);
0286                fShapeOutput.resize(shapeInput2[0]);
0287                for (size_t i = 0; i < fShapeOutput.size(); i++) {
0288                   fShapeOutput[i] = Dim{ std::string("s_") + fNOutput + "_" + std::to_string(i)};
0289                }
0290             }
0291          } else {
0292             throw std::runtime_error("TMVA Reshape Op 2nd input Tensor " + fNInput2 + " is not found in model");
0293          }
0294       } else if (!fAttrAxes.empty()) {
0295          // case fNShape is empty and axes are provided as attributes (e.g. for Unsqueeze)
0296          std::cout << "attribute axes exists\n";
0297          fShapeOutput = ShapeInference({fShapeInput})[0];
0298       } else if (fOpMode == Flatten || fOpMode == Squeeze) {
0299          fShapeOutput = ShapeInference({fShapeInput})[0];
0300       } else {
0301          throw std::runtime_error("TMVA Reshape Op : Invalid Input/Attribute data");
0302       }
0303       // check if output is constant or not
0304       if (model.IsInitializedTensor(fNData) && model.GetTensorType(fNData) == ETensorType::INT64) {
0305          fIsOutputConstant = true;
0306          auto inputData = static_cast<int64_t*>(model.GetInitializedTensorData(fNData).get());
0307          auto o_shape = ConvertShapeToInt(fShapeOutput);
0308          if (ConvertShapeToLength(ConvertShapeToInt(fShapeInput)) != ConvertShapeToLength(o_shape) )
0309             throw std::runtime_error("TMVA Reshape Op : Invalid Input/Output lengths");
0310          model.AddConstantTensor<int64_t>(fNOutput, o_shape, inputData);
0311          if (model.Verbose()) {
0312             std::cout << Name() << " : " << fNData << " " << ConvertDimShapeToString(fShapeInput) << " -->  " << fNOutput << " (constant) " << ConvertDimShapeToString(fShapeOutput)  << " : " <<
0313             ConvertValuesToString(ConvertShapeToLength(o_shape), inputData) << std::endl;
0314          }
0315       }
0316       // for shape tensors we can have it if output shape is size==1 or a scalar
0317       else if (model.IsShapeTensor(fNData) && fShapeOutput.size() <=1) {
0318          fIsOutputConstant = true;
0319          auto inputData = model.GetShapeTensorValues(fNData);
0320          model.AddShapeTensor(fNOutput, inputData);
0321          if (model.Verbose()) {
0322             std::cout << Name() << " : " << fNData << " " << ConvertDimShapeToString(fShapeInput) << " -->  " << fNOutput << " (shape) " << ConvertDimShapeToString(fShapeOutput)  << " : " <<
0323             ConvertDimShapeToString(inputData) << std::endl;
0324          }
0325       }
0326       else {
0327          // non-constant case
0328          model.AddIntermediateTensor(fNOutput, model.GetTensorType(fNData), fShapeOutput);
0329          if (model.Verbose())
0330             std::cout << Name() << " : " << fNData << " " << ConvertDimShapeToString(fShapeInput) << " -->  "<< fNOutput << "  " << ConvertDimShapeToString(fShapeOutput)  << std::endl;
0331       }
0332    }
0333
0334    std::string Generate(std::string opName) override {
0335       if (fIsOutputConstant) return "";  //no op for constant tensors
0336
0337       std::stringstream out;
0338       std::string opType = "Reshape";
0339       if (fOpMode == Flatten)
0340          opType = "Flatten";
0341       else if (fOpMode == Squeeze)
0342          opType = "Squeeze";
0343       else if (fOpMode == Unsqueeze)
0344          opType = "Unsquueze";
0345
0346       out << SP << "///--------" << opType << " operator " << opName << " --> " << ConvertDimShapeToString(fShapeOutput) << "\n";
0347
0348       // in case of dynamic output shape we need to set the shape value from input shape tensor
0349       // and take case of the zero values
0350       if (fDynamicShape) {
0351          for (size_t i = 0; i < fShapeOutput.size(); i++) {
0352             // since fNInput2 values are int64_t, should we check if they are negative?
0353             out << SP << "size_t " << fShapeOutput[i].param << " = " << "tensor_" << fNInput2 << "[" << i << "];\n";
0354             if (!fAllowZero)
0355                out << SP << "if (tensor_" << fNInput2 << "[" << i << "] <= 0 ) "
0356                          <<  fShapeOutput[i].param << " = " <<  fShapeInput[i] << ";\n";
0357          }
0358       }
0359
0360       // output of reshape is same as input
0361       auto lengthOut = ConvertDimShapeToLength(fShapeOutput);
0362       auto lengthIn = ConvertDimShapeToLength(fShapeInput);
0363       if (lengthOut != lengthIn) {
0364          // check needs to be done at run-time
0365          out << SP << "if (" << lengthOut << "!=" << lengthIn << ")\n";
0366          out << "throw std::runtime_error(\"TMVA SOFIE Reshape Op : output lengths is different than input one\");\n";
0367       }
0368
0369
0370       out << SP << "std::copy( tensor_" << fNData << ", tensor_" << fNData << " + " << lengthIn << ", " << "tensor_" << fNOutput
0371           << ");\n";
0372       return out.str();
0373    }
0374 };
0375
0376 }//SOFIE
0377 }//Experimental
0378 }//TMVA
0379
0380
0381 #endif //TMVA_SOFIE_ROPERATOR_RESHAPE