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