Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Release GPU memory after tensorflow session is closed
def run_release_gpu(func):
def parallel_wrapper(output_dict, *argv, **kwargs):
ret = func(*argv, **kwargs)
if ret is not None:
output_dict['ret'] = ret
def outer_wrapper(*argv, **kwargs):
same_process = kwargs.pop('same_process', False)
if same_process:
return func(*argv, **kwargs)
with multiprocessing.Manager() as manager:
output = manager.dict()
args = (output, ) + argv
p = multiprocessing.Process(target=parallel_wrapper, args=args, kwargs=kwargs)
p.start()
p.join()
ret_val = output.get('ret', None)
return ret_val
return outer_wrapper
@run_release_gpu
def run_computations(a, b, c):
with tf.Graph().as_default():
with tf.Session() as sess:
pass
a = run_computations(1, 2, 3)
@arainboldt

This comment has been minimized.

Copy link

@arainboldt arainboldt commented Nov 5, 2020

Hi, thanks for sharing this on stack overflow. Unfortunately, it still seems very relevant as a work-around, but I can't seem to get it to work as Process needs to pickle the target func, which in this case is local, and thus cannot be pickled. Did you ever get this to work?

@MInner

This comment has been minimized.

Copy link
Owner Author

@MInner MInner commented Nov 20, 2020

It's being a while since I posted this, no idea whether it is still a viable workaround.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment