simcop2387/faceswap
1
0
Fork 0
mirror of https://github.com/deepfakes/faceswap synced 2025-06-07 10:43:27 -04:00
Code Issues Releases Wiki Activity

Adding allow_growth option (#140)

Browse source
This commit is contained in:
Clorr 2018-02-14 16:23:21 +01:00 • committed by GitHub
parent 0085b5b8a7
commit 51f1993d93
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 16 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

16
scripts/train.py
View file

@ -79,6 +79,11 @@ class TrainingProcessor(object):
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 @@ class TrainingProcessor(object):
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 @@ class TrainingProcessor(object):
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=''):