From 562ae83394f0eac97774a0a0981b4e3ac81be19f Mon Sep 17 00:00:00 2001
From: Christian Chapman-Bird <c.chapmanbird@btinternet.com>
Date: Mon, 27 Sep 2021 13:59:48 +0100
Subject: [PATCH] model load byte open fix, rmse flattens inputs

---
 EMRI_DET/nn/model_creation.py    | 4 ++--
 EMRI_DET/validate.py             | 3 +--
 emri_data/scripts/model_test.py  | 6 +++---
 emri_data/scripts/model_train.py | 2 +-
 4 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/EMRI_DET/nn/model_creation.py b/EMRI_DET/nn/model_creation.py
index d6d2ece..f39cd7d 100644
--- a/EMRI_DET/nn/model_creation.py
+++ b/EMRI_DET/nn/model_creation.py
@@ -47,8 +47,8 @@ def create_mlp(input_features, output_features, neurons, layers, activation, mod
     return model
 
 
-def load_mlp(model_name, get_state_dict=False):
+def load_mlp(model_name, device, get_state_dict=False):
     model = pickle.load(open(get_script_path()+f'/../models/{model_name}/function.pickle', "rb"))  # load blank model
     if get_state_dict:
-        model.load_state_dict(torch.load(open(get_script_path()+f'/../models/{model_name}/model.pth')))
+        model.load_state_dict(torch.load(open(get_script_path()+f'/../models/{model_name}/model.pth', "rb"), map_location=device))
     return model
diff --git a/EMRI_DET/validate.py b/EMRI_DET/validate.py
index aad5b99..190a63f 100644
--- a/EMRI_DET/validate.py
+++ b/EMRI_DET/validate.py
@@ -78,13 +78,12 @@ def compute_rmse(comparison_sets):
         rmse (double): root-mean-square-error between the truth and the model output for this dataset.
     """
     truth, pred = comparison_sets
-    rmse = np.sqrt(np.mean((truth - pred) ** 2))
+    rmse = np.sqrt(np.mean((truth.flatten() - pred.flatten()) ** 2))
     return rmse
 
 
 def test_threshold_accuracy(comparison_sets, threshold):
     truth, pred = comparison_sets
-
     out_classified = np.zeros(shape=pred.size)
     out_classified[pred.flatten() >= threshold] = 1
 
diff --git a/emri_data/scripts/model_test.py b/emri_data/scripts/model_test.py
index 1a0c48e..0f6ac89 100644
--- a/emri_data/scripts/model_test.py
+++ b/emri_data/scripts/model_test.py
@@ -9,7 +9,7 @@ import pandas as pd
 
 device = 'cuda:0'
 model_name = 'silu_1'
-mlp = load_mlp(model_name, get_state_dict=True).to(device)
+mlp = load_mlp(model_name, device, get_state_dict=True).to(device)
 mlp.eval()
 
 df = pd.read_csv('../schwarz_data/grid_dataframe.csv')
@@ -17,12 +17,12 @@ x_inds = [0, 1, 2, 3, 4, 5, 6, 7, 8]  # [0,1,2,4,5,6,7,8,9]
 y_inds = [9]
 
 xdata = df.iloc[:,x_inds].to_numpy()
-ydata = np.log(df.iloc[:,y_inds].to_numpy())
+ydata = np.log(df.iloc[:,y_inds].to_numpy() + 1e-5)
 
 xmeanstd = np.load(f'../models/{model_name}/xdata_mean_std.npy')
 ymeanstd = np.load(f'../models/{model_name}/ydata_mean_std.npy')
 
-net_out = run_on_dataset(mlp,[xdata, ydata],device=device,y_transform_fn=np.exp)
+net_out = run_on_dataset(mlp,[xdata, ydata],device=device,y_transform_fn=np.exp, n_batches=50)
 
 rmse = compute_rmse([np.exp(ydata),np.exp(net_out)])
 print('RMSE: ', rmse)
diff --git a/emri_data/scripts/model_train.py b/emri_data/scripts/model_train.py
index dcee102..812bb75 100644
--- a/emri_data/scripts/model_train.py
+++ b/emri_data/scripts/model_train.py
@@ -5,7 +5,7 @@ from EMRI_DET.nn.model_creation import create_mlp
 import pandas as pd
 
 if __name__ == '__main__':
-    device = "cuda:2"
+    device = "cuda:0"
     fp = '../schwarz_data/{}'
     #['logM', 'logq', 'a', 'p0', 'e', 'Y0', 'thetaS', 'phiS', 'thetaK', 't', 'SNR']
     train_inds = [0,1,2,3,4,5,6,7,8]#[0,1,2,4,5,6,7,8,9]
-- 
GitLab