add models folder

shirecoding · Aug 7, 2020 · 1f0fdcd · 1f0fdcd
1 parent 1e7ead1
commit 1f0fdcd
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/modelmaker/resources/templates/default/models/README.md b/modelmaker/resources/templates/default/models/README.md
@@ -0,0 +1 @@
+save trained models here
diff --git a/modelmaker/resources/templates/default/scripts/train.py b/modelmaker/resources/templates/default/scripts/train.py
@@ -1,3 +1,4 @@
+import os
 import numpy as np
 from tensorflow import keras
 from {{ package_name }}.models import SimpleClassification
@@ -46,7 +47,7 @@ def _generate_random_shape(img_size):
 
 # training parameters
 batch_size = 32
-epochs = 50
+epochs = 5
 
 # generate training data
 train_data = ShapesClassDatagen(400, batch_size)
@@ -63,4 +64,12 @@ def _generate_random_shape(img_size):
 model.fit(train_data, steps_per_epoch=len(train_data), epochs=epochs)
 
 # save model
-model.save("/Users/benjaminhon/simple_model")
+model_folder = os.path.join(
+    os.path.dirname(
+        os.path.dirname(
+            os.path.realpath(__file__)
+        )
+    ),
+    "models"
+)
+model.save(os.path.join(model_folder, 'simple_model'))