Pytorch – 没有使用with torch.no_grad()造成测试网络时显存爆炸的问题
1 显存爆炸的问题
最近使用以下示例代码测试自定义深度学习网络时耗光了所有显存,出现了梯度爆炸的问题。
model.eval()
for batch_idx, data in enumerate(tqdm(data_loader)):
image = data[0].to('cuda:0')
......
经过排查原因是没有加上with torch.no_grad()
语句停止梯度更新,从而导致了显存爆炸的问题,正确的示例代码如下
model.eval()
with torch.no_grad()
for batch_idx, data in enumerate(tqdm(data_loader)):
image = data[0].to('cuda:0')
......
2 model.train、model.eval和with torch.no_grad
model.train()
会将网络中的模块设置为训练模式,此时,如果神经网络中BN(batch normalization)层和Dropout层,那么这两个层将会起作用,防止网络出现过拟合的问题。
model.eval()
则会将网络设置为测试模式,此时,不会启用神经网络中的BN(batch normalization)层和Dropout层,model.eval()是保证BN层直接使用全部训练数据的均值和方差,即测试过程中要保证BN层的均值和方差不变。对于Dropout层,model.eval()是利用到了所有网络连接,即不进行随机舍弃神经元。
with torch.no_grad()
会将网络中Tensor的属性全部设置为False,并停止Autograd引擎,禁止梯度反向传播,以起到加速和节省显存的作用。它的作用是将该with语句包裹起来的部分停止梯度的更新,从而节省了GPU算力和显存,但是并不会影响dropout和BN层的行为。因此,测试的时候加上此语句也不会影响测试精度的,只是停止了梯度更新而已。在测试和验证阶段,使用with torch.no_grad()
会使得网络有更快的推理速度和内存使用,这使得我们在验证和测试网络时可以使用更大的batch size。
本文作者:StubbornHuang
版权声明:本文为站长原创文章,如果转载请注明原文链接!
原文标题:Pytorch – 没有使用with torch.no_grad()造成测试网络时显存爆炸的问题
原文链接:https://www.stubbornhuang.com/2322/
发布于:2022年08月23日 10:29:14
修改于:2023年06月21日 18:13:30
当前分类随机文章推荐
- Pytorch - torch.stack参数详解与使用 阅读1188次,点赞0次
- Pytorch - 使用torchsummary/torchsummaryX/torchinfo库打印模型结构、输出维度和参数信息 阅读2153次,点赞1次
- Pytorch - 使用Pyav解码视频文件并将视频帧转换为Pytorch tensor作为网络模型输入数据 阅读4576次,点赞0次
- Pytorch - 使用torch.onnx.export将Pytorch模型导出为ONNX模型 阅读8602次,点赞0次
- Pytorch - 内置的CTC损失函数torch.nn.CTCLoss参数详解与使用示例 阅读2259次,点赞1次
- Pytorch - 模型微调时删除原有模型中的某一层的方法 阅读3347次,点赞0次
- Pytorch – 使用torch.matmul()替换torch.einsum(‘nkctv,kvw->nctw’,(a,b))算子模式 阅读1391次,点赞0次
- Pytorch - torch.distributed.init_process_group函数详解 阅读1026次,点赞0次
- Pytorch - torch.topk参数详解与使用 阅读549次,点赞0次
- Pytorch - Pytoch结合Tensorboard实现数据可视化 阅读444次,点赞0次
全站随机文章推荐
- Python - 使用ffmepg批量转换某个文件夹以及所有子文件夹下所有的视频,修改其帧率/码率/分辨率到另一文件夹,并保留原有文件夹结构 阅读3298次,点赞0次
- Docker - 镜像操作入门指南 阅读68次,点赞0次
- C++ - 速通nlohmann json,nlohmann json使用教程 阅读186次,点赞0次
- Git - 断点续传git大项目 阅读41次,点赞0次
- 资源分享 - Python网络数据采集 (美 Ryan Mitchell著 陶俊杰 陈小莉译) 高清PDF下载 阅读3217次,点赞0次
- Unity - Color32[]转为byte[]字节数组 阅读3757次,点赞1次
- OpenCV - 在图像处理教程使用最广的测试美女图片Lenna,Lenna原图 阅读515次,点赞0次
- 资源分享 - Essential Mathematics for Games and Interactive Applications(First Edition) 英文高清PDF下载 阅读2270次,点赞0次
- 资源分享 - Digital Lighting & Rendering , First Edition 英文高清PDF下载 阅读1326次,点赞0次
- 资源分享 - Handbook of Digital Image Synthesis - Scientific Foundations of Rendering 英文高清PDF下载 阅读1582次,点赞0次
评论
169