├── crypto_nn
    ├── lib
    │   └── readme.txt
    ├── include
    │   ├── readme.txt
    │   ├── relu.h
    │   └── layer.h
    ├── weights
    │   ├── example_weights.txt
    │   └── readme.txt
    ├── USE.txt
    ├── test_dataset
    │   └── test_set.data
    ├── Makefile
    └── src
    │   ├── relu.cpp
    │   ├── layer.cpp
    │   └── main.cpp
├── relu
    ├── relu_graphic.png
    └── relu_approx.m
├── training
    ├── test.data
    ├── train.data
    └── train_cryptodl.py
└── README.md


/crypto_nn/lib/readme.txt:
--------------------------------------------------------------------------------
1 | Here you should put fhe.a
2 | 


--------------------------------------------------------------------------------
/crypto_nn/include/readme.txt:
--------------------------------------------------------------------------------
1 | Here you should put HElib header files
2 | 


--------------------------------------------------------------------------------
/relu/relu_graphic.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/alessandro-nori/crypto_nn/HEAD/relu/relu_graphic.png


--------------------------------------------------------------------------------
/crypto_nn/weights/example_weights.txt:
--------------------------------------------------------------------------------
1 | 0.2846 -0.2691 -0.3774
2 | 0.2151 -0.4314 1.0593
3 | -0.2613 0.3433 0.4949
4 | 
5 | 


--------------------------------------------------------------------------------
/crypto_nn/weights/readme.txt:
--------------------------------------------------------------------------------
1 | Each line represents a neuron [w0 w1 b]
2 | from layer0 (the one connected with the input) to layer1 (output)
3 | 


--------------------------------------------------------------------------------
/crypto_nn/include/relu.h:
--------------------------------------------------------------------------------
 1 | #ifndef _relu_H
 2 | #define _relu_H
 3 | 
 4 | #include "FHE.h"
 5 | 
 6 | vector<Ctxt> relu(vector<Ctxt> input);
 7 | vector<int64_t> relu(vector<int64_t> input);
 8 | 
 9 | int64_t get_scale();
10 | 
11 | #endif
12 | 


--------------------------------------------------------------------------------
/crypto_nn/USE.txt:
--------------------------------------------------------------------------------
 1 | To build crypto_nn, you will need to have HElib, GMP and NTL libraries installed.
 2 | Now that you have these libraries, you have to copy header files from Helib/src
 3 | directory to crypto_nn/include and fhe.a to crypto_nn/lib.
 4 | 
 5 | On the command line in the crypto_nn directory:
 6 | 
 7 |    make
 8 | 
 9 | will compile and build bin/crypto_nn.
10 | 


--------------------------------------------------------------------------------
/relu/relu_approx.m:
--------------------------------------------------------------------------------
 1 | %%
 2 | clc
 3 | clear all
 4 | close all
 5 | %%
 6 | 
 7 | f = @(x) 1./(1+exp(-x)); % sigmoid function
 8 | x = -600:1:600;
 9 | y = f(x);
10 | n = 1; % polynomial derivative degree
11 | p_d = polyfit(x, y, n);
12 | 
13 | l = 200; % interval
14 | xx = -l:1:l;
15 | relu = max(0, xx);
16 | yy = polyval(p_d, x);
17 | 
18 | p = [p_d 50]
19 | yy2 = polyval(p, xx);
20 | plot(xx, yy2, xx, relu)
21 | legend('polynomial ReLU (degree 2)', 'ReLU');
22 | 


--------------------------------------------------------------------------------
/crypto_nn/include/layer.h:
--------------------------------------------------------------------------------
 1 | #ifndef _layer_H
 2 | #define _layer_H
 3 | 
 4 | #include "FHE.h"
 5 | 
 6 | class layer {
 7 | private:
 8 |   uint8_t n_input, n_output;
 9 | 
10 |   // weight matrix
11 |   vector<vector<int>> w;
12 |   // bias array
13 |   vector<int> b;
14 | 
15 |   int64_t scale;
16 | 
17 |   layer();
18 | 
19 | public:
20 |   layer(uint8_t n_input, uint8_t n_output, vector<vector<int>>& w, vector<int>& b, int64_t scale = 1);
21 |   vector<Ctxt> feed_forward(vector<Ctxt> input);
22 |   vector<long> feed_forward(vector<long> input);
23 | };
24 | 
25 | #endif
26 | 


--------------------------------------------------------------------------------
/training/test.data:
--------------------------------------------------------------------------------
 1 | 20 12 0
 2 | 10 13 0
 3 | 1 17 1
 4 | 1 13 1
 5 | 16 2 1
 6 | 0 13 1
 7 | 20 10 0
 8 | 14 5 1
 9 | 11 6 0
10 | 15 18 1
11 | 9 3 1
12 | 12 9 0
13 | 6 2 1
14 | 17 3 1
15 | 15 11 0
16 | 5 16 1
17 | 1 0 1
18 | 17 18 1
19 | 5 19 1
20 | 2 20 1
21 | 6 13 1
22 | 9 10 0
23 | 2 19 1
24 | 2 17 1
25 | 11 5 1
26 | 7 20 1
27 | 2 10 1
28 | 10 14 1
29 | 8 4 1
30 | 18 10 0
31 | 1 3 1
32 | 11 19 1
33 | 3 5 1
34 | 20 9 0
35 | 18 20 1
36 | 19 19 1
37 | 10 17 1
38 | 12 20 1
39 | 13 0 1
40 | 2 14 1
41 | 4 1 1
42 | 1 7 0
43 | 8 19 1
44 | 16 3 1
45 | 9 11 0
46 | 1 17 1
47 | 13 4 1
48 | 0 15 1
49 | 14 7 0
50 | 16 7 0
51 | 


--------------------------------------------------------------------------------
/crypto_nn/test_dataset/test_set.data:
--------------------------------------------------------------------------------
 1 | 20 12 0
 2 | 10 13 0
 3 | 1 17 1
 4 | 1 13 1
 5 | 16 2 1
 6 | 0 13 1
 7 | 20 10 0
 8 | 14 5 1
 9 | 11 6 0
10 | 15 18 1
11 | 9 3 1
12 | 12 9 0
13 | 6 2 1
14 | 17 3 1
15 | 15 11 0
16 | 5 16 1
17 | 1 0 1
18 | 17 18 1
19 | 5 19 1
20 | 2 20 1
21 | 6 13 1
22 | 9 10 0
23 | 2 19 1
24 | 2 17 1
25 | 11 5 1
26 | 7 20 1
27 | 2 10 1
28 | 10 14 1
29 | 8 4 1
30 | 18 10 0
31 | 1 3 1
32 | 11 19 1
33 | 3 5 1
34 | 20 9 0
35 | 18 20 1
36 | 19 19 1
37 | 10 17 1
38 | 12 20 1
39 | 13 0 1
40 | 2 14 1
41 | 4 1 1
42 | 1 7 0
43 | 8 19 1
44 | 16 3 1
45 | 9 11 0
46 | 1 17 1
47 | 13 4 1
48 | 0 15 1
49 | 14 7 0
50 | 16 7 0
51 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | # Crypto Neural Network
 2 | 
 3 | Crypto_nn is a very simple example of neural network that can perform classification over encrypted data using homomorphic encryption.
 4 | The idea is taken from [CryptoDL: Deep Neural Networks over Encrypted Data](https://arxiv.org/pdf/1711.05189.pdf) by Ehsan Hesamifard, Hassan Takabi, Mehdi Ghasemi where
 5 | you can find all the details.
 6 | 
 7 | ## Activation function
 8 | To use activation functions within HE schemes, they should be approximated in a form which is implemented using only addition and multiplication (e.g. polynomial).
 9 | 
10 | In this example I simulated the ReLU function as presented in the paper and I obtained the following approximation:
11 | 
12 | 0.0012x<sup>2</sup> + 0.5x + 52
13 | 
14 | ![alt text](https://github.com/ironordnassela/crypto_nn/blob/master/relu/relu_graphic.png)
15 | 


--------------------------------------------------------------------------------
/crypto_nn/Makefile:
--------------------------------------------------------------------------------
 1 | IDIR = include
 2 | ODIR = build
 3 | SDIR = src
 4 | LDIR = lib
 5 | BDIR = bin
 6 | 
 7 | TARGET = bin/crypto_nn
 8 | 
 9 | CC = g++
10 | CFLAGS = -g -O2 -std=c++11 -pthread -DFHE_THREADS -DFHE_BOOT_THREADS -fmax-errors=2
11 | 
12 | LD = g++
13 | AR = ar
14 | ARFLAGS=rv
15 | GMP=-lgmp
16 | NTL=-lntl
17 | 
18 | LDLIBS = -L/usr/local/lib $(NTL) $(GMP) -lm
19 | 
20 | _DEPS = layer.h relu.h
21 | DEPS = $(patsubst %,$(IDIR)/%,$(_DEPS))
22 | 
23 | _OBJ = main.o layer.o relu.o 
24 | OBJ = $(patsubst %,$(ODIR)/%,$(_OBJ))
25 | 
26 | $(TARGET): $(OBJ)
27 | 	@mkdir -p $(BDIR)
28 | 	$(CC) $(CFLAGS) -o $@ $^ $(LDIR)/fhe.a $(LDLIBS)
29 | 
30 | $(ODIR)/main.o: $(SDIR)/main.cpp
31 | 	@mkdir -p $(ODIR)
32 | 	@$(CC) -c -o $@ $< $(CFLAGS)
33 | 
34 | $(ODIR)/%.o: $(SDIR)/%.cpp $(IDIR)/%.h
35 | 	@mkdir -p $(ODIR)
36 | 	@$(CC) -c -o $@ $< $(CFLAGS)
37 | 
38 | .PHONY: clean
39 | 
40 | clean:
41 | 	@rm -rf $(ODIR) $(TARGET) 
42 | 


--------------------------------------------------------------------------------
/crypto_nn/src/relu.cpp:
--------------------------------------------------------------------------------
 1 | #include "../include/relu.h"
 2 | #include <NTL/lzz_pXFactoring.h>
 3 | 
 4 | // this variable is used to get int coefficient for the ReLU approximation
 5 | static int64_t scale = 10000;
 6 | 
 7 | static int64_t c0 = 520000; // constant
 8 | static int64_t c1 = 5000; // coefficient for degree 1
 9 | static int64_t c2 = 12;  // coefficient for degree 2
10 | 
11 | 
12 | vector<Ctxt> relu(vector<Ctxt> input) {
13 |   vector<Ctxt> output;
14 | 
15 |   for (int i=0; i<input.size(); i++) {
16 |     Ctxt x2 = input[i];
17 |     x2*=x2;
18 |     x2.multByConstant(to_ZZX(c2));
19 | 
20 |     Ctxt x1 = input[i];
21 |     x1.multByConstant(to_ZZX(c1));
22 | 
23 |     output.push_back(x2);
24 |     output[i] += x1;
25 |     output[i].addConstant(to_ZZX(c0));
26 |   }
27 | 
28 |   return output;
29 | }
30 | 
31 | vector<int64_t> relu(vector<int64_t> input) {
32 |   vector<int64_t> output(input.size());
33 | 
34 |   for (int i=0; i<input.size(); i++) {
35 |     output[i] = input[i]*input[i]*c2+input[i]*c1+c0;
36 |   }
37 | 
38 |   return output;
39 | }
40 | 
41 | int64_t get_scale() {
42 |   return scale;
43 | }
44 | 


--------------------------------------------------------------------------------
/crypto_nn/src/layer.cpp:
--------------------------------------------------------------------------------
 1 | #include "../include/layer.h"
 2 | #include <NTL/lzz_pXFactoring.h>
 3 | 
 4 | using namespace std;
 5 | 
 6 | layer::layer(uint8_t n_input, uint8_t n_output, vector<vector<int>>& w, vector<int>& b, int64_t scale)
 7 |             : n_input(n_input), n_output(n_output), w(w), b(b), scale(scale) {}
 8 | 
 9 | vector<Ctxt> layer::feed_forward(vector<Ctxt> input) {
10 |   vector<Ctxt> output;
11 | 
12 |   for (uint8_t i=0; i<n_output; i++) {
13 |     for (uint8_t j=0; j<n_input; j++) {
14 |       Ctxt temp = input[j];
15 |       temp.multByConstant(to_ZZX(w[i][j]));
16 |       if (j==0) output.push_back(temp);
17 |       else output[i] += temp;
18 |     }
19 | 
20 |     output[i].addConstant(to_ZZX((long)b[i]*scale));
21 | 
22 |   }
23 | 
24 |   return output;
25 | }
26 | 
27 | vector<long> layer::feed_forward(vector<long> input) {
28 |   vector<long> output(n_output, 0);
29 | 
30 |   for (uint8_t i=0; i<n_output; i++) {
31 |     for (uint8_t j=0; j<n_input; j++) {
32 |       output[i] += input[j]*w[i][j];
33 |       // cout << output[i] << " " << input[j] << " " << int(w[i][j]) << " " << int(b[i]) << endl;
34 |     }
35 | 
36 |     output[i] += (long)b[i]*scale;
37 | 
38 |   }
39 | 
40 |   return output;
41 | }
42 | 


--------------------------------------------------------------------------------
/training/train.data:
--------------------------------------------------------------------------------
  1 | 10 15 1
  2 | 9 19 1
  3 | 4 19 1
  4 | 7 12 0
  5 | 11 10 0
  6 | 7 17 1
  7 | 1 2 1
  8 | 8 14 1
  9 | 14 8 0
 10 | 2 12 1
 11 | 7 16 1
 12 | 19 16 0
 13 | 17 15 0
 14 | 15 15 0
 15 | 15 4 1
 16 | 3 20 1
 17 | 1 12 1
 18 | 13 20 1
 19 | 15 15 0
 20 | 12 18 1
 21 | 19 14 0
 22 | 16 7 0
 23 | 19 8 0
 24 | 0 12 1
 25 | 10 6 1
 26 | 20 4 1
 27 | 18 4 1
 28 | 5 9 0
 29 | 14 11 0
 30 | 9 19 1
 31 | 9 3 1
 32 | 4 18 1
 33 | 18 15 0
 34 | 8 15 1
 35 | 19 14 0
 36 | 12 2 1
 37 | 1 17 1
 38 | 18 18 1
 39 | 2 8 0
 40 | 20 19 1
 41 | 15 1 1
 42 | 3 1 1
 43 | 13 12 0
 44 | 0 11 1
 45 | 7 12 0
 46 | 3 8 0
 47 | 20 17 0
 48 | 13 0 1
 49 | 10 12 0
 50 | 3 9 0
 51 | 18 10 0
 52 | 1 10 1
 53 | 11 14 0
 54 | 16 3 1
 55 | 7 1 1
 56 | 5 6 1
 57 | 1 20 1
 58 | 17 12 0
 59 | 16 11 0
 60 | 13 11 0
 61 | 1 11 1
 62 | 8 5 1
 63 | 3 11 1
 64 | 3 9 0
 65 | 14 16 1
 66 | 11 10 0
 67 | 1 8 0
 68 | 4 17 1
 69 | 11 18 1
 70 | 9 10 0
 71 | 4 2 1
 72 | 11 10 0
 73 | 13 7 0
 74 | 1 8 0
 75 | 14 5 1
 76 | 11 3 1
 77 | 1 0 1
 78 | 14 14 0
 79 | 17 11 0
 80 | 10 8 0
 81 | 6 20 1
 82 | 12 5 1
 83 | 17 5 1
 84 | 20 5 1
 85 | 2 11 1
 86 | 6 4 1
 87 | 9 9 0
 88 | 8 9 0
 89 | 1 15 1
 90 | 12 2 1
 91 | 13 1 1
 92 | 7 16 1
 93 | 6 20 1
 94 | 4 7 0
 95 | 1 6 1
 96 | 19 3 1
 97 | 1 18 1
 98 | 9 0 1
 99 | 1 1 1
100 | 14 2 1
101 | 5 11 0
102 | 14 20 1
103 | 19 12 0
104 | 2 14 1
105 | 17 3 1
106 | 0 16 1
107 | 14 9 0
108 | 7 17 1
109 | 2 19 1
110 | 17 19 1
111 | 18 18 1
112 | 14 19 1
113 | 18 8 0
114 | 14 4 1
115 | 10 9 0
116 | 6 20 1
117 | 15 15 0
118 | 16 13 0
119 | 20 11 0
120 | 4 1 1
121 | 13 6 0
122 | 10 5 1
123 | 16 18 1
124 | 0 5 1
125 | 6 13 1
126 | 11 0 1
127 | 4 20 1
128 | 18 0 1
129 | 17 2 1
130 | 18 14 0
131 | 15 2 1
132 | 17 2 1
133 | 5 8 0
134 | 15 9 0
135 | 10 11 0
136 | 12 7 0
137 | 2 2 1
138 | 11 2 1
139 | 0 9 1
140 | 3 20 1
141 | 1 18 1
142 | 1 5 1
143 | 12 2 1
144 | 12 7 0
145 | 13 4 1
146 | 2 9 0
147 | 11 4 1
148 | 20 2 1
149 | 16 2 1
150 | 10 14 1
151 | 9 8 0
152 | 18 19 1
153 | 4 16 1
154 | 7 18 1
155 | 15 2 1
156 | 8 7 0
157 | 12 20 1
158 | 12 0 1
159 | 19 8 0
160 | 11 14 0
161 | 9 5 1
162 | 10 14 1
163 | 20 6 0
164 | 12 0 1
165 | 15 17 1
166 | 4 7 0
167 | 5 11 0
168 | 2 12 1
169 | 18 12 0
170 | 12 4 1
171 | 10 18 1
172 | 14 2 1
173 | 6 12 1
174 | 8 17 1
175 | 10 17 1
176 | 8 14 1
177 | 8 15 1
178 | 1 7 0
179 | 9 12 0
180 | 5 12 1
181 | 1 14 1
182 | 16 8 0
183 | 19 19 1
184 | 16 19 1
185 | 15 19 1
186 | 0 18 1
187 | 2 1 1
188 | 18 2 1
189 | 3 11 1
190 | 3 16 1
191 | 10 6 1
192 | 19 19 1
193 | 16 5 1
194 | 5 7 0
195 | 20 9 0
196 | 14 0 1
197 | 2 2 1
198 | 8 3 1
199 | 18 13 0
200 | 3 4 1
201 | 


--------------------------------------------------------------------------------
/crypto_nn/src/main.cpp:
--------------------------------------------------------------------------------
  1 | #include "../include/FHE.h"
  2 | #include <NTL/lzz_pXFactoring.h>
  3 | 
  4 | #include <fstream>
  5 | #include <iostream>
  6 | 
  7 | #include "../include/layer.h"
  8 | #include "../include/relu.h"
  9 | 
 10 | #define MAXQ 127
 11 | 
 12 | int quantize(float x, int maxq = MAXQ, float max = 1) {
 13 |   int xq = int(x*maxq/max);
 14 |   return xq;
 15 | }
 16 | 
 17 | int main(int argc, char **argv) {
 18 | 
 19 |     if (argc < 3) {
 20 |       cerr << "usage: " << argv[0] << " <weights_file_path> <test_dataset>" << endl;
 21 |       exit(-1);
 22 |     }
 23 | 
 24 |     long m=0, p=1308355967371729  , r=1; // Native plaintext space
 25 |     // Computations will be 'modulo p'
 26 |     long L=11;          // Levels
 27 |     long c=3;           // Columns in key switching matrix
 28 |     long w=64;          // Hamming weight of secret key
 29 |     long d=0;
 30 |     long security = 80;
 31 |     ZZX G;
 32 |     m = FindM(security,L,c,p, d, 0, 0);
 33 | 
 34 |     FHEcontext context(m, p, r);
 35 |     // initialize context
 36 |     buildModChain(context, L, c);
 37 |     // modify the context, adding primes to the modulus chain
 38 |     FHESecKey secretKey(context);
 39 |     // construct a secret key structure
 40 |     const FHEPubKey& publicKey = secretKey;
 41 | 
 42 |     // an "upcast": FHESecKey is a subclass of FHEPubKey
 43 | 
 44 |     //if(0 == d)
 45 |     G = context.alMod.getFactorsOverZZ()[0];
 46 | 
 47 |    secretKey.GenSecKey(w);
 48 |    // actually generate a secret key with Hamming weight w
 49 | 
 50 |    cout << "Generated secret key" << endl;
 51 | 
 52 |    //
 53 |    int n_input = 2, n_H = 2, n_output = 1;
 54 | 
 55 |    ifstream f;
 56 |    string file_name = argv[1];
 57 |    f.open(file_name);
 58 |    if (!f) {
 59 |      cout << "can't load neural network parameters from " << file_name << endl;
 60 |      return -1;
 61 |    }
 62 | 
 63 | 
 64 |    // initialization hidden layer
 65 |    vector<vector<int>> weights1(n_H);  // weights
 66 |    vector<int> bias1(n_H); // bias
 67 |    for (int i=0; i<n_H; i++) {
 68 |      float w1, w2, b;
 69 |      f >> w1 >> w2 >> b;
 70 |      vector<int> w(n_input);
 71 |      w[0] = quantize(w1);
 72 |      w[1] = quantize(w2);
 73 |      bias1[i] = quantize(b);
 74 |      // cout << w[0] << " " << w[1] << " " << bias1[i] << endl;
 75 |      weights1[i] = w;
 76 |    }
 77 | 
 78 |    // initialization output layer
 79 |    vector<vector<int>> weights2(n_output);
 80 |    vector<int> bias2(n_output);
 81 |    for (int i=0; i<n_output; i++) {
 82 |      float w1, w2, b;
 83 |      f >> w1 >> w2 >> b;
 84 |      vector<int> w(n_H);
 85 |      w[0] = quantize(w1);
 86 |      w[1] = quantize(w2);
 87 |      bias2[i] = quantize(b);
 88 |      // cout << w[0] << " " << w[1] << " " << bias2[i] << endl;
 89 |      weights2[i] = w;
 90 |    }
 91 | 
 92 |    f.close();
 93 | 
 94 |    layer l1(2, 2, weights1, bias1);
 95 |    layer l2(2, 1, weights2, bias2, get_scale()*127);
 96 | 
 97 |    cout << "Neural network parameters loaded" << endl;
 98 | 
 99 |    f.open(argv[2]);
100 |    if (!f) {
101 |      cout << "can't load data for test" << endl;
102 |      return -1;
103 |    }
104 | 
105 |    ofstream fout;
106 |    fout.open("pred.txt");
107 |    if (!f) {
108 |      cout << "can't output the predictions" << endl;
109 |      return -1;
110 |    }
111 | 
112 |    long x1, x2, pred;
113 |    while (f >> x1 >> x2 >> pred) {
114 | 
115 |      // this should be done client side and cipher inputs should be sent to the Cloud
116 |      Ctxt ctx1(publicKey);
117 |      Ctxt ctx2(publicKey);
118 | 
119 |      publicKey.Encrypt(ctx1, to_ZZX(x1));
120 |      publicKey.Encrypt(ctx2, to_ZZX(x2));
121 |      //
122 | 
123 |      // what the Cloud receives
124 |      vector<Ctxt> inputC;
125 |      inputC.push_back(ctx1);
126 |      inputC.push_back(ctx2);
127 | 
128 |      vector<Ctxt> outputC = relu(l1.feed_forward(inputC));
129 | 
130 |      outputC = l2.feed_forward(outputC);
131 |      // now the Cloud can send back the cipher output to the client which can decrypt it
132 | 
133 |      ZZX outputZ;
134 |      secretKey.Decrypt(outputZ, outputC[0]);
135 | 
136 |      /*
137 |      * plain output is scaled of 127*127 (due to weights quantization) * 10000 (due to integer coefficient of ReLU)
138 |      * in my solution the Cloud should also send this value to the client in order to allow it to convert the output to real from integer
139 |      */
140 | 
141 |      cout << "plain input: " << x1 << " " << x2 << endl;
142 |      cout << "plain output: " << coeff(outputZ, 0) << endl;
143 |      cout << endl;
144 | 
145 |      fout << x1 << " " << x2 << " " << coeff(outputZ, 0) << endl;
146 | 
147 |    }
148 | 
149 |    f.close();
150 | 
151 |    cout << "output scale: " << MAXQ*MAXQ*get_scale() << endl;
152 | 
153 |    return 0;
154 | }
155 | 


--------------------------------------------------------------------------------
/training/train_cryptodl.py:
--------------------------------------------------------------------------------
  1 | import sys
  2 | 
  3 | import matplotlib.pyplot as plt
  4 | import numpy as np
  5 | 
  6 | import torch
  7 | 
  8 | #ReLU approximation coefficients
  9 | a = 0.1524
 10 | b = 0.5
 11 | c = 0.409
 12 | 
 13 | class ReLUApproxim(torch.autograd.Function):
 14 | 
 15 |     @staticmethod
 16 |     def forward(ctx, x):
 17 |         ctx.save_for_backward(x)
 18 |         return a*(x**2)+b*x+c
 19 | 
 20 |     @staticmethod
 21 |     def backward(ctx, grad_output):
 22 |         x, = ctx.saved_tensors
 23 |         grad_input = grad_output.clone()
 24 |         grad_input = grad_input*(2*a*x + b)
 25 |         return grad_input
 26 | 
 27 | class TwoLayerNet(torch.nn.Module):
 28 |     def __init__(self, D_in, H, D_out):
 29 |         super(TwoLayerNet, self).__init__()
 30 |         self.linear1 = torch.nn.Linear(D_in, H)
 31 |         self.linear2 = torch.nn.Linear(H, D_out)
 32 | 
 33 |     def forward(self, x):
 34 |         relu = ReLUApproxim.apply
 35 |         h_relu = relu(self.linear1(x))
 36 |         y_pred = self.linear2(h_relu)
 37 |         return y_pred
 38 | 
 39 |     def get_parameters(self):
 40 |         w = []
 41 |         w.append(list(self.linear1.parameters()))
 42 |         w.append(list(self.linear2.parameters()))
 43 |         return w
 44 | 
 45 |     def save_parameters(self, filename):
 46 |         # layer0
 47 |         params = list(self.linear1.parameters())
 48 |         w1 = params[0][0].detach().numpy()
 49 |         w2 = params[0][1].detach().numpy()
 50 |         b1 = params[1][0].item()
 51 |         b2 = params[1][0].item()
 52 | 
 53 |         #layer1
 54 |         params = list(self.linear2.parameters())
 55 |         w3 = params[0][0].detach().numpy()
 56 |         b3 = params[1][0].detach().item()
 57 | 
 58 |         f = open(filename, 'w')
 59 |         f.write(str(w1[0]) + ' ' + str(w1[1]) + ' ' + str(b1) + '\n')
 60 |         f.write(str(w2[0]) + ' ' + str(w2[1]) + ' ' + str(b2) + '\n')
 61 |         f.write(str(w3[0]) + ' ' + str(w3[1]) + ' ' + str(b3) + '\n')
 62 |         f.close()
 63 | 
 64 | 
 65 | def load_data(filename):
 66 |     try:
 67 |         f = open(filename, "r")
 68 |     except FileNotFoundError:
 69 |         print("File not found!!!")
 70 |         exit()
 71 |     
 72 |     lines = f.readlines()
 73 |     X = []
 74 |     t = []
 75 | 
 76 |     for l in lines:
 77 |         x1, x2, y = [int(s) for s in l.strip().split(" ")]
 78 |         
 79 |         x = torch.tensor([x1, x2], dtype=torch.float32)
 80 |         y = torch.tensor([y], dtype=torch.float32)
 81 | 
 82 |         X.append(x)
 83 |         t.append(y)
 84 |     
 85 |     return X, t
 86 | 
 87 | 
 88 | def test(model, X, t):
 89 |     correct = 0
 90 | 
 91 |     for x, target in zip(X, t):
 92 |             pred = model(x)
 93 | 
 94 |             pred = 0 if pred.item() < 0.5 else 1
 95 |             if pred == target.item():
 96 |                 correct+=1
 97 |     
 98 |     return correct*100/len(t)
 99 | 
100 | def plot_predictions(model, X):
101 |     # lines parameters
102 |     m1 = -0.1
103 |     b1 = 7
104 |     m2 = 0.5
105 |     b2 = 9
106 | 
107 |     predictions = []
108 |     inputs = []
109 | 
110 |     for x in X:
111 |         inputs.append(x.numpy())
112 |         pred = model(x)
113 |         pred = 'r' if pred.item() < 0.5 else 'b'
114 |         predictions.append(pred)
115 | 
116 |     inputs = np.array(inputs)
117 |     predictions = np.array(predictions)
118 | 
119 |     axis0 = np.linspace(0,20,200)
120 |     line1 = axis0*m1+b1
121 |     line2 = axis0*m2+b2
122 |     line2[line2<0] = 0
123 | 
124 |     ax = plt.axes()
125 |     ax.set(xlim=(0, 20), ylim=(0, 20), xlabel='x1', ylabel='x2')
126 |     ax.plot(axis0, line1, 'r')
127 |     ax.plot(axis0, line2, 'r')
128 |     ax.scatter(inputs[:,0], inputs[:,1], c=predictions)
129 |     axisx = np.linspace(0,20,200)
130 |     ax.fill_between(axisx, 0, 20,
131 |                     color='blue', alpha=0.2)
132 |     ax.fill_between(axis0, line2, line1,
133 |                     color='red', alpha=0.2)
134 | 
135 |     plt.show()
136 | 
137 | 
138 | def main():
139 | 
140 |     if len(sys.argv) > 1:
141 |         # output file name
142 |         fout = sys.argv[1]
143 |     else:
144 |         fout = 'weights.txt'
145 | 
146 |     # D_in is input dimension
147 |     # D_out is output dimension
148 |     # H is hidden dimension
149 |     D_in, H, D_out = 2, 2, 1
150 | 
151 |     model = TwoLayerNet(D_in, H, D_out)
152 |     loss_fn = torch.nn.MSELoss()
153 |     optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
154 | 
155 |     X, t = load_data('train.data')
156 |     X_test, t_test = load_data('test.data')
157 |     print('Training and Test data loaded')
158 |     accuracy = 0.00
159 | 
160 |     print('Training phase...')
161 |     epoch = 0
162 |     while accuracy < 95.0:
163 |         for x, target in zip(X, t):
164 |             pred = model(x)
165 |             loss = loss_fn(pred, target)
166 | 
167 |             optimizer.zero_grad()
168 |             loss.backward()
169 |             optimizer.step()
170 |         
171 |         accuracy = test(model, X_test, t_test)
172 |         epoch+=1
173 |         print('Epoch', epoch, 'accuracy:', accuracy)
174 | 
175 | 
176 |     model.save_parameters(fout)
177 | 
178 |     plot_predictions(model, X_test)
179 | 
180 |     
181 | 
182 | 
183 | if __name__ == '__main__':
184 |     main()


--------------------------------------------------------------------------------