Bandit安全工具对PyTorch模型加载的误报问题解析

在Python安全审计工具Bandit的使用过程中，我们发现其对PyTorch模型加载函数torch.load()存在一个值得商榷的检测逻辑。本文将深入分析该问题的技术背景、安全考量以及最佳实践方案。

问题背景

Bandit的安全规则B614(pytorch_load_save)会标记所有使用torch.load()的代码，认为这可能存在反序列化安全风险。然而，当开发者明确设置weights_only=True参数时，PyTorch实际上已经内置了安全机制。

技术原理

PyTorch的torch.load()函数在2.6版本之前默认允许加载任意Python对象，这确实可能带来安全风险。但从2.6版本开始，PyTorch引入了weights_only参数：

1. 当weights_only=True时，加载器仅允许反序列化已知安全的类型（如张量、数字等基础类型）
2. 该参数在PyTorch 2.6+版本中默认值为True
3. 安全模式会严格限制可加载的对象类型，防止恶意代码执行

安全建议

基于PyTorch的安全机制演进，我们建议：

1. 对于PyTorch 2.6+版本：

   • 可以安全使用默认参数
   • Bandit无需标记任何torch.load()调用

2. 对于旧版本：

   • 必须显式设置weights_only=True
   • Bandit应仅标记未设置该参数或设为False的情况

3. 通用最佳实践：

   • 尽量升级到PyTorch 2.6+版本
   • 加载模型时显式声明安全参数
   • 仅加载可信来源的模型文件

实现方案

Bandit检测逻辑应调整为：

• 检查torch.load()调用是否包含weights_only参数
• 当参数存在且为True时，不触发告警
• 其他情况（参数缺失或为False）触发B614警告

这种改进既保持了安全审计的严谨性，又避免了过度告警对开发体验的影响。

总结

安全工具需要与时俱进地理解底层框架的安全机制。PyTorch通过weights_only参数已经解决了模型加载的主要安全隐患，安全工具应该据此优化检测逻辑，在保障安全性的同时减少误报，帮助开发者更高效地编写安全代码。