将 PyTorch 模型转换为 Core ML

大家好 我是 Steve 我是 Apple 的一名工程师 嗨 我是 Paul 我也是一名工程师 这本视频中 我们要详细为大家介绍 Core ML 的新内容之一 将 PyTorch 模型转化为 Core ML

在 WWDC 2020 上 我们宣布了对 Core ML 转换器进行改造 这改善了转换过程的许多方面

我们扩展了对资料库最常用的 深度学习社区的支持 我们重新设计了转换器架构 来改善用户体验 利用新的内存表示 我们统一了 API 所以只有一个调用可以从模型源 唤出转换指令 如果你已经看过了 我非常推荐你看看 讲解新转换器架构的细节视频

但在这段视频中 我要专注于模型转换 从新的 PyTorch 深度学习框架 模型构建开始

所以你可能是一名 ML 工程师 你一直在努力 用 PyTorch 来训练模型 你可能是一名 app 开发者 在线上发现了一个厉害的 PyTorch 模型 现在你想把这个模型用到 app 里

现在的问题是 你要怎样将 PyTorch 模型 转换为 Core ML 模型？

旧的 Core ML 转换器 会要求你将模型导出到 ONNX 作为过程中的一个步骤 如果你用了这个转换器 就可能会碰到一些局限 ONNX 是一种开放标准 开发和引入新功能的速度可能很慢 更棘手的是 像 PyTorch 这样的 ML 框架 将最新的模型功能导出到 ONNX 需要时间来添加支持 所以在用旧的转换器时 你可能会发现自己无法将 PyTorch 模型 导到 ONNX 上 向 Core ML 的转换也就无从谈起了

去除掉这个额外的依赖项 仅仅是新 Core ML 转换器中的 其中一个变动

所以在本视频中 我们将会详细了解全新 PyTorch 模型 转换路径的细节 我们将会介绍将 PyTorch 模型 转化为 Core ML 的 不同方式 包括一些实际的转换示例 最后我还会分享一些有用的建议 供你遵循 以防你在过程中遇到困难

那么 现在我们来深入了解一下 新的转换过程

从你希望转换的 PyTorch 模型开始 你将用到 PyTorch 的 JIT 模块 来转换一种叫做 TorchScript 的表现 如果你好奇 JIT 是什么 它是“Just In Time”的缩写 现在手中有了 TorchScript 的模型 你将唤出新的 Core ML 转换器 产生一个可以拖进你的 app 中的 ML 模型 稍后在本视频中 我将会深入介绍 TorchScript 的转换过程 但是现在让我们着眼于 新的 Core ML 转换器的工作原理

这个转换器是用 Python 编译的 唤出它只需要几行代码 你只需要给它提供一个模型 它可以是一个 TorchScript 对象 或者是保存在磁盘上的路径 它也是对模型输入端的说明 还可以包含有关模型输出的一些信息 但并不是必须的

转换器的工作原理是 迭代 TorchScript 图中的操作 并将它们逐个转换为 其 Core ML 的对应值 有时一个 TorchScript 操作 可能会被转换为 多个 Core ML 操作 有时 图形优化通道 可能会检测到一个已知模式 把几个操作合并成一个

有时 一个模型包含的自定义操作 可能是转换器无法理解的 但也没关系 转换器拥有可扩展的设计 所以为新的操作添加定义并非难事 很多时候 你可以把这个操作表达为现有操作的组合 我们称之为“复合操作” 但如果这还不够 你也可以编写自定义 Swift 实现 在转换过程中以此当做目标 我不会在本视频中 详细说明如何做到这一点 请访问我们的在线资源 获取示例 以及操作步骤

现在我已经概述了整个转换过程 是时候回到开头 深入讲解如何从 PyTorch 模型中获取 TorchScript 模型 PyTorch 有两种方法可以做到这一点 第一种叫做“追踪” 第二种叫做“脚本编制”

我们首先来看看追踪模型是什么意思

追踪是通过唤出 PyTorch 的 JIT 模块的 追踪方法来实现的 如此代码片段所示 我们传递一个 Pythorch 模型 和一个示例输入 它返回模型和 TorchScript 表示

这个调用的真正作用是什么呢？ 主动追踪通过模型的前向传递 运行示例输入 并捕获输入通过模型层级时 唤出的操作 所有这些操作的组合 会变成模型的 TorchScript 表现 现在当你选择一个示例输入追踪时 所使用的的数据 最好与日常使用中模型会看到的相似 例如 你可以使用一个验证数据样本 或者以你的 app 呈现给模型的那样 捕捉数据 你也可以使用随机数据 如果是这样 确保输入值的范围 以及张量的形状与模型期待的一致

让我们通过一个例子 将这一切更加具体化 我要介绍一下我的同事 Paul 他将带领你完成 从 Pythorch 到 Core ML 的 完整分割模型转换 谢谢你 Steve 假如我有一个分割模型 打算在设备上运行 如果你不了解分割模型的作用 我来说明一下 它会获取图像 并为图像的每个像素 分配一个类概率分数 我的模型如何在设备端运行呢？ 我将把我的模型转换为 Core ML 模型 为此 我首先追踪我的 PyTorch 模型 使用 PyTorch 的 JIT 追踪模块 将它转换为 TorchScript 形式

然后 使用新的 Core ML 转换器 将 TorchScript 模型 转换为 Core ML 模型 最后 我将展示最终的 Core ML 模型 如何无缝集成到 Xcode 中 我们看一下这个过程在代码中是怎样的

在此 Jupyter Notebook 中 我将把幻灯片中提及到的 PyTorch 分割模型 转换为 Core ML 模型 如果你想自己试试这段代码 可以在与此视频相关的代码片段中找到 首先 我导入一些 将用于此演示版本的依赖项

接下来 我从 torchvision 和示例输入中 载入 ResNet-101 分割模型： 在本例中是一张狗和猫的图片

PyTorch 模型采用张量对象 而不是 PIL Image 对象 所以我用 transforms.ToTensor 将图像转换为张量 该模型还需要张量有一个额外的维度 来表示批处理大小 因此我也将其添加了进来

如幻灯片中所述 Core ML 转换器 与 TorchScript 模型接洽良好 为此 我使用了 Torch.JIT 模块的 追踪方法 可将 PyTorch 模型 转换为 TorchScript 模型

不好 追踪方法出现了例外 正如例外说明中所说的 “只能从追踪函数中 输出张量或张量的元组”

这是 PyTorch JIT 模块的局限性 这里的问题是我的模型正在返回代码字典 我通过将模型包装在 PyTorch 模块中 来解决此问题 仅从输出代码字典中提取张量值

在这里说明一下我的包装类 是从 PyTorch 的模块类继承的

我定义了包含 ResNet-101 的模型属性 正如上面所使用的那样

在此包装类的前进方法中 我用名为 “out” 的密钥 为返回代码字典建立索引 并只返回张量输出

现在模型返回张量而不是代码字典 就会追踪成功

现在是我使用 新的 Core ML 转换器的时候了 首先 我需要定义我的输入及其预处理

我将输入定义为 ImageType 并通过预处理以使用 ImageNet 统计信息 对图像进行归一化 并将其值缩小到 0 到 1 之间 ResNet-101 所需要的就是这种预处理

接下来 我只需调用 Core ML Tools 的转换方法 传入 TorchScript 模型和输入定义

在转换后 我将设置模型的元数据 以便其他程序 例如 Xcode 可以理解它 我将模型的类型设置为分割 并按模型的秩序枚举类

那么 我转换后的模型成功了吗？ 我可以很容易通过 Xcode 将模型可视化 首先 我将保存模型

现在我需要做的就是 在访达中点击我保存的模型 模型将由 Xcode 打开

在这里我可以查看其元数据 包括输入形状和类型

为了可视化模型的输出 我将转到预览标签 并拖动我的狗和猫的示例图像

看来我的模型已成功分割了 这张图片中的宠物

ResNet-101 能够被追踪 但是有些模型不只是被追踪 为了说明其他模型如何进行转换 我将交回给 Steve 说明

谢谢 Paul 好的 我想我们已经很好地掌握了 如何使用追踪来进行转换 但是 Pythorch 提供了另一种 获取 TorchScript 的方法 所以现在让我们来详细探讨这个方法 它叫做“脚本编制” 脚本编制通过采用 PyTorch 模型 直接编译为 TorchScript 操作来运行 请记住 追踪是在数据流经模型时捕捉的 但和追踪一样 脚本编制模型也非常简单 只需唤出 PyTorch 的 JIT 模块的 脚本编制方法 并为其提供模型

好了 我已经向你展示了 两种获取 TorchScript 表示形式的方法 你可能想知道这两种方法各用于什么情况 如果模型包含控制流 则必须使用脚本编制 让我们看一个例子来了解原因 在这里这个模型有分支和循环 脚本编制将捕捉所有的分支和循环 因为它是直接编译模型的 如果我们追踪模型 我们得到的只是 给定输入通过模型的路径 你可以看到这并不能捕捉整个模型

如果确实需要编写模型脚本 那么如果尽可能多地追踪模型 并且只编写模型中需要它的部分 通常会得到最好的结果 这是因为追踪通常会比脚本编制 生成更简单的表现 让我们看看如何通过 查看一些代码来应用这个理念 在本示例中 我有一个 在一个循环中以固定的次数 运行一些代码块的模型 我已经将循环的主体 分离成一个可以很容易被追踪的对象 然后 我可以将脚本编制 作为一个整体应用到模型中 我们所做的基本上就是将脚本编制 限制在需要其控制流的数位 然后追踪其他所有内容 这种追踪和脚本编制相结合的做法 之所以奏效 是因为它们都会跳过 已经转换为 TorchScript 的代码

现在我们来看一个 使用脚本编制的具体示例 我将交回给 Paul 他将带你展示如何转换语言模型 你好 假设我有一个语句完成模型 我想将其转换为 Core ML 模型 以便它可以在设备上运行 在某些情况下 语句完成是一项任务 包括获取语句片段 并使用模型来预测接下来可能出现的单词 那么从计算步骤来看 这是什么样子的呢？ 我将从语句片段中的几个单词开始 然后将它们传递到所谓的编码器中 可以将这些单词转换为 模型可以理解的表示形式 在本例中 是一个整数标记序列 接下来 我将把这个标记序列传递给模型 它将预测序列中的下一个标记 我将继续为模型提供部分构造的句子 并在结尾添加新的标记 直到模型预测到一个特殊的句末标记 这意味着我的句子已经完成

现在我有了一个完整的标记语句 我将其通过解码器传递 解码器会将标记转换回单词 这个示意图的中间部分 所完成的标记列表 是我将要转换成的 Core ML 模型 编码器和解码器分开处理 让我们通过查看一些伪代码 来确保我们理解了现在的情况 模型的核心是下一个标记预测器 为此 我将使用 Hugging Face 的 GPT2 模型 预测器将标记列表作为输入 并为下一个标记提供预测 接下来 我将围绕预测器包装一些控制流 以继续操作 直到看到句末标记

在此循环内 我将预测的标记附加到运行列表中 并将其作为每次循环中预测器的输入 当预测器返回句末标记时 我将返回完整的句子进行解码 现在来看看整个过程的编码 让我们深入到 Jupyter Notebook 在这个笔记本中 我将构建一个语言模型 该模型采用语句片段并完成语句 先把输入部分完成

这是我模型的代码

我的模型继承自 torch.Module 并包含句末标记 next_token_predictor 模型 和表示句子开头的默认标记的属性

与幻灯片中一样 在前进方法中 我编写了一个可获取标记列表 并预测下一个标记列表的循环体 循环将继续 直到生成句末标记 这时 我们将返回该句子

如前所述 我的下一个标记预测器将是 GPT2 它将驻留在循环体中

我将按照把追踪循环 与对整个模型 进行脚本编制分开的做法来进行 因此 我将只在下一个标记预测器上 运行 JIT 追踪程序 它需要一个标记列表作为输入 因此为了进行追踪 我只传递一个随机标记列表

我可以看到追踪程序发出警告 告诉我这个追踪可能不会泛化到其他输入 请注意 此警告来自 PyTorch 的 JIT 追踪程序 而不是 Core ML 稍后在“疑难解答”部分将对此进行解释 但由于实际上没有问题 因此我现在将忽略此警告

通过追踪循环体的大部分内容 我可以实例化我的语句完成模型 并应用 JIT 脚本编写器 来准备将其转换为 Core ML

正如分割演示中所示 我现在将 TorchScript 模型 转换为 Core ML 模型

现在我看看我的模型能不能完成一个句子 我创建一个句子片段： 本例是 “曼哈顿桥是” 然后我通过 GPT2 内含的编码器 来获取片段的编码 然后将该标记列表转换为 Torch 张量

接下来 从 Core ML 模型中将输入打包 运行所述模型 并用 GPT2 内含的解码器解码输出

很好 Core ML 模型能够完成句子 看起来它生成了一个 关于曼哈顿大桥的叙述

在追踪和编制模型脚本 以使其转换为 Core ML 格式时 你可能会遇到一些障碍 我将交回给 Steve 为你提供帮助

在我们进行总结之前 我想回顾一下我们在将 PyTorch 模型 转换为 Core ML 时遇到的障碍 并介绍一些故障排除技巧和最佳做法

回想一下分割演示中的情形 还记得我们在追踪过程中 遇到了一个错误吗？ 这是因为我们的模型返回了代码字典 并且 JIT 追踪只能处理张量或张量元组

我们在演示中展示的解决方案 是在模型周围创建一个 解压模型本机输出的薄包装器 请记住 在这个示例中 模型返回代码字典 因此这里我们访问 表示推断结果的代码字典密钥 并返回该张量 当然 这种方法也适用于 如果我们想从代码字典中 访问并返回多个条目 或者需要解包其他类型的容器

然后在语言模型演示过程中 我们遇到了一个追踪程序警告 说追踪可能不会泛化到其他输入 而且我们看到追踪程序 笼统地列印出有问题的代码行

那到底是怎么回事？ 如果我们查看模型源代码来理解此警告 我们会看到模型正在基于另一张量的大小 对一个张量进行切片 获取张量的大小结果是 会导致一个纯 Python 值 换句话说 不是 PyTorch 张量 并且追踪程序警告其无法追踪 对这些纯 Python 值执行的数学运算 但是 在本例中 追踪程序发出的警告有些严重 实际上并没有问题

在追踪对纯 Python 值 进行操作的代码时 一个很好的经验法则是 追踪程序只会正确捕捉 内置的 Python 操作

下面有一些例子来帮助阐释这个法则 让我们仔细思考一下 然后根据这个经验法则 看看它们是否会被正确地追踪

第一个示例与我们在演示过程中看到的 非常相似 并且由于应用了内置操作 在本例中是添加操作 因此将实现正确的追踪

第二个示例也将实现正确追踪 在本例中使用模运算符 这也是一个内置操作

但是第三个示例无法正确追踪 JIT 追踪程序 不知道 math.sqrt 函数资料库的作用 于是追踪层将记录一个常量值 而不是进行计算张量大小和平方根的操作

但是 通过对模型进行简单的修复 用 Python 的内置幂运算符 替换 math.sqrt 就可以实现正确的追踪 现在让我们看一下 编制模型脚本会失败的情况 该模型以空列表开始 并向其连续添加一组固定的整数 请记住 这不是一个非常有用的模型 我只是用它来说明故障情况 如果我对此模型编制脚本 则会收到运行错误提示 提示类型不匹配 JIT 脚本编制器需要类型信息 来将模型转换为 TorchScript 并且可以很好地从上下文推断对象类型 但有时这是不可能的 如果脚本编制器无法确定对象的类型 它将假定该对象是张量 在本例中 它假设这个列表是张量列表 而实际上它是作为整数列表构建的 那么我该怎么做以帮助脚本编制器呢？ 我既可以计入变量有意义的初始值 也可以使用类型注释 在这里 我对模型进行了调整 以展示两者的示例

还有最后一件事 在追踪之前 一定要确保模型处于评估模式 这确保所有层级都配置为推断 而不是训练 对于大多数层级 这并不重要 但例如 如果你的模型中有一个退出层 设置评估模式将确保它被禁用 当转换器遇到被禁用的操作时 将会把它们视为传递操作

我们在这段视频中已经介绍了很多内容 但是你可以在与视频相关的链接中 找到更多信息 包括 Core ML 转换器的说明文档 关于定制操作转换的信息 以及许多详细的 TorchScript 示例

我们非常高兴能为 Pythorch 模型的转换 提供一流的支持 希望你会发现新的 Core ML 转换器 将为你的 PyTorch 模型 提供更广泛的支持 使你能够优化设备上的模型执行 并真正为你提供最大限度的支持 使你的模型转换更加容易 感谢观看