PyTorch-Image-Models 中 torch.load 的安全加载机制解析

背景介绍

在深度学习模型训练和推理过程中，模型权重的加载是一个关键环节。PyTorch 框架提供了 torch.load() 函数用于加载保存的模型检查点。随着 PyTorch 2.4.0 版本的发布，该函数引入了一个重要的安全特性变更，这对 PyTorch-Image-Models (timm) 库的使用产生了直接影响。

问题现象

当用户将 PyTorch 升级到 2.4.0 版本后，在使用 timm 库加载模型权重时会收到如下警告信息：

You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling...

这个警告信息明确指出，当前默认的 weights_only=False 设置存在潜在安全风险，因为在反序列化过程中可能会执行恶意代码。PyTorch 团队计划在未来的版本中将默认值改为 True，以增强安全性。

技术原理

weights_only 参数是 PyTorch 引入的一项重要安全特性：

1. 安全模式(weights_only=True)：仅允许加载包含张量、数字、字符串、列表和字典等基本数据类型的检查点文件，禁止加载任意 Python 对象
2. 非安全模式(weights_only=False)：使用 Python 的 pickle 模块进行完全反序列化，可能执行恶意代码

PyTorch 团队建议对所有不受信任的模型文件使用安全模式加载。对于 timm 库而言，所有官方提供的模型检查点都只包含权重数据，因此完全可以安全地使用 weights_only=True 模式。

兼容性考量

在实现这一变更时，开发团队面临一个重要挑战：向后兼容性。因为 weights_only 参数是在较新的 PyTorch 版本中引入的，旧版本中并不存在这个参数。直接添加该参数会导致旧版本 PyTorch 抛出参数不存在的错误。

解决方案

经过讨论，timm 库采用了以下稳健的解决方案：

1. 使用 try-except 块来检测当前 PyTorch 版本是否支持 weights_only 参数
2. 对于支持的版本，显式设置 weights_only=True 以启用安全模式
3. 对于不支持的旧版本，回退到原始加载方式

这种实现方式既解决了新版本中的警告问题，又确保了与旧版本的兼容性，同时遵循了安全最佳实践。

实际影响

这一变更对用户的主要影响包括：

1. 消除了冗长的安全警告，使输出更加简洁
2. 增强了模型加载过程的安全性
3. 确保了对旧版本 PyTorch 的兼容性

值得注意的是，在使用 weights_only=True 模式时，如果检查点文件中包含训练状态等复杂对象（如优化器状态），可能会导致部分数据无法加载。这在模型推理场景下通常不是问题，但在恢复训练时可能需要特别注意。

最佳实践

基于这一变更，建议 timm 库用户：

1. 对于纯推理场景，优先使用 weights_only=True 模式
2. 对于训练恢复场景，确认检查点文件内容后选择合适的加载模式
3. 定期更新 timm 库以获取最新的安全改进

这一改进体现了 timm 库对安全性和用户体验的持续关注，同时也展示了如何优雅地处理框架版本演进带来的兼容性挑战。