ModelScan：防御模型序列化攻击的工具

机器学习（ML）模型在互联网上公开共享，在团队内部和团队之间传播。基础模型的兴起导致公共ML模型越来越多地被用于进一步训练/微调。ML模型越来越多地被用于做出关键决策和支持关键任务应用。

尽管如此，模型还没有像扫描收件箱中的PDF文件那样严格地被扫描。

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

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

简介

如果您准备开始扫描您的模型，很简单：

pip install modelscan

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

modelscan -p /path/to/model_file.pkl

为什么要扫描模型

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

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

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

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

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

这些攻击非常容易执行，您可以在我们的📓notebooks文件夹中查看工作示例。

入门

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格式的模型扫描器，需要安装额外的依赖：

pip install 'modelscan[ tensorflow, h5py ]'

通过CLI使用ModelScan

ModelScan通过CLI支持以下参数：

请记住，模型和其他任何形式的数字媒体一样，在使用来自不可信来源的内容之前，您应该对其进行扫描。

CLI退出代码

CLI退出状态代码如下：

• 0：扫描成功完成，未发现漏洞
• 1：扫描成功完成，发现漏洞
• 2：扫描失败，modelscan在扫描过程中抛出错误
• 3：未向工具传递支持的文件
• 4：使用错误，CLI传递了无效或不完整的选项

理解结果

一旦扫描完成，如果发现问题，您将看到类似这样的输出：

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

这对于使用来说是坚决不行的。

在您的ML流程和CI/CD流程中集成ModelScan

临时扫描是一个很好的第一步，请让自己、同事和朋友养成习惯，每当他们下载新模型进行探索时都要这样做。但这对于改善生产MLOps流程的安全性还不够。

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

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

下图中的红色块在传统ML流程中突出显示了这一点。

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

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

深入探讨

在📓notebooks文件夹中，您可以探索一些展示如何针对各种ML框架（如TensorFlow和PyTorch）执行模型序列化攻击的笔记本。

要深入了解这些攻击的具体工作原理，请查看🖹模型序列化攻击解释器。

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

许可

版权所有 2023 Protect AI

根据Apache许可证2.0版（"许可证"）获得许可； 除非符合许可证的规定，否则您不得使用此文件。 您可以在以下位置获取许可证的副本：

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

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

致谢

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

贡献

我们非常欢迎您为我们的开源ModelScan项目做出贡献。 如果您想贡献，请遵循贡献页面上的详细说明。