[tmva][sofie][gsoc] Add InstanceNormalization operator

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"

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 <cmath>
+#include <sstream>
+
+namespace TMVA{
+namespace Experimental{
+namespace SOFIE{
+
+template <typename T>
+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<size_t> fShapeX;
+	std::vector<size_t> fShapeScale;
+	std::vector<size_t> fShapeB;
+	std::vector<size_t> fShapeMean;
+	std::vector<size_t> fShapeVar;
+	std::vector<size_t> 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<T, float>::value){
+			fType = "float";
+		}
+		else{
+			throw
+				std::runtime_error("TMVA SOFIE Encountered unsupported type parsing a BatchNormalization operator");
+		}
+	}
+
+
+	std::vector<ETensorType> TypeInference(std::vector<ETensorType> input) {
+		ETensorType out = input[0];
+		return {out};
+	}
+
+	std::vector<std::vector<size_t>> ShapeInference(std::vector<std::vector<size_t>> 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<float*>(original_B.get());
+				float *original_scale = static_cast<float*>(original_S.get());
+				float *original_mean = static_cast<float*>(original_M.get());
+				float *original_var = static_cast<float*>(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<channels; ch++){
+					for(h=0; h<height; h++){
+						for(w=0; w<width; w++){
+							new_bias[bs*channels*height*width + ch*height*width + h*width + w] = original_bias[ch];
+							new_scale[bs*channels*height*width + ch*height*width + h*width + w] = original_scale[ch];
+							new_mean[bs*channels*height*width + ch*height*width + h*width + w] = original_mean[ch];
+							new_var[bs*channels*height*width + ch*height*width + h*width + w] = original_var[ch];
+						}
+					}
+				}
+				size_t Batchoffset = channels*height*width;
+				for(bs = 1; bs<batchSize; bs++){
+					std::copy(new_bias, new_bias+Batchoffset, new_bias+(bs*Batchoffset));
+					std::copy(new_scale, new_scale+Batchoffset, new_scale+(bs*Batchoffset));
+					std::copy(new_mean, new_mean+Batchoffset, new_mean+(bs*Batchoffset));
+					std::copy(new_var, new_var+Batchoffset, new_var+(bs*Batchoffset));
+				}
+				//// new_var =1. / sqrt(input_var + fepsilon)
+				for(size_t i=0; i<n; i++){
+					new_var[i] = 1./sqrt(new_var[i] + fepsilon);	
+				}
+				std::vector<size_t> new_bias_shape = {batchSize,channels,height,width};
+				std::shared_ptr<void> new_bias_ptr(new_bias, std::default_delete<float[]>());
+				std::shared_ptr<void> new_scale_ptr(new_scale, std::default_delete<float[]>());
+				std::shared_ptr<void> new_mean_ptr(new_mean, std::default_delete<float[]>());
+				std::shared_ptr<void> new_var_ptr(new_var, std::default_delete<float[]>());
+				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 ="<<batchSize * channels * height * width<<";\n";
+		out << "\t" << "const int "<<OpName<< "_incx = 1;\n";
+		out << "\t" << "const int "<<OpName<< "_incy = 1;\n";
+		out << "\t" << "BLAS::scopy_(&N, " << "tensor_" << fNX <<", &" << OpName << "_incx," << "tensor_" << fNY <<", &" << OpName << "_incy);\n\n";
+
+		//// blas saxpy (Y = -Bmean + Y)
+		out << "\t" << "float "<<OpName<< "_alpha = -1;\n";
+		out << "\t" << "BLAS::saxpy_(&N, &" << OpName << "_alpha, " << "tensor_" << fNMean << ", &" << OpName << "_incx," << "tensor_" << fNY <<", &" << OpName << "_incy);\n\n";
+
+		//// Y *= scale*var
+		out << "\t" << "for (size_t i = 0; i < " << n << "; i++) {\n";
+		out << "\t" << "\t" << "tensor_" << fNY << "[i] *= tensor_" << fNScale << "[i] * tensor_" << fNVar << "[i]; \n";
+		out << "\t" << "}\n";
+
+		//// blas saxpy (Y = Bbias + Y)
+		out << "\t" <<OpName<< "_alpha = 1;\n";
+		out << "\t" << "BLAS::saxpy_(&N, &" << OpName << "_alpha, " << "tensor_" << fNB << ", &" << OpName << "_incx, " << "tensor_" << fNY <<", &" << OpName << "_incy);\n\n";
+        // std::cout<<out;
+		return out.str();
+	}
+
+};
+
+}//SOFIE
+}//Experimental
+}//TMVA
+
+
+#endif //TMVA_SOFIE_ROPERATOR_BatchNormalization

tmva/sofie/inc/TMVA/ROperator_InstanceNormalization.hxx

@@ -0,0 +1,158 @@
+#ifndef TMVA_SOFIE_ROPERATOR_InstanceNormalization
+#define TMVA_SOFIE_ROPERATOR_InstanceNormalization
+
+#include "SOFIE_common.hxx"
+#include "ROperator.hxx"
+#include "RModel.hxx"
+
+
+#include <cmath>
+#include <sstream>
+
+namespace TMVA{
+namespace Experimental{
+namespace SOFIE{
+
+template <typename T>
+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<size_t> fShapeX;
+	std::vector<size_t> fShapeScale;
+	std::vector<size_t> fShapeB;
+	std::vector<size_t> 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<T, float>::value){
+			fType = "float";
+		}
+		else{
+			throw
+				std::runtime_error("TMVA SOFIE Encountered unsupported type parsing a InstanceNormalization operator");
+		}
+	}
+
+
+	std::vector<ETensorType> TypeInference(std::vector<ETensorType> input) {
+		ETensorType out = input[0];
+		return {out};
+	}
+
+	std::vector<std::vector<size_t>> ShapeInference(std::vector<std::vector<size_t>> 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 + "<<fepsilon<<" ))* " << "tensor_" << fNScale <<"[c] + " << "tensor_" <<fNB<<"[c];\n";
+		out << "\t" << "\t" << "\t" << "\t" << "}\n";
+		out << "\t" << "\t" << "\t" << "}\n";
+		out << "\t" << "\t" << "}\n";	
+		out << "\t" << "}\n";
+		// std::cout<<out;
+		return out.str();
+	}
+
+};
+
+}//SOFIE
+}//Experimental
+}//TMVA
+
+
+#endif //TMVA_SOFIE_ROPERATOR_InstanceNormalization

tmva/sofie/src/RModel.cxx

@@ -188,6 +188,8 @@ namespace SOFIE{
             } else if (routine == "Axpy") {
                fGC += ("\textern \"C\" void saxpy_(const int * n, const float * alpha, const float * x,\n"
                        "\t                         const int * incx, float * y, const int * incy);\n");
+            } else if (routine == "Copy") {
+               fGC += ("\textern \"C\" void scopy_(const int *n, const float* x, const int *incx, float* y, const int* incy);\n");
             }
          }
          fGC += ("}//BLAS\n");