ML-CUDA out of memory (torch.load())解决

本文介绍了在使用Pytorch时遇到CUDA内存不足问题的解决方法。当使用`torch.load()`加载模型时,由于.pth文件保存了训练时所用GPU的信息,即使更改设备设置,仍可能导致错误。解决方案是在加载模型时指定`map_location='cuda:0'`,将模型映射到目标GPU,从而避免错误。此方法帮助博主成功解决了跨GPU加载模型导致的内存冲突问题。

ML-CUDA out of memory (torch.load())解决

CUDA out of memory-torch.load

1. 杂话

 在上次的CUDA out of memory的问题解决之后(CUDA out of memory (Training)解决),这次又是遇到了新的问题。
 使用Pytorch的朋友可能都知道,我们是可以使用torch.load()来加载我们之前训练好的模型的。

2. 问题

 于是今天,我在使用torch.load()的时候就又碰到了爆显存的问题

 但是我仔细看了一下我的代码,我在加载模型的时候,好像也没有用到那块被占用的显卡啊
(自从上次的报错之后,我和实验室的另一位同学达成了一人使用一块显卡的和平相处协议)

 感谢这位大哥的帖子,帮我解决了问题pytorch加载模型torch.load报cuda out of memory解决方案

3. 解决

 简单来说就是,在torch.save()得到的pth文件里,会保存着你之前是在那张显卡训练这个模型的信息,而你在加载这个模型的时候,就会经过这张显卡,这个时候,哪怕你最终不是使用这张显卡计算,但是这张显卡现在是被占用的话,就会报CUDA out of memory的错误。。。

 也就是说,因为我和大哥一人一张显卡,之前我一直都是在cuda:1进行训练的,但是今天1卡被大哥用了,就算我把device改为cuda:0,在加载模型的时候也还是会报错,这个时候你就需要在代码中加入一个参数,来直接把模型映射到你要使用的显卡上

于是代码就从
torch.
根据Mend.io提供的安全报告,PyTorch 2.3.0版本存在多个高危漏洞(CVE-2024-48063等),主要涉及反序列化漏洞、内存损坏和拒绝服务攻击。以下是解决方案和详细说明: --- ### **1. 漏洞影响分析** | CVE编号 | 严重性 | 漏洞类型 | 影响范围 | |------------------|--------|------------------------------|------------------------| | CVE-2024-48063 | 9.8 | 反序列化未信任数据 | 所有PyTorch 2.3.0用户 | | CVE-2025-2999/8 | 5.3 | 内存损坏(RNN工具函数) | `unpack_sequence`等函数| | CVE-2025-3001/0 | 5.3 | 内存损坏(LSTM/JIT模块) | `torch.lstm_cell`等 | | CVE-2025-2148 | 5.0 | 内存损坏(Profiler回调) | JIT调试功能 | | CVE-2025-3136/21 | 3.3 | 内存损坏(CUDA/JIT模块) | GPU内存分配、模型加载 | | CVE-2025-2953 | 3.3 | 拒绝服务(MKL-DNN池化层) | 量化推理场景 | --- ### **2. 修复方案** #### **方案1:升级到安全版本** - **最新稳定版**:PyTorch 2.4.0+(已修复CVE-2024-48063等) ```bash pip install --upgrade torch>=2.4.0 ``` - **验证安装**: ```python import torch print(torch.__version__) # 应输出≥2.4.0 ``` #### **方案2:临时缓解措施(若无法升级)** - **禁用高危功能**: - 避免使用`torch.load()`加载不可信模型(CVE-2024-48063)。 - 替换RNN相关函数为安全实现: ```python # 替代pad_packed_sequence(示例) from torch.nn.utils.rnn import PackedSequence def safe_pad(packed_seq): return packed_seq[0] # 需根据实际需求调整 ``` - **限制CUDA使用**: ```python import os os.environ['CUDA_VISIBLE_DEVICES'] = '' # 强制使用CPU ``` #### **方案3:依赖隔离** - 使用虚拟环境隔离问题版本: ```bash conda create -n pytorch_safe python=3.9 conda activate pytorch_safe pip install torch==2.4.0 ``` --- ### **3. 长期安全建议** 1. **启用依赖项扫描**: - 集成Mend.io、Snyk等工具到CI/CD流程。 - 示例GitHub Action配置: ```yaml - name: Scan for vulnerabilities uses: snyk/actions/python@master with: command: test ``` 2. **监控安全公告**: - 订阅[PyTorch安全公告](https://github.com/pytorch/pytorch/security/advisories)。 3. **最小化权限**: - 运行PyTorch服务时使用非root用户。 ---
评论 1
成就一亿技术人!
拼手气红包6.0元
还能输入1000个字符
 
红包 添加红包
表情包 插入表情
 条评论被折叠 查看
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值