模型导出和降低¶

本节介绍将 PyTorch 模型转换为 ExecuTorch 使用的运行时格式的过程。此过程通常称为“导出”，因为它使用 PyTorch 的导出功能将 PyTorch 模型转换为适合在设备上执行的格式。此过程会生成一个 .pte 文件，该文件已针对特定后端进行了设备上执行优化。

先决条件¶

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

导出和降低过程¶

将模型导出并降低为 .pte 格式的过程通常包括以下步骤

选择要定位的后端。
准备 PyTorch 模型，包括输入和形状规范。
使用 torch.export.export 导出模型。
使用 to_edge_transform_and_lower 为目标后端优化模型。
通过调用 to_executorch 并序列化输出来创建 .pte 文件。

量化 - 使用较低精度来减少推理时间和内存占用的过程 - 通常也在这一阶段进行。有关更多信息，请参阅量化概览。

硬件后端¶

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

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

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

可用后端¶

下面列出了常用的硬件后端。对于移动设备，请考虑在 Android 上使用 XNNPACK，在 iOS 上使用 XNNPACK 或 Core ML。要为特定后端创建 .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() 设置为评估模式。除非执行设备上训练，否则模型应始终在评估模式下导出。此模式将某些具有训练特定行为的操作（例如 batch norm 或 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 文件，该文件可以在移动设备上运行。

支持可变输入大小（动态形状）¶

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

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

在示例模型中，输入按照批次、通道、高度和宽度（NCHW）的标准惯例以 4 维张量形式提供。形状为 [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])

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

后续步骤¶

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

对于高级用例，请参阅以下内容

有关量化模型以减少推理时间和内存占用的信息，请参阅量化概览。
有关控制内存放置和规划的信息，请参阅内存规划。
有关编写自定义编译器 passes 的信息，请参阅自定义编译器 Passes。
有关导出生成的中间表示的信息，请参阅导出 IR 规范。