Add fit()

Author: MagicCube

src/ai/models/Model.js

@@ -17,12 +17,18 @@ export default class Model {
     return this.predict([inputX]);
   }
 
-  train(inputXs, inputYs, iterationCount = 100) {
+  train(inputXs, inputYs) {
     throw new Error(
       'Abstract method must be implemented in the derived class.'
     );
   }
 
+  fit(inputXs, inputYs, iterationCount = 100) {
+    for (let i = 0; i < iterationCount; i += 1) {
+      this.train(inputXs, inputYs);
+    }
+  }
+
   loss(predictedYs, labels) {
     const meanSquareError = predictedYs
       .sub(tensor(labels))

src/ai/models/genetic/GeneticModel.js

@@ -7,6 +7,10 @@ export default class GeneticModel extends Model {
     this.mutate(offspring);
   }
 
+  fit(chromosomes) {
+    this.train(chromosomes);
+  }
+
   select(chromosomes) {
     const parents = [chromosomes[0], chromosomes[1]];
     return parents;

src/ai/models/nn/NNModel.js

@@ -47,12 +47,10 @@ export default class NNModel extends Model {
     return prediction;
   }
 
-  train(inputXs, inputYs, iterationCount = 100) {
-    for (let i = 0; i < iterationCount; i += 1) {
-      this.optimizer.minimize(() => {
-        const predictedYs = this.predict(inputXs);
-        return this.loss(predictedYs, inputYs);
-      });
-    }
+  train(inputXs, inputYs) {
+    this.optimizer.minimize(() => {
+      const predictedYs = this.predict(inputXs);
+      return this.loss(predictedYs, inputYs);
+    });
   }
 }