include/tkDNN/Layer.h

@@ -180,6 +180,65 @@ public:
 };
 
 
+class LayerBNWgs : public Layer {
+public:
+    LayerBNWgs(Network* net, int input, int output, std::string fname_weights);
+    ~LayerBNWgs();
+
+    int inputs, outputs;
+    std::string weights_path;
+
+    dnnType* bias_h, * bias_d;
+    dnnType* power_h = nullptr;
+    dnnType* scales_h = nullptr, * scales_d = nullptr;
+    dnnType* mean_h = nullptr, * mean_d = nullptr;
+    dnnType* variance_h = nullptr, * variance_d = nullptr;
+
+    __half* bias16_h = nullptr, * bias16_d = nullptr;
+    __half* power16_h = nullptr, * power16_d = nullptr;
+    __half* scales16_h = nullptr, * scales16_d = nullptr;
+    __half* mean16_h = nullptr, * mean16_d = nullptr;
+    __half* variance16_h = nullptr, * variance16_d = nullptr;
+
+
+    void releaseHost(bool release32 = true, bool release16 = true) {
+        if (release32) {
+            if (bias_h != nullptr) { delete[]     bias_h;     bias_h = nullptr; }
+            if (scales_h != nullptr) { delete[]   scales_h;   scales_h = nullptr; }
+            if (mean_h != nullptr) { delete[]     mean_h;     mean_h = nullptr; }
+            if (variance_h != nullptr) { delete[] variance_h; variance_h = nullptr; }
+            if (power_h != nullptr) { delete[]    power_h;    power_h = nullptr; }
+        }
+        if (net->fp16 && release16) {
+            if (bias16_h != nullptr) { delete[]     bias16_h;     bias16_h = nullptr; }
+            if (scales16_h != nullptr) { delete[]   scales16_h;   scales16_h = nullptr; }
+            if (mean16_h != nullptr) { delete[]     mean16_h;     mean16_h = nullptr; }
+            if (variance16_h != nullptr) { delete[] variance16_h; variance16_h = nullptr; }
+            if (power16_h != nullptr) { delete[]    power16_h;    power16_h = nullptr; }
+
+        }
+    }
+
+    void releaseDevice(bool release32 = true, bool release16 = true) {
+        if (release32) {
+            if (bias_d != nullptr) { cudaFree(bias_d);     bias_d = nullptr; }
+            if (scales_d != nullptr) { cudaFree(scales_d);   scales_d = nullptr; }
+            if (mean_d != nullptr) { cudaFree(mean_d);     mean_d = nullptr; }
+            if (variance_d != nullptr) { cudaFree(variance_d); variance_d = nullptr; }
+        }
+        if (net->fp16 && release16) {
+            if (bias16_d != nullptr) { cudaFree(bias16_d);     bias16_d = nullptr; }
+            if (scales16_d != nullptr) { cudaFree(scales16_d);   scales16_d = nullptr; }
+            if (mean16_d != nullptr) { cudaFree(mean16_d);     mean16_d = nullptr; }
+            if (variance16_d != nullptr) { cudaFree(variance16_d); variance16_d = nullptr; }
+            if (power16_d != nullptr) { cudaFree(power16_d);    power16_d = nullptr; }
+        }
+    }
+
+
+
+};
+
 /**
     Input layer (it doesn't need weights)
 */

src/LayerBNWgs.cpp

@@ -0,0 +1,88 @@
+#include <iostream>
+#include <string.h>
+
+#include "Layer.h"
+#include "kernels.h"
+
+namespace tk { namespace dnn {
+	LayerBNWgs::LayerBNWgs(Network* net, int input, int output, std::string fname_weights) : Layer(net) {
+		this->inputs = inputs;
+		this->outputs = output;
+		this->weights_path = fname_weights;
+
+		std::cout << "Reading BatchNorm O = " << outputs << std::endl;
+		int seek = 0;
+		readBinaryFile(weights_path.c_str(), outputs, &bias_h, &bias_d, seek);
+		seek += outputs;
+		readBinaryFile(weights_path.c_str(), outputs, &scales_h, &scales_d, seek);
+		seek += outputs;
+		readBinaryFile(weights_path.c_str(), outputs, &mean_h, &mean_d, seek);
+		seek += outputs;
+		readBinaryFile(weights_path.c_str(), outputs, &variance_h, &variance_d, seek);
+		seek += outputs;
+
+		float eps = TKDNN_BN_MIN_EPSILON;
+
+		power_h = new dnnType[outputs];
+		for (int i = 0; i < outputs; i++) power_h[i] = 1.0f;
+
+		for (int i = 0; i < outputs; i++)
+			mean_h[i] = mean_h[i] / -sqrt(eps + variance_h[i]);
+
+		for (int i = 0; i < outputs; i++)
+			variance_h[i] = 1.0f / sqrt(eps + variance_h[i]);
+
+		if (!net->fp16)
+			return;
+
+		int b_size = outputs;
+		bias16_h = new __half[b_size];
+		cudaMalloc(&bias16_d, b_size * sizeof(__half));
+		float2half(bias_d, bias16_d, b_size);
+		cudaMemcpy(bias16_h, bias16_d, b_size * sizeof(__half), cudaMemcpyDeviceToHost);
+
+		power16_h = new __half[b_size];
+		mean16_h = new __half[b_size];
+		variance16_h = new __half[b_size];
+		scales16_h = new __half[b_size];
+
+		cudaMalloc(&power16_d, b_size * sizeof(__half));
+		cudaMalloc(&mean16_d, b_size * sizeof(__half));
+		cudaMalloc(&variance16_d, b_size * sizeof(__half));
+		cudaMalloc(&scales16_d, b_size * sizeof(__half));
+
+		//temporary buffers
+		float* tmp_d;
+		cudaMalloc(&tmp_d, b_size * sizeof(float));
+
+		//init power array of ones
+		cudaMemcpy(tmp_d, power_h, b_size * sizeof(float), cudaMemcpyHostToDevice);
+		float2half(tmp_d, power16_d, b_size);
+		cudaMemcpy(power16_h, power16_d, b_size * sizeof(__half), cudaMemcpyDeviceToHost);
+
+		//mean array
+		cudaMemcpy(tmp_d, mean_h, b_size * sizeof(float), cudaMemcpyHostToDevice);
+		float2half(tmp_d, mean16_d, b_size);
+		cudaMemcpy(mean16_h, mean16_d, b_size * sizeof(__half), cudaMemcpyDeviceToHost);
+
+		//convert variance
+
+		cudaMemcpy(tmp_d, variance_h, b_size * sizeof(float), cudaMemcpyHostToDevice);
+		float2half(tmp_d, variance16_d, b_size);
+		cudaMemcpy(variance16_h, variance16_d, b_size * sizeof(__half), cudaMemcpyDeviceToHost);
+
+		//convert scales
+		float2half(scales_d, scales16_d, b_size);
+		cudaMemcpy(scales16_h, scales16_d, b_size * sizeof(__half), cudaMemcpyDeviceToHost);
+
+		cudaFree(tmp_d);
+
+
+	}
+
+	LayerBNWgs::~LayerBNWgs() {
+		releaseHost();
+		releaseDevice();
+	}
+
+} }