Adding allow_growth option (deepfakes#140)

scripts/train.py

@@ -79,6 +79,11 @@ def parse_arguments(self, description, subparser, command):
                             type=int,
                             default=64,
                             help="Batch size, as a power of 2 (64, 128, 256, etc)")
+        parser.add_argument('-ag', '--allow-growth',
+                            action="store_true",
+                            dest="allow_growth",
+                            default=False,
+                            help="Sets allow_growth option of Tensorflow to spare memory on some configs")
         parser.add_argument('-ep', '--epochs',
                             type=int,
                             default=1000000,
@@ -122,6 +127,9 @@ def process(self):
         thr.join() # waits until thread finishes
 
     def processThread(self):
+        if self.arguments.allow_growth:
+            self.set_tf_allow_growth()
+
         print('Loading data, this may take a while...')
         # this is so that you can enter case insensitive values for trainer
         trainer = self.arguments.trainer
@@ -164,6 +172,14 @@ def processThread(self):
             print(e)
             exit(1)
 
+    def set_tf_allow_growth(self):
+        import tensorflow as tf
+        from keras.backend.tensorflow_backend import set_session
+        config = tf.ConfigProto()
+        config.gpu_options.allow_growth = True
+        config.gpu_options.visible_device_list="0"
+        set_session(tf.Session(config=config))
+
     preview_buffer = {}
 
     def show(self, image, name=''):