added retraining function

luickk · Sep 21, 2017 · 8744d1a · 8744d1a
1 parent 6163010
commit 8744d1a
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/train.py b/train.py
@@ -16,6 +16,7 @@
 EPOCHS = 20000
 
 MODEL_NAME = 'Pete'
+MODEL_TRAIN = 'model_data/Pete'
 
 #training data path
 path = 'D:/data_ai/'
@@ -24,6 +25,9 @@ def main():
     model = cnn_model.googlenet(WIDTH, HEIGHT, 3, LR, output=8, model_name=MODEL_NAME)
 
 
+    if MODEL_TRAIN:
+        model.load(MODEL_TRAIN)
+
     files = sorted(glob.glob('{}/*.npy*'.format(path)), key=data_prog.numericalSort)
     for e in tqdm(range(EPOCHS)):
         for f in files: