training scripts fixed

Nolte-Group · Jan 3, 2025 · a48f78b · a48f78b
1 parent 5a26e4e
commit a48f78b
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/train.py b/train.py
@@ -94,6 +94,12 @@ def get_data(spark, split=[0.99, 0.005, 0.005]):
     testy = np.array(test_df.select("treatment_index").collect()).astype(int)
     _testy = np.zeros((len(testy), int(index_max+1)))
     for index, value in enumerate(testy):
+        _testy[index, value] = 1.
+    testy = _testy
+    del _testy
+
+    return (selected, (trainx, trainy), (valx, valy), (testx, testy))
+
 
 def main():
     # jax mesh setup
@@ -146,7 +152,7 @@ def main():
                                    np.argmax(test_set[1], axis=1))
     plt.imshow(conf_matrix, origin="upper")
     plt.gca().set_aspect("equal")
-    plt.savefig("confusion_matrix.png")
+    plt.savefig("/app/workdir/confusion_matrix.png")
 
     return
 

diff --git a/train_cpu.py b/train_cpu.py
@@ -140,7 +140,7 @@ def main():
                                    np.argmax(test_set[1], axis=1))
     plt.imshow(conf_matrix, origin="upper")
     plt.gca().set_aspect("equal")
-    plt.savefig("confusion_matrix.png")
+    plt.savefig("/app/workdir/confusion_matrix.png")
 
     return