DeepLabCut中PyTorch模型加载的安全警告及解决方案

背景介绍

在计算机视觉和深度学习领域，DeepLabCut是一个广泛使用的开源工具包，用于动物姿态估计和行为分析。近期，使用DeepLabCut PyTorch版本(rc4)运行analyze_videos时，系统会显示一个关于torch.load函数的安全警告。

问题分析

这个警告源于PyTorch 2.4版本引入的一项安全改进。当使用torch.load加载模型时，如果weights_only参数设置为False(当前默认值)，系统会使用Python的pickle模块，这存在潜在的安全风险。恶意构造的pickle数据可能在反序列化过程中执行任意代码。

PyTorch官方计划在未来版本中将weights_only的默认值改为True，这将限制反序列化过程中可执行的函数。除非用户通过torch.serialization.add_safe_globals明确允许，否则将不允许加载任意对象。

解决方案

对于DeepLabCut用户，有以下几种解决方案：

1. 环境变量设置：在PyTorch 2.4及以上版本中，可以通过设置环境变量强制使用weights_only模式：

import os
os.environ["TORCH_FORCE_WEIGHTS_ONLY_LOAD"] = "1"

2. 代码修改：等待DeepLabCut官方在正式发布前更新代码，直接指定weights_only=True参数。

3. 安全实践：仅加载自己创建的模型检查点，不要加载来自不可信来源的模型文件。

技术细节

weights_only模式的工作原理是限制反序列化过程中可加载的对象类型，只允许加载张量、列表、字典等基本数据结构，而禁止加载任意Python对象。这种模式显著提高了模型加载的安全性，特别是在共享模型或从外部来源获取模型时。

最佳实践建议

1. 对于研究机构和企业的安全计算环境，建议立即采用weights_only模式
2. 在模型共享时，告知接收方模型的来源和完整性验证方法
3. 定期更新DeepLabCut和PyTorch到最新版本，以获取最新的安全修复

总结

这个安全警告反映了深度学习社区对模型安全性的日益重视。虽然当前DeepLabCut版本仍使用默认的加载方式，但用户可以通过简单的方法提前适应这一变化。理解并实施这些安全措施对于保护研究数据和计算环境至关重要。