Release GPU memory after tensorflow session is closed

run_release_gpu.py

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)

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?

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