root/TMVA/ROperator_BatchNormalization.hxx

0001 #ifndef TMVA_SOFIE_ROPERATOR_BatchNormalization
0002 #define TMVA_SOFIE_ROPERATOR_BatchNormalization
0003
0004 #include "SOFIE_common.hxx"
0005 #include "ROperator.hxx"
0006 #include "RModel.hxx"
0007
0008
0009 #include <cmath>
0010 #include <sstream>
0011
0012 namespace TMVA{
0013 namespace Experimental{
0014 namespace SOFIE{
0015
0016 template <typename T>
0017 class ROperator_BatchNormalization final : public ROperator
0018 {
0019
0020 private:
0021
0022    /* Attributes */
0023    float fepsilon = 1e-05;
0024    float fmomentum = 0.9;
0025    std::size_t ftraining_mode = 0;
0026
0027    std::string fNX;
0028    std::string fNScale;
0029    std::string fNB;
0030    std::string fNMean;
0031    std::string fNVar;
0032    std::string fNY;
0033    EActivationType fActivation;
0034
0035    std::vector<size_t> fShapeX;
0036    std::vector<size_t> fShapeScale;
0037    std::vector<size_t> fShapeB;
0038    std::vector<size_t> fShapeMean;
0039    std::vector<size_t> fShapeVar;
0040    std::vector<size_t> fShapeY;
0041
0042    std::string fType;
0043
0044 public:
0045    ROperator_BatchNormalization() = delete;
0046
0047    /* Constructor */
0048    ROperator_BatchNormalization( float epsilon, float momentum, std::size_t training_mode,
0049    std::string nameX, std::string nameScale, std::string nameB,
0050    std::string nameMean, std::string nameVar, std::string nameY, EActivationType activation=EActivationType::UNDEFINED):
0051    fepsilon(epsilon), fmomentum(momentum), ftraining_mode(training_mode),
0052    fNX(UTILITY::Clean_name(nameX)), fNScale(UTILITY::Clean_name(nameScale)),
0053    fNB(UTILITY::Clean_name(nameB)), fNMean(UTILITY::Clean_name(nameMean)),
0054    fNVar(UTILITY::Clean_name(nameVar)), fNY(UTILITY::Clean_name(nameY)), fActivation(activation)
0055    {
0056       fInputTensorNames = { fNX };
0057       fOutputTensorNames = { fNY };
0058
0059       if(std::is_same<T, float>::value){
0060       fType = "float";
0061       }
0062       else{
0063           throw
0064               std::runtime_error("TMVA SOFIE Encountered unsupported type parsing a BatchNormalization operator");
0065       }
0066    }
0067
0068
0069    std::vector<ETensorType> TypeInference(std::vector<ETensorType> input) override {
0070       ETensorType out = input[0];
0071       return {out};
0072    }
0073
0074    std::vector<std::vector<size_t>> ShapeInference(std::vector<std::vector<size_t>> input) override {
0075       if (input.size() != 5 ) {
0076          throw
0077          std::runtime_error("TMVA SOFIE BatchNormalization Op Shape inference need 5 input tensors");
0078       }
0079       for(size_t i = 0; i < input.size(); i++) {
0080          if (input[i].size() != 4) {
0081             throw
0082             std::runtime_error("TMVA SOFIE BatchNormalization Op Shape inference only accept tensor with 4 dimensions");
0083          }
0084       }
0085
0086       auto ret = input;
0087       return ret;
0088    }
0089
0090    void Initialize(RModel& model) override {
0091       if (!model.CheckIfTensorAlreadyExist(fNX)) {
0092          throw
0093             std::runtime_error("TMVA SOFIE BatchNormalization op Input Tensor " + fNX + " fnx is not found in model");
0094       }
0095       if (!model.CheckIfTensorAlreadyExist(fNScale)) {
0096          throw
0097             std::runtime_error("TMVA SOFIE BatchNormalization op Input Tensor " + fNScale + " fns is not found in model");
0098       }
0099       if (!model.CheckIfTensorAlreadyExist(fNB)) {
0100          throw
0101             std::runtime_error("TMVA SOFIE BatchNormalization op Input Tensor " + fNB + " fnb is not found in model");
0102       }
0103       if (!model.CheckIfTensorAlreadyExist(fNMean)) {
0104          throw
0105             std::runtime_error("TMVA SOFIE BatchNormalization op Input Tensor " + fNMean + " fnm is not found in model");
0106       }
0107       if (!model.CheckIfTensorAlreadyExist(fNVar)) {
0108          throw
0109             std::runtime_error("TMVA SOFIE BatchNormalization op Input Tensor " + fNVar + " fnv is not found in model");
0110       }
0111
0112       fShapeX = model.GetTensorShape(fNX);
0113
0114       if (fShapeX.size() <  2 || fShapeX.size() > 4) {
0115          throw
0116             std::runtime_error("TMVA SOFIE BatchNormalization Op input tensor " + fNX + " fnx has wrong shape : " + ConvertShapeToString(fShapeX));
0117       }
0118
0119       fShapeScale = model.GetTensorShape(fNScale);
0120       fShapeB = model.GetTensorShape(fNB);
0121       fShapeMean = model.GetTensorShape(fNMean);
0122       fShapeVar = model.GetTensorShape(fNVar);
0123       fShapeY = fShapeX;
0124       model.AddIntermediateTensor(fNY, model.GetTensorType(fNX), fShapeY);
0125
0126       if (fShapeB.size() == 1) {
0127          // Broadcast scale, bias, input_mean and input_var to shape_X
0128          auto original_B = model.GetInitializedTensorData(fNB);
0129          auto original_S = model.GetInitializedTensorData(fNScale);
0130          auto original_M = model.GetInitializedTensorData(fNMean);
0131          auto original_V = model.GetInitializedTensorData(fNVar);
0132          size_t batchSize = fShapeX[0];
0133          size_t channels = fShapeX[1];
0134          size_t height = (fShapeX.size() > 2) ? fShapeX[2] : 1;
0135          size_t width = (fShapeX.size() > 3) ? fShapeX[3] : 1;
0136          size_t n = batchSize * channels * height * width;
0137          if (fType == "float") {
0138             float *original_bias = static_cast<float *>(original_B.get());
0139             float *original_scale = static_cast<float *>(original_S.get());
0140             float *original_mean = static_cast<float *>(original_M.get());
0141             float *original_var = static_cast<float *>(original_V.get());
0142             float *new_bias = new float[n];
0143             float *new_scale = new float[n];
0144             float *new_mean = new float[n];
0145             float *new_var = new float[n];
0146             size_t bs = 0, ch = 0, h = 0, w = 0;
0147             for (ch = 0; ch < channels; ch++) {
0148                for (h = 0; h < height; h++) {
0149                   for (w = 0; w < width; w++) {
0150                      new_bias[bs * channels * height * width + ch * height * width + h * width + w] = original_bias[ch];
0151                      new_scale[bs * channels * height * width + ch * height * width + h * width + w] =
0152                         original_scale[ch];
0153                      new_mean[bs * channels * height * width + ch * height * width + h * width + w] = original_mean[ch];
0154                      new_var[bs * channels * height * width + ch * height * width + h * width + w] = original_var[ch];
0155                   }
0156                }
0157             }
0158             size_t Batchoffset = channels * height * width;
0159             for (bs = 1; bs < batchSize; bs++) {
0160                std::copy(new_bias, new_bias + Batchoffset, new_bias + (bs * Batchoffset));
0161                std::copy(new_scale, new_scale + Batchoffset, new_scale + (bs * Batchoffset));
0162                std::copy(new_mean, new_mean + Batchoffset, new_mean + (bs * Batchoffset));
0163                std::copy(new_var, new_var + Batchoffset, new_var + (bs * Batchoffset));
0164             }
0165             //// new_var =1. / sqrt(input_var + fepsilon)
0166             for (size_t i = 0; i < n; i++) {
0167                new_var[i] = 1. / sqrt(new_var[i] + fepsilon);
0168                new_scale[i] *= new_var[i]; // include var in new scale
0169             }
0170             std::vector<size_t> new_bias_shape = {batchSize, channels, height, width};
0171             std::shared_ptr<void> new_bias_ptr(new_bias, std::default_delete<float[]>());
0172             std::shared_ptr<void> new_scale_ptr(new_scale, std::default_delete<float[]>());
0173             std::shared_ptr<void> new_mean_ptr(new_mean, std::default_delete<float[]>());
0174             std::shared_ptr<void> new_var_ptr(new_var, std::default_delete<float[]>());
0175             model.UpdateInitializedTensor(fNB, model.GetTensorType(fNB), new_bias_shape, new_bias_ptr);
0176             model.UpdateInitializedTensor(fNScale, model.GetTensorType(fNScale), new_bias_shape, new_scale_ptr);
0177             model.UpdateInitializedTensor(fNMean, model.GetTensorType(fNMean), new_bias_shape, new_mean_ptr);
0178             model.UpdateInitializedTensor(fNVar, model.GetTensorType(fNVar), new_bias_shape, new_var_ptr);
0179             fShapeB = model.GetTensorShape(fNB);
0180             fShapeScale = model.GetTensorShape(fNScale);
0181             fShapeMean = model.GetTensorShape(fNMean);
0182             fShapeVar = model.GetTensorShape(fNVar);
0183          }
0184       }
0185    }
0186
0187    std::string Generate(std::string OpName) override {
0188       OpName = "op_" + OpName;
0189       if (fShapeX.empty()){
0190          throw std::runtime_error("TMVA SOFIE Batch Normalization called to Generate without being initialized first");
0191       }
0192
0193       std::stringstream out;
0194       //// Batch Norm op
0195       size_t batchSize = fShapeX[0];
0196       size_t channels = fShapeX[1];
0197       size_t height = (fShapeX.size() > 2) ? fShapeX[2] : 1;
0198       size_t width = (fShapeX.size() > 3) ? fShapeX[3] : 1;
0199       size_t n = batchSize * channels * height * width;
0200
0201       //// copy X into Y
0202       out << "\n\n//---- BatchNorm\n";
0203       out << SP << "constexpr int " << OpName << "_N =" << batchSize * channels * height * width << ";\n";
0204       out << SP << "constexpr int "<<OpName<< "_incx = 1;\n";
0205       out << SP << "constexpr int "<<OpName<< "_incy = 1;\n";
0206       out << SP << "BLAS::scopy_(&" << OpName << "_N, " << "tensor_" << fNX << ", &" << OpName << "_incx," << "tensor_" << fNY << ", &" << OpName << "_incy);\n\n";
0207
0208       //// blas saxpy (Y = -Bmean + Y)
0209       out << SP << "float "<<OpName<< "_alpha = -1;\n";
0210       out << SP << "BLAS::saxpy_(&" << OpName << "_N, &" << OpName << "_alpha, " << "tensor_" << fNMean << ", &" << OpName << "_incx,"
0211          << "tensor_" << fNY <<", &" << OpName << "_incy);\n\n ";
0212
0213       //// Y *= scale*var
0214       out << SP << "for (size_t i = 0; i < " << n << "; i++) {\n";
0215       // scale tensor contains already the var
0216       out << SP << SP << "tensor_" << fNY << "[i] *= tensor_" << fNScale << "[i]; \n";
0217       out << SP << "}\n";
0218
0219       //// blas saxpy (Y = Bbias + Y)
0220       out << SP <<OpName<< "_alpha = 1;\n";
0221       out << SP << "BLAS::saxpy_(&" << OpName << "_N, &" << OpName << "_alpha, " << "tensor_" << fNB << ", &" << OpName << "_incx, "
0222          << "tensor_" << fNY << ", &" << OpName << "_incy);\n\n";
0223
0224       if(fActivation == EActivationType::RELU){
0225          out << SP << "for (int id = 0; id < " << ConvertShapeToLength(fShapeY) << " ; id++){\n";
0226          out << SP << SP << "tensor_" << fNY << "[id] = ((tensor_" << fNY << "[id] > 0 )? tensor_" << fNY << "[id] : 0);\n";
0227          out << SP << "}\n";
0228       }
0229       return out.str();
0230    }
0231
0232    std::vector<std::string> GetBlasRoutines() override { return { std::string("Copy"), std::string("Axpy") }; }
0233 };
0234
0235 }//SOFIE
0236 }//Experimental
0237 }//TMVA
0238
0239
0240 #endif //TMVA_SOFIE_ROPERATOR_BatchNormalization