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

最新推荐文章于 2024-07-28 09:34:41 发布

转载最新推荐文章于 2024-07-28 09:34:41 发布 · 382 阅读

0 ·

CC 4.0 BY-SA版权

原文链接：http://www.cnblogs.com/liuzhan709/p/10053009.html

文章标签：

#人工智能 #python

本文介绍如何使用torch.no_grad()函数提高模型在验证集和测试集上的评估效率。通过在不需要梯度计算的情况下禁用自动求导，可以显著减少内存消耗和加速计算过程。文章提供了一个具体的代码示例，展示了如何在PyTorch中实现这一优化。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

原因在于没有使用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)

转载于:https://www.cnblogs.com/liuzhan709/p/10053009.html