prehistoric-systems/comixify
1
0
Fork 1
mirror of https://github.com/maciej3031/comixify.git synced 2026-03-11 08:54:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix GPU memory not released

Browse source
This commit is contained in:
Maciej Pęśko 2018-11-02 22:47:00 +01:00
parent ce0e178763
commit 7fcb951b1c
1 changed files with 10 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
style_transfer/style_transfer.py
View file

@ -11,6 +11,8 @@ from torch.autograd import Variable
from keras_contrib.layers import InstanceNormalization
from CartoonGAN.network.Transformer import Transformer
from utils import profile
import gc
from keras import backend as K
class StyleTransfer():
@ -43,14 +45,14 @@ class StyleTransfer():
@classmethod
def _comix_gan_stylize(cls, frames):
# comixGAN_cache_key = 'comixGAN_model_cache'
# comixGAN_model = cache.get(comixGAN_cache_key) # get model from cache
comixGAN_cache_key = 'comixGAN_model_cache'
comixGAN_model = cache.get(comixGAN_cache_key) # get model from cache
# if comixGAN_model is None:
if comixGAN_model is None:
# load pretrained model
comixGAN_model = load_model(settings.COMIX_GAN_MODEL_PATH,
custom_objects={'InstanceNormalization': InstanceNormalization})
# cache.set(comixGAN_cache_key, comixGAN_model, None) # None is the timeout parameter. It means cache forever
comixGAN_model = load_model(settings.COMIX_GAN_MODEL_PATH,
custom_objects={'InstanceNormalization': InstanceNormalization})
cache.set(comixGAN_cache_key, comixGAN_model, None) # None is the timeout parameter. It means cache forever
frames = cls._resize_images(frames, size=450)
batch_size = 5
@ -59,7 +61,8 @@ class StyleTransfer():
batch_of_frames = np.stack(frames[i:i+batch_size]) / 255
stylized_batch_of_imgs = comixGAN_model.predict(batch_of_frames)
stylized_imgs.append(255*stylized_batch_of_imgs)
K.clear_session()
gc.collect()
return list(np.concatenate(stylized_imgs, axis=0))
@classmethod