Home > AI

在异常终止CUDA程序的时候, 自动回收GPU显存

linker   ·   发表于 27天前   ·   AI
# 在异常终止CUDA程序的时候, 自动回收GPU显存

import functools, traceback, sys
def gpu_mem_restore(func):
"Reclaim GPU RAM if CUDA out of memory happened, or execution was interrupted"
@functools.wraps(func)
def wrapper(*args, **kwargs):
try:
return func(*args, **kwargs)
except:
type, val, tb = sys.exc_info()
traceback.clear_frames(tb)
raise type(val).with_traceback(tb) from None
return wrapper

@gpu_mem_restore
def fit():
raise Exception("fake fit...")

if __name__ == '__main__':
fit()


3 Reply   |  Until 27天前 | 26 View

linker
发表于 27天前

查看cuDNN版本:

torch.backends.cudnn.version()

评论列表

  • 加载数据中...

编写评论内容

linker
发表于 27天前

开启cuDNN@pytorch

torch.backends.cudnn.benchmark=True

评论列表

  • 加载数据中...

编写评论内容

linker
发表于 27天前

if args.use_gpu and torch.cuda.is_available():
    device = torch.device('cuda')
    torch.backends.cudnn.benchmark = True
else:
    device = torch.device('cpu')

评论列表

  • 加载数据中...

编写评论内容
LoginCan Publish Content
程序耗时 0.0241秒 内存:372KB