目录
  • 文档 >
  • 模型导出和降低
快捷方式

模型导出与降低

本节介绍将 PyTorch 模型转换为 ExecuTorch 使用的运行时格式的过程。这个过程通常被称为“导出”，因为它利用 PyTorch 的导出功能将 PyTorch 模型转换为适合设备上执行的格式。这个过程会生成一个 .pte 文件，该文件针对使用特定后端的设备上执行进行了优化。如果使用程序-数据分离，还会生成相应的 .ptd 文件，其中仅包含模型中的权重/常量。

先决条件

导出需要安装 ExecuTorch Python 库，通常通过运行 pip install executorch 来安装。有关更多信息，请参阅 安装。这个过程假定你有一个 PyTorch 模型，可以从 Python 中实例化它，并且可以提供示例输入张量来运行模型。

导出与降低过程

将模型导出并降低到 .pte 格式的过程通常涉及以下步骤：

  1. 选择目标后端。

  2. 准备 PyTorch 模型，包括输入和形状规范。

  3. 使用 torch.export.export 导出模型。

  4. 使用 to_edge_transform_and_lower 为目标后端优化模型。

  5. 通过调用 to_executorch 并序列化输出来创建 .pte 文件。


量化——即使用降低精度来减少推理时间和内存占用——通常也在此阶段完成。有关更多信息，请参阅 量化概述。

硬件后端

ExecuTorch 后端为特定的硬件目标提供硬件加速。为了在目标硬件上实现最大性能，ExecuTorch 在导出和降低过程中会针对特定后端优化模型。这意味着生成的 .pte 文件是针对特定硬件专门化的。为了部署到多个后端，例如 iOS 上的 Core ML 和 Android 上的 Arm CPU，通常会为每个后端生成一个专用的 .pte 文件。

硬件后端的选择取决于模型的部署目标硬件。每个后端都有特定的硬件要求和模型支持级别。有关更多详细信息，请参阅每个硬件后端的文档。

作为 .pte 文件创建过程的一部分，ExecuTorch 会识别模型中支持给定后端的（分区）部分。这些部分会由后端提前处理以支持高效执行。如果模型中存在不支持委托的部分，则会使用 CPU 上的可移植回退实现来执行。这可以在并非所有模型算子都得到后端支持的情况下实现部分模型加速，但可能会对性能产生负面影响。此外，还可以指定多个分区器并设置优先级。例如，这允许 GPU 上不支持的算子通过 XNNPACK 在 CPU 上运行。

可用后端

下面列出了常用的硬件后端。对于移动端，可以考虑使用 XNNPACK 处理 Android，使用 XNNPACK 或 Core ML 处理 iOS。要为特定后端创建 .pte 文件，请将相应的分区器类传递给 to_edge_transform_and_lower。有关更多信息，请参阅相应的后端文档和下面的 导出与降低 部分。

模型准备

导出过程接受标准的 PyTorch 模型，通常是 torch.nn.Module。这可以是自定义模型定义，也可以是来自现有源的模型，例如 TorchVision 或 HuggingFace。有关降低 TorchVision 模型的示例，请参阅 ExecuTorch 入门。

模型导出是从 Python 完成的。这通常通过 Python 脚本或交互式 Python Notebook（如 Jupyter 或 Colab）来完成。下面的示例展示了一个简单的 PyTorch 模型的实例化和输入。输入被准备为 torch.Tensors 的元组，模型可以使用这些输入运行。

import torch

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.seq = torch.nn.Sequential(
            torch.nn.Conv2d(1, 8, 3),
            torch.nn.ReLU(),
            torch.nn.Conv2d(8, 16, 3),
            torch.nn.ReLU(),
            torch.nn.AdaptiveAvgPool2d((1,1))
       )
        self.linear = torch.nn.Linear(16, 10)

    def forward(self, x):
        y = self.seq(x)
        y = torch.flatten(y, 1)
        y = self.linear(y)
        return y

model = Model().eval()
inputs = (torch.randn(1,1,16,16),)
outputs = model(*inputs)
print(f"Model output: {outputs}")

请注意，模型使用 .eval() 设置为评估模式。模型应始终在评估模式下导出，除非进行设备上训练。此模式会将某些具有训练特定行为的操作（如批归一化或 dropout）配置为使用推理模式配置。

导出与降低

要实际导出和降低模型，请依次调用 export、to_edge_transform_and_lower 和 to_executorch。这将生成一个可以序列化到文件的 ExecuTorch 程序。将所有内容整合起来，使用 XNNPACK 委托来降低上面示例模型以获得移动 CPU 性能，可以如下完成：

import torch
from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
from executorch.exir import to_edge_transform_and_lower
from torch.export import Dim, export

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.seq = torch.nn.Sequential(
            torch.nn.Conv2d(1, 8, 3),
            torch.nn.ReLU(),
            torch.nn.Conv2d(8, 16, 3),
            torch.nn.ReLU(),
            torch.nn.AdaptiveAvgPool2d((1,1))
        )
        self.linear = torch.nn.Linear(16, 10)

    def forward(self, x):
        y = self.seq(x)
        y = torch.flatten(y, 1)
        y = self.linear(y)
        return y

model = Model()
inputs = (torch.randn(1,1,16,16),)
dynamic_shapes = {
    "x": {
        2: Dim("h", min=16, max=1024),
        3: Dim("w", min=16, max=1024),
    }
}

exported_program = export(model, inputs, dynamic_shapes=dynamic_shapes)
executorch_program = to_edge_transform_and_lower(
    exported_program,
    partitioner = [XnnpackPartitioner()]
).to_executorch()

with open("model.pte", "wb") as file:
    file.write(executorch_program.buffer)

这将生成一个 model.pte 文件，可以在移动设备上运行。

要生成 model.pte、model.ptd 对，并将权重放入 model.ptd 中，请将以下转换函数添加到将常量标记为外部：

from executorch.exir.passes.external_constants_pass import (
    delegate_external_constants_pass,
)
partial_function = partial(
    delegate_external_constants_pass,
    ep=exported_program,
    gen_tag_fn=lambda x: "model", # This is the filename the weights will be saved to. In this case, weights will be saved as "model.ptd"
)

executorch_program = to_edge_transform_and_lower(
    exported_program,
    transform_passes = [partial_function],
    partitioner = [XnnpackPartitioner()]
).to_executorch()

保存 PTD 文件

executorch_program.write_tensor_data_to_file(output_directory)

它将保存到文件 model.ptd，文件名来自转换过程中的 gen_tag_fn。

支持可变输入尺寸（动态形状）

PyTorch 导出过程使用提供的示例输入来跟踪模型并了解每个步骤中张量的大小和类型。除非另有说明，否则导出将假定固定输入大小等于示例输入，并将使用此信息来优化模型。

许多模型需要支持可变输入尺寸。要支持这一点，导出接受一个 dynamic_shapes 参数，它会告知编译器哪些维度可以变化及其边界。这采用嵌套字典的形式，其中键对应于输入名称，值指定每个输入的边界。

在示例模型中，输入是 4 维张量，遵循批次、通道、高度和宽度 (NCHW) 的标准约定。形状为 [1, 3, 16, 16] 的输入表示 1 个批次，3 个通道，高度和宽度为 16。

假设您的模型支持 16x16 到 1024x1024 之间的图像尺寸。形状边界可以指定如下：

dynamic_shapes = {
    "x": {
        2: Dim("h", min=16, max=1024),
        3: Dim("w", min=16, max=1024),
    }
}

ep = torch.export.export(model, inputs, dynamic_shapes=dynamic_shapes)

在上面的示例中，"x" 对应于 Model.forward 中的参数名称。2 和 3 键对应于维度 2 和 3，即高度和宽度。由于没有指定批次和通道维度，因此这些值根据示例输入固定。

ExecuTorch 使用形状边界来优化模型和规划模型执行的内存。因此，建议将维度上限设置为不超过需要，因为更高的上限会增加内存消耗。

对于更复杂的用例，动态形状规范允许维度之间的数学关系。有关动态形状规范的更多信息，请参阅 表达动态性。

测试模型

在集成运行时代码之前，通常会从 Python 中测试导出的模型。这可用于评估模型准确性并在移动到目标设备之前进行健全性检查。请注意，并非所有硬件后端都可以从 Python 使用，因为它们可能需要专用硬件才能运行。有关硬件要求和模拟器可用性的更多信息，请参阅特定的后端文档。本示例中使用的 XNNPACK 委托在主机上始终可用。

from executorch.runtime import Runtime

runtime = Runtime.get()

input_tensor = torch.randn(1, 3, 32, 32)
program = runtime.load_program("model.pte")
method = program.load_method("forward")
outputs = method.execute([input_tensor])

Pybindings 目前不支持加载程序和数据。要运行包含 PTE 和 PTD 组件的模型，请使用 扩展模块。在 executorch-examples 中还有一个端到端演示。

有关更多信息，请参阅 运行时 API 参考。

高级主题

虽然许多模型可以按照上述步骤“正常工作”，但一些更复杂的模型可能需要额外的导出工作。这些包括带有状态的模型以及具有复杂控制流或自回归生成模型的。有关这些技术的示例用法，请参阅 Llama 模型。

状态管理

某些类型的模型会维护内部状态，例如 Transformer 中的 KV 缓存。有两种方法可以在 ExecuTorch 中管理状态。第一种是将状态作为模型输入和输出带出，有效地使核心模型无状态。这有时被称为将状态作为 IO 管理。

第二种方法是直接利用模型内的可变缓冲区。可变缓冲区可以使用 PyTorch 的 register_buffer API 在 nn.Module 上注册。缓冲区的存储由框架管理，并且在方法执行结束时，对缓冲区的任何修改都会被写回。

可变缓冲区有几个限制：

  • 可变性的导出可能很脆弱。

    • 如果您遇到与梯度相关的导出时错误，请考虑在将张量分配给缓冲区之前显式调用 detach()。

    • 确保对可变缓冲区进行的任何操作都使用就地操作（通常以 _ 结尾）。

    • 不要重新分配缓冲区变量。而是使用 copy_ 来更新整个缓冲区内容。

  • 可变缓冲区不会在 .pte 内的多个方法之间共享。

  • 就地操作会被替换为非就地变体，并且在方法执行结束时将结果张量写回。当使用 index_put_ 时，这可能会成为性能瓶颈。

  • 并非所有后端都支持缓冲区突变，这可能会导致图中断和内存传输回 CPU。

突变支持是实验性的，未来可能会发生变化。

动态控制流

如果控制流的路径在导出时不是固定的，则认为它是动态的。当 if 或 loop 条件取决于张量的值时，通常是这种情况，例如，当生成结束序列标记时终止的生成器循环。形状相关的控制流也可能是动态的，如果张量形状取决于输入。

要使动态 if 语句可导出，可以使用 torch.cond 来编写。目前 ExecuTorch 不支持动态循环。启用此类模型的一般方法是将循环体导出为方法，然后从应用程序代码处理循环逻辑。这对于处理自回归模型中的生成器循环（如 Transformer 增量解码）很常见。

多方法模型

ExecuTorch 允许将多个方法打包到一个 .pte 文件中。这对于更复杂的模型架构（如编码器-解码器模型）可能很有用。

要在 .pte 中包含多个方法，必须使用 torch.export.export 单独导出每个方法，从而为每个方法生成一个 ExportedProgram。这些可以作为字典传递给 to_edge_transform_and_lower。

encode_ep = torch.export.export(...)
decode_ep = torch.export.export(...)
lowered = to_edge_transform_and_lower({
    "encode": encode_ep,
    "decode": decode_ep,
}).to_executorch()

在运行时，可以通过方法名称传递给 Module 类的 load_method 和 execute。

多方法 .pte 有几个注意事项：

  • 方法是单独进行内存规划的。目前，激活内存不会在方法之间重用。对于高级用例，可以使用 自定义内存计划 或 自定义内存分配器 来重叠分配。

  • 方法之间不会共享可变缓冲区。

  • PyTorch 导出目前不允许导出模块上 forward 之外的其他方法。为了解决这个问题，通常会为每个方法创建包装器 nn.Modules。


class EncodeWrapper(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model

    def forward(self, *args, **kwargs):
        return self.model.encode(*args, **kwargs)

class DecodeWrapper(torch.nn.Module):
    # ...

encode_ep = torch.export.export(EncodeWrapper(model), ...)
decode_ep = torch.export.export(DecodeWrapper(model), ...)
# ...

下一步

PyTorch 和 ExecuTorch 的导出与降低 API 提供了高度的可定制性，以满足不同硬件和模型的需求。有关更多信息，请参阅 torch.export 和 导出 API 参考。

有关高级用例，请参阅以下内容：

文档

访问全面的 PyTorch 开发者文档

查看文档

教程

为初学者和高级开发者提供深入的教程

查看教程

资源

查找开发资源并让您的问题得到解答

查看资源

为了分析流量并优化您的体验，我们在此网站上使用 Cookie。点击或浏览即表示您同意我们使用 Cookie。作为本网站的当前维护者，Facebook 的 Cookie 政策适用。了解更多信息，包括可用的控件： Cookie 政策。