ModelScan：防范模型序列化攻击

机器学习（ML）模型在互联网上公开共享，并在团队内部和团队之间共享。基础模型的兴起导致公开的 ML 模型越来越多地被用于进一步的训练/微调。ML 模型越来越多地被用来做出关键决策，并为关键任务应用提供支持。 尽管如此，模型尚未像收件箱中的 PDF 文件那样受到严格的扫描。

这种情况需要改变，而合适的工具是第一步。

ModelScan 是 Protect AI 的一个开源项目，用于扫描模型以确定它们是否包含不安全的代码。它是第一个支持多种模型格式的模型扫描工具。ModelScan 目前支持：H5、Pickle 和 SavedModel 格式。这可以保护你使用 PyTorch、TensorFlow、Keras、Sklearn、XGBoost 时的安全，更多格式即将推出。

TL;DR

如果你已准备好开始扫描模型，方法很简单：

pip install modelscan

安装完成后，扫描一个模型：

modelscan -p /path/to/model_file.pkl

为什么你应该扫描模型

模型通常由自动化管道创建，也可能来自数据科学家的笔记本电脑。无论哪种情况，模型在使用前都需要从一台机器移动到另一台机器。将模型保存到磁盘的过程称为序列化。

模型序列化攻击是指在分发之前，在序列化（保存）过程中将恶意代码添加到模型内容中——这是特洛伊木马的现代版本。

该攻击通过利用模型的保存和加载过程来实现。当你使用 model = torch.load(PATH) 加载模型时，PyTorch 会打开文件内容并开始运行其中的代码。在你加载模型的那一刻，漏洞利用就已经执行了。

模型序列化攻击可用于执行：

• 凭证窃取（用于在你的环境中向其他系统写入和读取数据的云凭证）
• 数据窃取（发送给模型的请求）
• 数据投毒（模型执行任务后发送的数据）
• 模型投毒（篡改模型本身的结果）

这些攻击执行起来非常简单，你可以在我们的 📓notebooks 文件夹中查看可运行的示例。

实施和自动化模型安全

ModelScan 提供强大的开源扫描能力。如果你需要全面的 AI 安全，请考虑 Guardian。它是我们的企业级模型扫描产品。

Guardian 的功能：

1. 尖端扫描：使用我们最新的扫描器、更广泛的模型支持和自动模型格式检测。
2. 主动安全：在 Hugging Face 模型进入你的环境之前定义并执行安全要求——无需更改代码。
3. 企业级覆盖：在整个组织中实施统一的安全态势，与你的 CI/CD 管道无缝集成。
4. 全面审计追踪：全面了解所有扫描和结果，使你能够有效识别和缓解威胁。

开始使用

ModelScan 的工作原理

如果使用你的机器学习框架加载模型会自动执行攻击，那么 ModelScan 如何在不加载恶意代码的情况下检查内容呢？

很简单，它像读取字符串一样逐字节读取文件内容，查找不安全的代码签名。这使得它非常快，扫描模型所需的时间仅相当于你的计算机从磁盘处理整个文件大小的时间（大多数情况下只需几秒）。而且它也是安全的。

ModelScan 将不安全代码分为以下等级：

• 严重
• 高
• 中
• 低

如果检测到问题，请立即联系模型的作者以确定原因。

在某些情况下，代码可能会嵌入模型中，以便数据科学家更容易重现结果，但这会让你面临攻击。请自行判断这是否适合你的工作负载。

支持哪些模型和框架？

这将持续扩展，请留意我们发布说明中的变更。

目前，ModelScan 支持任何基于 Pickle 的格式以及许多其他格式：

安装

ModelScan 作为 Python 包安装在你的系统上（支持 Python 3.9 至 3.12）。如上所示，你可以通过在终端中运行以下命令来安装它：

pip install modelscan

要将其包含在项目的依赖项中，以便所有人都可以使用，请将其添加到你的 requirements.txt 或 pyproject.toml 中，如下所示：

modelscan = ">=0.1.1"

用于 Tensorflow 或 HD5 格式模型的扫描器需要通过 extras 安装：

pip install 'modelscan[ tensorflow, h5py ]'

通过 CLI 使用 ModelScan

ModelScan 通过 CLI 支持以下参数：

请记住，模型就像任何其他形式的数字媒体一样，在使用前应扫描来自任何不可信来源的内容。

CLI 退出代码

CLI 的退出状态代码如下：

• 0：扫描成功完成，未发现漏洞
• 1：扫描成功完成，发现漏洞
• 2：扫描失败，modelscan 在扫描期间抛出错误
• 3：没有向工具传递受支持的文件
• 4：用法错误，CLI 收到无效或不完整的选项

在 Python 中以编程方式使用 ModelScan

虽然 ModelScan 可以轻松地通过 CLI 使用，但你也可以将其直接集成到你的 Python 应用程序或工作流中。

from modelscan.modelscan import ModelScan
from modelscan.settings import DEFAULT_SETTINGS

# Initialize ModelScan with default settings
scanner = ModelScan(settings=DEFAULT_SETTINGS)

# Scan a model file or directory 
results = scanner.scan("/path/to/model_file.pkl")

# Check if issues were found
if scanner.issues.all_issues:
    print(f"Found {len(scanner.issues.all_issues)} issues!")
    
    # Access issues by severity
    issues_by_severity = scanner.issues.group_by_severity()
    for severity, issues in issues_by_severity.items():
        print(f"{severity}: {len(issues)} issues")
        
# Generate a report (default is console output)
scanner.generate_report()

你可以使用自己的设置自定义扫描行为：

# Start with default settings and customize
custom_settings = DEFAULT_SETTINGS.copy()

# Update settings as needed
custom_settings["reporting"]["module"] = "modelscan.reporting.json_report.JSONReport"
custom_settings["reporting"]["settings"]["output_file"] = "scan_results.json"

# Initialize with custom settings
scanner = ModelScan(settings=custom_settings)

理解结果

扫描完成后，如果发现问题，你将看到如下输出：

这里我们有一个模型，其中包含 ReadFile 和 WriteFile 的不安全操作符。显然，我们不希望模型随意读写文件。我们现在应该联系该模型的创建者，确定他们期望这个功能做什么。在这种特定情况下，它允许攻击者读取我们的 AWS 凭证并将其写入另一个位置。

这绝对禁止使用。

将 ModelScan 集成到你的 ML 管道和 CI/CD 管道中

临时扫描是很好的第一步，请让自己、同事和朋友养成习惯，每当我们拉取一个新模型进行探索时都进行扫描。但这还不足以提高生产 MLOps 流程的安全性。

模型扫描需要执行多次，以实现以下目标：

1. 在加载所有预训练模型进行进一步工作之前扫描它们，防止受损模型影响你的模型构建或数据科学环境。
2. 在训练后扫描所有模型，以检测危害新模型的供应链攻击。
3. 在部署到端点之前扫描所有模型，确保模型在存储后未被入侵。

下面的红色模块在传统 ML 管道中突出了这一点。

对于微调或对 LLM、基础模型或外部模型的任何修改，流程也是相同的。

如果模型部署在 ML 管道之外，也请将扫描嵌入 CI/CD 系统的部署过程中，以确保模型在部署时的使用安全。

深入探讨

在 📓notebooks 文件夹中，你可以探索多个 notebook，它们展示了如何针对 TensorFlow 和 PyTorch 等各种 ML 框架执行模型序列化攻击。

要更深入地了解这些攻击究竟是如何运作的，请查看 🖹 模型序列化攻击解释器。

如果你遇到其他在静态上下文中评估模型的方法，请联系我们，我们很乐意了解更多！

许可证

版权所有 2024 Protect AI

根据 Apache 许可证 2.0 版（“许可证”）获得许可； 除非遵守许可证，否则你不得使用此文件。 你可以在以下位置获取许可证副本：

http://www.apache.org/licenses/LICENSE-2.0

除非适用法律要求或书面同意，否则根据许可证分发的软件按“原样”基础分发，不附带任何明示或暗示的保证或条件。请参阅许可证，了解许可证项下管理权限和限制的特定语言。

致谢

我们深受 Matthieu Maitre 的启发，他构建了 PickleScan。我们感谢这项工作，并通过 ModelScan 对其进行了重大扩展。ModelScan 以与 PickleScan 相似的精神开源。

贡献

我们非常欢迎你为我们的开源 ModelScan 项目做出贡献。如果你想贡献，请查看 贡献页面 上的详细信息。