pytorch代码中同时包含训练和测试代码时显存爆炸

原因在于没有使用torch.no_grad()函数。在查看验证集和测试集表现时,应使用类似这样的代码

def evaluate(data_loader):
    with torch.no_grad():
        mean_acc, mean_iou = 0, 0
        for i, (img, gnd) in enumerate(data_loader):
            if torch.cuda.is_available():
                img = img.cuda(device=device)
                gnd = gnd.cuda(device=device)
            out = model(img)
            .......
        return mean_acc / len(data_loader), mean_iou / len(data_loader)

posted @ 2018-12-02 13:01  技术流选手  阅读(1824)  评论(0编辑  收藏  举报