(1) In PyTorch 2.6, we changed the default value of the `weights_only` argument in `torch.load` from `False` to `True`. Re-running `torch.load` with `weights_only` set to `False` will likely succeed, but it can result in arbitrary code execution. Do it on
时间: 2025-03-30 20:09:12 浏览: 121
### PyTorch 2.6 中 `torch.load` 函数 `weights_only` 参数的影响
在 PyTorch 的开发过程中,安全性一直是开发者关注的重点之一。特别是在处理模型保存与加载的过程中,`torch.load` 是一个常用的功能,但它也伴随着一定的安全风险。
#### 背景说明
当使用 `torch.load()` 加载模型时,默认情况下会依赖 Python 的 `pickle` 模块来反序列化数据[^3]。然而,`pickle` 模块本身具有执行任意代码的能力,这意味着如果加载的数据来自不可信的源,则可能引发严重的安全隐患。为了应对这一问题,在未来的 PyTorch 版本中(如假设中的 2.6),计划将 `weights_only` 默认设置为 `True`,从而仅允许加载权重而忽略其他对象。
#### 影响分析
当前版本中,`weights_only=False` 表示不限制加载的内容范围,除了模型权重外还可能加载额外的对象或元数据。这种行为虽然提供了更大的灵活性,但也增加了被恶意利用的风险。一旦 `weights_only` 设为 `True`,则只会加载张量形式的权重数据,其余部分会被丢弃。对于某些特定场景下的复杂模型结构恢复来说,这可能导致兼容性问题。
#### 如何避免安全风险?
为了避免上述提到的由于使用 `pickle` 带来的潜在威胁,可以采取如下措施:
1. **显式指定参数**
在调用 `torch.load` 方法时,应始终明确设定 `weights_only=True` 来增强程序运行期间的安全保障水平。
2. **验证输入文件来源合法性**
确保所要读取的 `.pt`, `.pth` 文件来自于可信赖渠道或者经过严格校验后再投入使用。
3. **更新至最新稳定版框架**
随着时间推移,官方团队持续改进和完善 API 功能及其内部实现逻辑;及时升级到新发行版本有助于获得更优性能表现的同时享受更多防护机制保护。
以下是调整后的代码实例展示如何正确配置该选项:
```python
import torch
# 明确指明 weights_only 参数以提高安全性
model_state_dict = torch.load('path/to/model.pth', map_location=torch.device('cpu'), weights_only=True)
# 如果需要重建整个模型而非仅仅导入其状态字典,请单独定义好对应类构造器之后再操作
```
通过以上方式能够有效降低因不当使用而导致意外后果发生的概率。
阅读全文
相关推荐








