binary_mnist.cpp

#include <bits/stdc++.h>

using namespace std;

float binarize(float x){
	float xx = 0.5f * x + 0.5f;
	xx = min(max(xx, 0.0f), 1.0f);
	return (xx >= 0.5f) ? 1.0f : -1.0f;
}

void BinaryConv2D(float *inputs, float* weights, float* outputs,
	int height, int width, int filter_num, int channel){
//kernel = 3x3, //padding=same
//dimension order: 
// - weight: k_h, k_w, channel, filter
// - inputs, outputs: height, width, channel
	for(int fi = 0; fi < filter_num; fi++){
		for(int y = 0; y < height; y++){
			for(int x = 0; x < width; x++){
				float sum = 0;
				for(int chan = 0; chan < channel; chan++){
					for(int yy = -1; yy <= 1; yy++){
						for(int xx = -1; xx <= 1; xx++){
							int py = y + yy;
							int px = x + xx;
							float x = 0;
							if(py < 0 || px < 0 || py >= height || px >= width) x = 0;
							else x = inputs[py*width*channel+px*channel+chan];
							int index = (3*channel*filter_num)*(yy+1) + channel*filter_num*(xx+1) + filter_num*chan + fi;
							float w = binarize(weights[index]);
							sum += w * x;
						}
					}
        }
				outputs[width*filter_num*y+filter_num*x+fi] = sum;//[28][28][32]
			}
    }
	}
}

void BatchNormalization(float* data, float* weights, int height, int width, int channel){
	for(int chan = 0; chan < channel; chan++){
		//gamma, beta, mean, variance
		float gamma = weights[chan];
		float beta = weights[chan + channel];
		float mean = weights[chan + channel*2];
		float variance = weights[chan + channel*3];
		const float eps = 1e-6;
		for(int i = 0; i < width*height; i++){
			data[i*channel+chan] = ((data[i*channel+chan] - mean) / sqrt(variance + eps)) * gamma + beta;
		}
	}
}

void BinaryActivation(float* data, int height, int width, int channel){
	for(int i = 0; i < width*height*channel; i++) data[i] = binarize(data[i]);
}

void MaxPooling2D(float* inputs, float* outputs, int height, int width, int channel){
	//2x2 -> 1x1
	int output_width = width/2;
	int output_height = height/2;
	for(int y = 0; y < output_height; y++){
		for(int x = 0; x < output_width; x++){
			for(int c = 0; c < channel; c++){
				float p1 = inputs[width*channel*(y*2)+channel*(x*2)+c];
				float p2 = inputs[width*channel*(y*2)+channel*(x*2+1)+c];
				float p3 = inputs[width*channel*(y*2+1)+channel*(x*2)+c];
				float p4 = inputs[width*channel*(y*2+1)+channel*(x*2+1)+c];
				outputs[output_width*channel*y+channel*x+c] = max(max(p1, p2), max(p3, p4));
			}
		}
	}
}

void BinaryDense(float* inputs, float* weights, float* outputs, int input_size, int output_size){
	for(int i = 0; i < output_size; i++){
		float sum = 0;
		for(int j = 0; j < input_size; j++){
			sum += inputs[j] * binarize(weights[j*output_size+i]);
		}
		outputs[i] = sum;
	}
}

void Activation_Softmax(float* data, int size){
	float exp_sum = 1e-9;
	for(int i = 0; i < size; i++){
		float exp_x = exp(data[i]);
		exp_sum += exp_x;
		data[i] = exp_x;
	}
	for(int i = 0; i < size; i++){
		data[i] /= exp_sum;
	}
}

int main(){
	//load inputs and weights
	char inputs_name[] = "extracted/input.bin";
	FILE *fp = fopen(inputs_name, "rb");
	const int INPUT_SIZE = 28*28;
	float inputs[INPUT_SIZE];
	fread(inputs, sizeof(float), INPUT_SIZE, fp);
	fclose(fp);
	for(int i = 0; i < 28; i++){
		for(int j = 0; j < 28; j++){
			cout << (inputs[i*28+j] == 0 ? "*" : " ");
		}
		cout << endl;
	}

	map<int, float*> conv2d_weights;
	vector<pair<int, int>> binaryconv2d_layers = {{0, 3*3*1*32}, {3,3*3*32*32}, {7,3*3*32*64}, {10, 3*3*64*64}};//index, size
	for(int i = 0; i < binaryconv2d_layers.size(); i++){
		int idx = binaryconv2d_layers[i].first;
		int size = binaryconv2d_layers[i].second;
		string weights_name = "extracted/weights_" + to_string(idx) + ".bin";
		fp = fopen(weights_name.c_str(), "rb");
		float* weights = (float*)malloc(size * sizeof(float));
		fread(weights, sizeof(float), size, fp);
		fclose(fp);
		conv2d_weights[idx] = weights;
	}
	map<int, float*> batchnorm_weights;
	vector<pair<int, int>> batchnorm_layers = {{1, 32}, {4, 32}, {8, 64}, {11, 64}, {16, 128}, {19, 10}};
	for(int i = 0; i < batchnorm_layers.size(); i++){
		int idx = batchnorm_layers[i].first;
		int size = batchnorm_layers[i].second;
		string weights_name = "extracted/batchnorm_" + to_string(idx) + ".bin";
		fp = fopen(weights_name.c_str(), "rb");
		float* weights = (float*)malloc(size * sizeof(float)*4);//gamma, beta, mean, variance
		fread(weights, sizeof(float), size*4, fp);
		fclose(fp);
		batchnorm_weights[idx] = weights;
	}
	map<int, float*> dense_weights;
	vector<pair<int, int>> dense_layers = {{15, 7*7*64*128}, {18, 128*10}};//index, size
	for(int i = 0; i < dense_layers.size(); i++){
		int idx = dense_layers[i].first;
		int size = dense_layers[i].second;
		string weights_name = "extracted/weights_" + to_string(idx) + ".bin";
		fp = fopen(weights_name.c_str(), "rb");
		float* weights = (float*)malloc(size * sizeof(float));
		fread(weights, sizeof(float), size, fp);
		fclose(fp);
		dense_weights[idx] = weights;
	}
	//load python outputs
	char outputs_name[] = "extracted/outputs.bin";
	fp = fopen(outputs_name, "rb");
	const int N = 10;
	float sample_outputs[N];
	fread(sample_outputs, sizeof(float), N, fp);

	float outputs[28*28*32];
	float outputs2[28*28*32];
	float outputs3[14*14*32];
	float outputs4[14*14*64];
	float outputs5[14*14*64];
	float outputs6[14*14*64];
	float outputs7[128];
	float outputs8[10];
	BinaryConv2D(inputs, conv2d_weights[0], outputs, 28, 28, 32, 1);
	BatchNormalization(outputs, batchnorm_weights[1], 28, 28, 32);
	BinaryActivation(outputs, 28, 28, 32);
	BinaryConv2D(outputs, conv2d_weights[3], outputs2, 28, 28, 32, 32);
	BatchNormalization(outputs2, batchnorm_weights[4], 28, 28, 32);
	BinaryActivation(outputs2, 28, 28, 32);
	MaxPooling2D(outputs2, outputs3, 28, 28, 32);
	BinaryConv2D(outputs3, conv2d_weights[7], outputs4, 14, 14, 64, 32);
	BatchNormalization(outputs4, batchnorm_weights[8], 14, 14, 64);
	BinaryActivation(outputs4, 14, 14, 64);
	BinaryConv2D(outputs4, conv2d_weights[10], outputs5, 14, 14, 64, 64);
	BatchNormalization(outputs5, batchnorm_weights[11], 14, 14, 64);
	BinaryActivation(outputs5, 14, 14, 64);
	MaxPooling2D(outputs5, outputs6, 14, 14, 64);
	BinaryDense(outputs6, dense_weights[15], outputs7, 7*7*64, 128);
	BatchNormalization(outputs7, batchnorm_weights[16], 1, 1, 128);
	BinaryActivation(outputs7, 1, 1, 128);
	BinaryDense(outputs7, dense_weights[18], outputs8, 128, 10);
	BatchNormalization(outputs8, batchnorm_weights[19], 1, 1, 10);
	Activation_Softmax(outputs8, 10);
	//verify
	bool ok = true;
	int ans = 0;
	float maxval = -1;
	for(int i = 0; i < 10; i++){
		if(abs(sample_outputs[i] - outputs8[i] > 1e-4)) ok = false;
		if(maxval < outputs8[i]){
			maxval = outputs8[i];
			ans = i;
		}
	}
	cout << (ok ? "OK" : "NG") << endl;
	cout << "Result : " << ans << " " << outputs8[ans] <<  endl;
	//free memory
	for(auto itr = conv2d_weights.begin(); itr != conv2d_weights.end(); itr++){
		float* ptr = itr->second;
		free(ptr);
	}
	for(auto itr = batchnorm_weights.begin(); itr != batchnorm_weights.end(); itr++){
		float* ptr = itr->second;
		free(ptr);
	}
	for(auto itr = dense_weights.begin(); itr != dense_weights.end(); itr++){
		float* ptr = itr->second;
		free(ptr);
	}
}