From b206a93d1226208c562f34f1dd874268c55d039c Mon Sep 17 00:00:00 2001
From: Pravez <paul.breton@enseirb-matmeca.fr>
Date: Tue, 29 May 2018 12:34:00 +0200
Subject: [PATCH] Added tests on convolutional layers

---
 package.json         |  2 +-
 test/model.spec.ts   | 29 ++++++++++++++++++++++++-----
 test/network.spec.ts |  9 +++++----
 3 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/package.json b/package.json
index dbc4b6c..fddd51d 100644
--- a/package.json
+++ b/package.json
@@ -26,7 +26,7 @@
     "setup": "npm install",
     "clean": "rm -rf node_modules/ dist/",
     "doc": "typedoc --theme markdown --out ./docs/ ./src/ --excludeExternals --externalPattern \"**/node_modules/**\" --ignoreCompilerErrors --name \"ReImproveJS\" --mode \"file\" --target ES5 --includeDeclarations",
-    "test": "mocha --reporter spec --compilers ts:ts-node/register 'test/*.spec.ts' --timeout 15000",
+    "test": "mocha --reporter spec --compilers ts:ts-node/register 'test/*.spec.ts' --timeout 120000",
     "build": "./scripts/build-npm.sh",
     "watch": "tsc -w -p ."
   },
diff --git a/test/model.spec.ts b/test/model.spec.ts
index d7fb9f3..68a1098 100644
--- a/test/model.spec.ts
+++ b/test/model.spec.ts
@@ -24,16 +24,35 @@ describe('Old model', () => {
 });
 
 const network = new ConvolutionalNeuralNetwork();
-network.InputShape = [5, 5, 1];
-network.addConvolutionalLayers([32, 64]);
+network.InputShape = [40, 40, 3];
+network.addConvolutionalLayer(32);
 network.addMaxPooling2DLayer();
-network.addNeuralNetworkLayers([128, 128, 2]);
-const nmodel = Model.FromNetwork(network);
-nmodel.compile({loss: 'meanSquaredError', optimizer: 'sgd'});
+network.addConvolutionalLayer(64);
+network.addMaxPooling2DLayer();
+network.addNeuralNetworkLayers([128, {type: 'dense', activation:'softmax', units:2}]);
+const nmodel = Model.FromNetwork(network, {stepsPerEpoch:10, epochs:1});
+nmodel.compile({loss: tf.losses.softmaxCrossEntropy, optimizer: 'adam'});
 
 describe('New model', () => {
     it('should have the right output size', () => {
         for (let i = 0; i < 10; ++i)
             expect(nmodel.randomOutput()).to.be.within(0, numActions);
     });
+
+    it('can be trained', async () => {
+        const x = tf.randomNormal([1, 40, 40, 3]);
+        const y = tf.tensor([[0, 1]]);
+
+        for(let i = 0;i < 5; ++i) {
+            await nmodel.fit(x, y);
+        }
+
+        let results = [];
+        for(let i = 0;i < 10; ++i)
+            results.push(nmodel.predict(x).getAction());
+
+        expect(results.reduce((p, c) => p + c)).to.be.greaterThan(7);
+    });
 });
+
+
diff --git a/test/network.spec.ts b/test/network.spec.ts
index 254109e..8fe0906 100644
--- a/test/network.spec.ts
+++ b/test/network.spec.ts
@@ -8,7 +8,7 @@ describe('Networks', () => {
         network.InputShape = [5, 5, 1];
         network.addConvolutionalLayers([32, 64]);
         network.addMaxPooling2DLayer({type: "maxpooling", strides: [5, 5]});
-        network.addNeuralNetworkLayers([128, 128, 2]);
+        network.addNeuralNetworkLayers([{type: 'dense', units: 256, name: 'test'}, 128, 2]);
 
         const layers = network.getLayers();
 
@@ -29,10 +29,11 @@ describe('Networks', () => {
             type: 'flatten'
         });
 
-        expect(layers[5]).to.be.deep.equal({
+        expect(layers[4]).to.be.deep.equal({
             type: 'dense',
-            units: 128,
-            activation: 'relu'
+            units: 256,
+            activation: 'relu',
+            name:'test'
         });
 
     });