PyKEEN项目中关于torch.load安全加载的优化解析

在PyKEEN项目的最新开发中，团队针对PyTorch 2.5及以上版本中torch.load函数的安全性问题进行了重要优化。本文将深入分析这一技术改进的背景、解决方案及其对项目安全性的提升。

背景分析

随着PyTorch 2.5版本的发布，其核心团队对模型加载机制进行了重要的安全升级。在之前的版本中，torch.load默认使用weights_only=False参数，这意味着它会隐式使用Python的标准pickle模块来反序列化数据。这种机制存在潜在的安全风险，因为pickle可以执行任意代码，恶意构造的模型文件可能导致代码注入攻击。

PyTorch官方安全文档明确指出，未来版本将会把weights_only参数的默认值改为True，这将限制反序列化过程中可执行的函数范围，显著提升安全性。

技术解决方案

PyKEEN开发团队及时响应这一变化，在项目代码中进行了以下关键改进：

1. 参数显式设置：在所有使用torch.load的地方显式添加了weights_only=True参数，确保代码行为在未来PyTorch版本更新后保持一致。

2. 安全全局变量管理：对于需要加载自定义对象的情况，团队使用了torch.serialization.add_safe_globals方法来明确允许特定的全局变量，在保证安全性的同时不损失功能性。

3. 前瞻性兼容：这些修改确保了PyKEEN项目能够平滑过渡到未来PyTorch版本中weights_only默认值变更的情况。

安全意义

这一改进对PyKEEN项目的安全性带来了显著提升：

• 防范代码注入：通过限制反序列化过程中可执行的代码范围，有效防止了潜在的恶意模型攻击。
• 明确安全边界：通过显式声明安全全局变量，建立了清晰的安全边界，使项目维护更加可控。
• 符合安全最佳实践：紧跟PyTorch官方的安全建议，体现了项目对安全性的高度重视。

开发者建议

对于使用PyKEEN的开发者，建议：

1. 及时更新到包含此优化的版本，以获得更好的安全性保障。
2. 在自定义模型加载逻辑时，遵循相同的安全实践，显式设置weights_only参数。
3. 对于需要加载自定义对象的情况，谨慎使用add_safe_globals方法，仅添加必要的安全全局变量。

这一优化体现了PyKEEN项目对安全性和前瞻性的重视，为知识图谱嵌入领域的开发者提供了更加安全可靠的工具基础。