From 3d82718f5a6609666c0627bfd39f597b324c74e7 Mon Sep 17 00:00:00 2001
From: Hunter Nichols <hznichol@ucsc.edu>
Date: Mon, 7 Jun 2021 12:26:45 -0700
Subject: [PATCH] Changed neural network model to be sklearn based

---
 compiler/characterizer/__init__.py       |  2 +-
 compiler/characterizer/lib.py            |  2 ++
 compiler/characterizer/neural_network.py | 28 ++++++++++--------------
 3 files changed, 15 insertions(+), 17 deletions(-)

diff --git a/compiler/characterizer/__init__.py b/compiler/characterizer/__init__.py
index a092ac1e..425ad5f2 100644
--- a/compiler/characterizer/__init__.py
+++ b/compiler/characterizer/__init__.py
@@ -12,13 +12,13 @@ from .lib import *
 from .delay import *
 from .elmore import *
 from .linear_regression import *
+from .neural_network import *
 from .setup_hold import *
 from .functional import *
 from .simulation import *
 from .measurements import *
 from .model_check import *
 from .analytical_util import *
-from .regression_model import *
 
 debug.info(1, "Initializing characterizer...")
 OPTS.spice_exe = ""
diff --git a/compiler/characterizer/lib.py b/compiler/characterizer/lib.py
index f7936413..434b2bf0 100644
--- a/compiler/characterizer/lib.py
+++ b/compiler/characterizer/lib.py
@@ -634,6 +634,8 @@ class lib:
                 from .linear_regression import linear_regression as model
             elif model_name_lc == "elmore":
                 from .elmore import elmore as model
+            elif model_name_lc == "neural_network":
+                from .neural_network import neural_network as model
             else:
                 debug.error("{} model not recognized. See options.py for available models.".format(OPTS.model_name))
 
diff --git a/compiler/characterizer/neural_network.py b/compiler/characterizer/neural_network.py
index 57db4580..6f1aa7c1 100644
--- a/compiler/characterizer/neural_network.py
+++ b/compiler/characterizer/neural_network.py
@@ -9,10 +9,7 @@
 from .regression_model import regression_model
 from globals import OPTS
 import debug
-
-from tensorflow import keras
-from tensorflow.keras import layers
-import tensorflow as tf
+from sklearn.neural_network import MLPRegressor
 
 
 class neural_network(regression_model):
@@ -20,21 +17,19 @@ class neural_network(regression_model):
     def __init__(self, sram, spfile, corner):
         super().__init__(sram, spfile, corner)
 
+    def get_model(self):
+        return MLPRegressor(solver='lbfgs', alpha=1e-5,
+                            hidden_layer_sizes=(40, 40, 40, 40), random_state=1)
+
     def generate_model(self, features, labels):
         """
-        Supervised training of model.
+        Training multilayer model
         """
         
-        model = keras.Sequential([
-            layers.Dense(32, activation=tf.nn.relu, input_shape=[features.shape[1]]),
-            layers.Dense(32, activation=tf.nn.relu),
-            layers.Dense(32, activation=tf.nn.relu),
-            layers.Dense(1)
-        ])
-
-        optimizer = keras.optimizers.RMSprop(0.0099)
-        model.compile(loss='mean_squared_error', optimizer=optimizer)
-        model.fit(features, labels, epochs=100, verbose=0)
+        flat_labels = np.ravel(labels)
+        model = self.get_model()
+        model.fit(features, flat_labels)
+        
         return model
         
     def model_prediction(self, model, features):    
@@ -44,5 +39,6 @@ class neural_network(regression_model):
         """
         
         pred = model.predict(features)
-        return pred
+        reshape_pred = np.reshape(pred, (len(pred),1))
+        return reshape_pred
         
\ No newline at end of file