diff --git a/tmva/sofie/inc/TMVA/OperatorList.hxx b/tmva/sofie/inc/TMVA/OperatorList.hxx index 9495719deeb4b..232c41584c4d3 100644 --- a/tmva/sofie/inc/TMVA/OperatorList.hxx +++ b/tmva/sofie/inc/TMVA/OperatorList.hxx @@ -2,3 +2,5 @@ #include "TMVA/ROperator_Gemm.hxx" #include "TMVA/ROperator_Relu.hxx" #include "TMVA/ROperator_Conv.hxx" +#include "ROperator_BatchNormalization.hxx" +#include "ROperator_InstanceNormalization.hxx" \ No newline at end of file diff --git a/tmva/sofie/inc/TMVA/ROperator_BatchNormalization.hxx b/tmva/sofie/inc/TMVA/ROperator_BatchNormalization.hxx new file mode 100644 index 0000000000000..d8f3509e1276b --- /dev/null +++ b/tmva/sofie/inc/TMVA/ROperator_BatchNormalization.hxx @@ -0,0 +1,222 @@ +#ifndef TMVA_SOFIE_ROPERATOR_BatchNormalization +#define TMVA_SOFIE_ROPERATOR_BatchNormalization + +#include "SOFIE_common.hxx" +#include "ROperator.hxx" +#include "RModel.hxx" + + +#include +#include + +namespace TMVA{ +namespace Experimental{ +namespace SOFIE{ + +template +class ROperator_BatchNormalization final : public ROperator +{ + +private: + + /* Attributes */ + float fepsilon = 1e-05; + float fmomentum = 0.9; + std::size_t ftraining_mode = 0; + + std::string fNX; + std::string fNScale; + std::string fNB; + std::string fNMean; + std::string fNVar; + std::string fNY; + + std::vector fShapeX; + std::vector fShapeScale; + std::vector fShapeB; + std::vector fShapeMean; + std::vector fShapeVar; + std::vector fShapeY; + + std::string fType; + +public: + ROperator_BatchNormalization() = delete; + + /* Constructor */ + ROperator_BatchNormalization( float epsilon, float momentum, std::size_t training_mode, + std::string nameX, std::string nameScale, std::string nameB, + std::string nameMean, std::string nameVar, std::string nameY): + fepsilon(epsilon), fmomentum(momentum), ftraining_mode(training_mode), + fNX(UTILITY::Clean_name(nameX)), fNScale(UTILITY::Clean_name(nameScale)), + fNB(UTILITY::Clean_name(nameB)), fNMean(UTILITY::Clean_name(nameMean)), + fNVar(UTILITY::Clean_name(nameVar)), fNY(UTILITY::Clean_name(nameY)) + { + if(std::is_same::value){ + fType = "float"; + } + else{ + throw + std::runtime_error("TMVA SOFIE Encountered unsupported type parsing a BatchNormalization operator"); + } + } + + + std::vector TypeInference(std::vector input) { + ETensorType out = input[0]; + return {out}; + } + + std::vector> ShapeInference(std::vector> input) { + if (input.size() != 5 ) { + throw + std::runtime_error("TMVA SOFIE BatchNormalization Op Shape inference need 5 input tensors"); + } + for(size_t i = 0; i < input.size(); i++) { + if (input[i].size() != 4) { + throw + std::runtime_error("TMVA SOFIE BatchNormalization Op Shape inference only accept tensor with 4 dimensions"); + } + } + + auto ret = input; + return ret; + } + + void Initialize(RModel& model){ + if (!model.CheckIfTensorAlreadyExist(fNX)) { + throw + std::runtime_error("TMVA SOFIE BatchNormalization op Input Tensor " + fNX + " fnx is not found in model"); + } + if (!model.CheckIfTensorAlreadyExist(fNScale)) { + throw + std::runtime_error("TMVA SOFIE BatchNormalization op Input Tensor " + fNScale + " fns is not found in model"); + } + if (!model.CheckIfTensorAlreadyExist(fNB)) { + throw + std::runtime_error("TMVA SOFIE BatchNormalization op Input Tensor " + fNB + " fnb is not found in model"); + } + if (!model.CheckIfTensorAlreadyExist(fNMean)) { + throw + std::runtime_error("TMVA SOFIE BatchNormalization op Input Tensor " + fNMean + " fnm is not found in model"); + } + if (!model.CheckIfTensorAlreadyExist(fNVar)) { + throw + std::runtime_error("TMVA SOFIE BatchNormalization op Input Tensor " + fNVar + " fnv is not found in model"); + } + + fShapeX = model.GetTensorShape(fNX); + if (fShapeX.size() != 4) { + throw + std::runtime_error("TMVA SOFIE BatchNormalization Op input tensor " + fNX + " fnx is not of 4 dimensions"); + } + + fShapeScale = model.GetTensorShape(fNScale); + fShapeB = model.GetTensorShape(fNB); + fShapeMean = model.GetTensorShape(fNMean); + fShapeVar = model.GetTensorShape(fNVar); + fShapeY = fShapeX; + model.AddIntermediateTensor(fNY, model.GetTensorType(fNX), fShapeY); + + if (fShapeB.size() == 1) { + // Broadcast scale, bias, input_mean and input_var to shape_X + auto original_B = model.GetInitializedTensorData(fNB); + auto original_S = model.GetInitializedTensorData(fNScale); + auto original_M = model.GetInitializedTensorData(fNMean); + auto original_V = model.GetInitializedTensorData(fNVar); + size_t batchSize = fShapeX[0], channels = fShapeX[1], height = fShapeX[2], width = fShapeX[3]; + size_t n = batchSize * channels * height * width; + if (fType == "float") { + float *original_bias = static_cast(original_B.get()); + float *original_scale = static_cast(original_S.get()); + float *original_mean = static_cast(original_M.get()); + float *original_var = static_cast(original_V.get()); + float *new_bias = new float[n]; + float *new_scale = new float[n]; + float *new_mean = new float[n]; + float *new_var = new float[n]; + size_t bs = 0, ch = 0, h = 0, w = 0; + for(ch=0; ch new_bias_shape = {batchSize,channels,height,width}; + std::shared_ptr new_bias_ptr(new_bias, std::default_delete()); + std::shared_ptr new_scale_ptr(new_scale, std::default_delete()); + std::shared_ptr new_mean_ptr(new_mean, std::default_delete()); + std::shared_ptr new_var_ptr(new_var, std::default_delete()); + model.UpdateInitializedTensor(fNB, model.GetTensorType(fNB), new_bias_shape, new_bias_ptr); + model.UpdateInitializedTensor(fNScale, model.GetTensorType(fNScale), new_bias_shape, new_scale_ptr); + model.UpdateInitializedTensor(fNMean, model.GetTensorType(fNMean), new_bias_shape, new_mean_ptr); + model.UpdateInitializedTensor(fNVar, model.GetTensorType(fNVar), new_bias_shape, new_var_ptr); + fShapeB = model.GetTensorShape(fNB); + fShapeScale = model.GetTensorShape(fNScale); + fShapeMean = model.GetTensorShape(fNMean); + fShapeVar = model.GetTensorShape(fNVar); + } + } + } + + + std::string Generate(std::string OpName){ + OpName = "op_" + OpName; + if (fShapeX.empty()){ + throw std::runtime_error("TMVA SOFIE Batch Normalization called to Generate without being initialized first"); + } + + std::stringstream out; + int length = 1; + for(auto& i: fShapeX){ + length *= i; + } + //// Batch Norm op + size_t batchSize = fShapeX[0], channels = fShapeX[1], height = fShapeX[2], width = fShapeX[3]; + size_t n = batchSize * channels * height * width; + + //// copy X into Y + out << "\t" << "const int N ="< +#include + +namespace TMVA{ +namespace Experimental{ +namespace SOFIE{ + +template +class ROperator_InstanceNormalization final : public ROperator +{ + +private: + + /* Attributes */ + float fepsilon = 1e-05; + + std::string fNX; + std::string fNScale; + std::string fNB; + std::string fNY; + + std::vector fShapeX; + std::vector fShapeScale; + std::vector fShapeB; + std::vector fShapeY; + + std::string fType; + +public: + ROperator_InstanceNormalization() = delete; + + /* Constructor */ + ROperator_InstanceNormalization( float epsilon, + std::string nameX, std::string nameScale, std::string nameB, + std::string nameY): + fepsilon(epsilon), + fNX(UTILITY::Clean_name(nameX)), fNScale(UTILITY::Clean_name(nameScale)), + fNB(UTILITY::Clean_name(nameB)), fNY(UTILITY::Clean_name(nameY)) + { + if(std::is_same::value){ + fType = "float"; + } + else{ + throw + std::runtime_error("TMVA SOFIE Encountered unsupported type parsing a InstanceNormalization operator"); + } + } + + + std::vector TypeInference(std::vector input) { + ETensorType out = input[0]; + return {out}; + } + + std::vector> ShapeInference(std::vector> input) { + if (input.size() != 3 ) { + throw + std::runtime_error("TMVA SOFIE InstanceNormalization Op Shape inference need 3 input tensors"); + } + for(size_t i = 0; i < input.size(); i++) { + if (input[i].size() != 4) { + throw + std::runtime_error("TMVA SOFIE InstanceNormalization Op Shape inference only accept tensor with 4 dimensions"); + } + } + + auto ret = input; + return ret; + } + + void Initialize(RModel& model){ + if (!model.CheckIfTensorAlreadyExist(fNX)) { + throw + std::runtime_error("TMVA SOFIE InstanceNormalization op Input Tensor " + fNX + " fnx is not found in model"); + } + if (!model.CheckIfTensorAlreadyExist(fNScale)) { + throw + std::runtime_error("TMVA SOFIE InstanceNormalization op Input Tensor " + fNScale + " fns is not found in model"); + } + if (!model.CheckIfTensorAlreadyExist(fNB)) { + throw + std::runtime_error("TMVA SOFIE InstanceNormalization op Input Tensor " + fNB + " fnb is not found in model"); + } + + fShapeX = model.GetTensorShape(fNX); + if (fShapeX.size() != 4) { + throw + std::runtime_error("TMVA SOFIE InstanceNormalization Op input tensor " + fNX + " fnx is not of 4 dimensions"); + } + + fShapeScale = model.GetTensorShape(fNScale); + fShapeB = model.GetTensorShape(fNB); + fShapeY = fShapeX; + model.AddIntermediateTensor(fNY, model.GetTensorType(fNX), fShapeY); + } + + + std::string Generate(std::string OpName){ + OpName = "op_" + OpName; + if (fShapeX.empty()){ + throw std::runtime_error("TMVA SOFIE Instance Normalization called to Generate without being initialized first"); + } + + std::stringstream out; + int length = 1; + for(auto& i: fShapeX){ + length *= i; + } + //// Instance Normalization operator + out << "\t" << "for (size_t n = 0; n < " << fShapeX[0] << "; n++) {\n"; + out << "\t" << "\t" << "for (size_t c = 0; c < " << fShapeX[1] << "; c++) {\n"; + + //// calculate mean + out << "\t" << "\t" << "\t" << "float "<< OpName<< "_mean = 0;\n"; + out << "\t" << "\t" << "\t" << "for (size_t h = 0; h < " << fShapeX[2] << "; h++) {\n"; + out << "\t" << "\t" << "\t" << "\t" << "for (size_t w = 0; w < " << fShapeX[3] << "; w++) {\n"; + out << "\t" << "\t" << "\t" << "\t" << "\t" << OpName<< "_mean += tensor_" << fNX << "[n * " << fShapeX[1] * fShapeX[2] * fShapeX[3] << " + c * "<< fShapeX[2] * fShapeX[3] << " + h * " << fShapeX[3] << " + w];\n"; + out << "\t" << "\t" << "\t" << "\t" << "}\n"; + out << "\t" << "\t" << "\t" << "}\n"; + out << "\t" << "\t" << "\t" << OpName<< "_mean = "<< OpName<< "_mean/"<<(fShapeX[2]*fShapeX[3])<<";\n"; + + //// calculate var + out << "\t" << "\t" << "\t" << "float "<< OpName<< "_var = 0;\n"; + out << "\t" << "\t" << "\t" << "for (size_t h = 0; h < " << fShapeX[2] << "; h++) {\n"; + out << "\t" << "\t" << "\t" << "\t" << "for (size_t w = 0; w < " << fShapeX[3] << "; w++) {\n"; + out << "\t" << "\t" << "\t" << "\t" << "\t" << OpName<< "_var += (tensor_" << fNX << "[n * " << fShapeX[1] * fShapeX[2] * fShapeX[3] << " + c * "<< fShapeX[2] * fShapeX[3] << " + h * " << fShapeX[3] << " + w] - "<< OpName<<"_mean) * (tensor_" << fNX << "[n * " << fShapeX[1] * fShapeX[2] * fShapeX[3] << " + c * "<< fShapeX[2] * fShapeX[3] << " + h * " << fShapeX[3] << " + w] - "<< OpName<<"_mean);\n"; + out << "\t" << "\t" << "\t" << "\t" << "}\n"; + out << "\t" << "\t" << "\t" << "}\n"; + out << "\t" << "\t" << "\t" << OpName<< "_var = "<< OpName<< "_var/"<<(fShapeX[2]*fShapeX[3])<<";\n"; + + //// in op + out << "\t" << "\t" << "\t" << "for (size_t h = 0; h < " << fShapeX[2] << "; h++) {\n"; + out << "\t" << "\t" << "\t" << "\t" << "for (size_t w = 0; w < " << fShapeX[3] << "; w++) {\n"; + out << "\t" << "\t" << "\t" << "\t" << "\t" << "tensor_" << fNY << "[n * " << fShapeX[1] * fShapeX[2] * fShapeX[3] << " + c * "<< fShapeX[2] * fShapeX[3] << " + h * " << fShapeX[3] << " + w] = ((tensor_" << fNX << "[n * " << fShapeX[1] * fShapeX[2] * fShapeX[3] << " + c * "<< fShapeX[2] * fShapeX[3] << " + h * " << fShapeX[3] << " + w] - " << OpName<<"_mean)/ std::sqrt(" << OpName<<"_var + "< make_ROperator_Transpose(const onnx::NodeProto& nodep std::unique_ptr make_ROperator_Relu(const onnx::NodeProto& nodeproto, const onnx::GraphProto& graphproto, std::unordered_map& tensor_type); std::unique_ptr make_ROperator_Gemm(const onnx::NodeProto& nodeproto, const onnx::GraphProto& graphproto, std::unordered_map& tensor_type); std::unique_ptr make_ROperator_Conv(const onnx::NodeProto& nodeproto, const onnx::GraphProto& graphproto, std::unordered_map& tensor_type); +std::unique_ptr make_ROperator_BatchNormalization(const onnx::NodeProto& nodeproto, const onnx::GraphProto& graphproto, std::unordered_map& tensor_type); +std::unique_ptr make_ROperator_InstanceNormalization(const onnx::NodeProto& nodeproto, const onnx::GraphProto& graphproto, std::unordered_map& tensor_type); using factoryMethodMap = std::unordered_map (*)(const onnx::NodeProto&, const onnx::GraphProto&, std::unordered_map&)>; @@ -35,7 +37,9 @@ const factoryMethodMap mapOptypeOperator = { {"Gemm", &make_ROperator_Gemm}, {"Transpose", &make_ROperator_Transpose}, {"Relu", &make_ROperator_Relu}, - {"Conv", &make_ROperator_Conv} + {"Conv", &make_ROperator_Conv}, + {"BatchNormalization", &make_ROperator_BatchNormalization}, + {"InstanceNormalization", &make_ROperator_InstanceNormalization} }; diff --git a/tmva/sofie_parsers/src/RModelParser_ONNX.cxx b/tmva/sofie_parsers/src/RModelParser_ONNX.cxx index 7a7849f3c1711..38af36c4e17c4 100644 --- a/tmva/sofie_parsers/src/RModelParser_ONNX.cxx +++ b/tmva/sofie_parsers/src/RModelParser_ONNX.cxx @@ -214,6 +214,81 @@ std::unique_ptr make_ROperator_Conv(const onnx::NodeProto& nodeproto, return op; } +std::unique_ptr make_ROperator_BatchNormalization(const onnx::NodeProto& nodeproto, const onnx::GraphProto& graphproto, std::unordered_map& tensor_type) { + + ETensorType input_type; + + auto input_name = nodeproto.input(0); + auto it = tensor_type.find(input_name); + if (it != tensor_type.end()) { + input_type = it->second; + } else { + throw + std::runtime_error("TMVA::SOFIE ONNX Parser BatchNorm op has input tensor " + input_name + " but its type is not yet registered"); + } + + std::unique_ptr op; + float fepsilon = 1e-05; + float fmomentum = 0.9; + std::size_t ftraining_mode = 0; + + switch(input_type) { + case ETensorType::FLOAT: + if (nodeproto.input_size() == 5) { + op.reset(new ROperator_BatchNormalization(fepsilon, fmomentum, ftraining_mode, nodeproto.input(0), nodeproto.input(1), nodeproto.input(2), nodeproto.input(3), nodeproto.input(4), nodeproto.output(0))); + } + break; + default: + throw + std::runtime_error("TMVA::SOFIE - Unsupported - Operator BatchNorm does not yet support input type " + std::to_string(static_cast(input_type))); + } + + ETensorType output_type = (op->TypeInference({input_type, input_type, input_type, input_type, input_type}))[0]; + auto it2 = tensor_type.find(nodeproto.output(0)); + if (it2 == tensor_type.end()) { + tensor_type[nodeproto.output(0)] = output_type; + } + + return std::move(op); +} + +std::unique_ptr make_ROperator_InstanceNormalization(const onnx::NodeProto& nodeproto, const onnx::GraphProto& graphproto, std::unordered_map& tensor_type) { + + ETensorType input_type; + + auto input_name = nodeproto.input(0); + auto it = tensor_type.find(input_name); + if (it != tensor_type.end()) { + input_type = it->second; + } else { + throw + std::runtime_error("TMVA::SOFIE ONNX Parser IN op has input tensor " + input_name + " but its type is not yet registered"); + } + + std::unique_ptr op; + float fepsilon = 1e-05; + + switch(input_type) { + case ETensorType::FLOAT: + if (nodeproto.input_size() == 3) { + op.reset(new ROperator_InstanceNormalization(fepsilon, nodeproto.input(0), nodeproto.input(1), nodeproto.input(2), nodeproto.output(0))); + } + break; + default: + throw + std::runtime_error("TMVA::SOFIE - Unsupported - Operator IN does not yet support input type " + std::to_string(static_cast(input_type))); + } + + ETensorType output_type = (op->TypeInference({input_type, input_type, input_type}))[0]; + auto it2 = tensor_type.find(nodeproto.output(0)); + if (it2 == tensor_type.end()) { + tensor_type[nodeproto.output(0)] = output_type; + } + + return std::move(op); +} + + } //INTERNAL @@ -347,6 +422,8 @@ RModel RModelParser_ONNX::Parse(std::string filename){ rmodel.AddBlasRoutines({"Gemm", "Gemv"}); } else if (op_type == "Conv") { rmodel.AddBlasRoutines({"Gemm", "Axpy"}); + } else if (op_type == "BatchNormalization") { + rmodel.AddBlasRoutines({"Copy", "Axpy"}); } }