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)


免責聲明!

本站轉載的文章為個人學習借鑒使用,本站對版權不負任何法律責任。如果侵犯了您的隱私權益,請聯系本站郵箱yoyou2525@163.com刪除。



 
粵ICP備18138465號   © 2018-2025 CODEPRJ.COM