Pytorch 训练与测试时爆显存 out of memory 的一个解决方案

首先给大家推荐一下我老师大神的人工智能教学网站。教学不仅零基础,通俗易懂,而且非常风趣幽默,还时不时有内涵黄段子!点这里可以跳转到网站

Pytorch 训练时有时候会因为加载的东西过多而爆显存,有些时候这种情况还可以使用cuda的清理技术进行修整,当然如果模型实在太大,那也没办法。

使用torch.cuda.empty_cache()删除一些不需要的变量代码示例如下:

try:    output = model(input)except RuntimeError as exception:    if "out of memory" in str(exception):        print("WARNING: out of memory")        if hasattr(torch.cuda, 'empty_cache'):            torch.cuda.empty_cache()    else:        raise exception

测试的时候爆显存有可能是忘记设置no_grad, 示例代码如下:

    with torch.no_grad():        for ii,(inputs,filelist) in tqdm(enumerate(test_loader), desc='predict'):            if opt.use_gpu:                inputs = inputs.cuda()                if len(inputs.shape) < 4:                    inputs = inputs.unsqueeze(1)             else:                if len(inputs.shape) < 4:                    inputs = torch.transpose(inputs, 1, 2)                    inputs = inputs.unsqueeze(1)

点这里可以跳转到人工智能网站

发表评论