PyTorch 深度学习——部署到生产环境本章涵盖以下内容：部署 PyTorch 模型的多种选项通过 Web 框架和

本章涵盖以下内容：

部署 PyTorch 模型的多种选项
通过 Web 框架和 API 部署模型
优化推理性能
将模型导出到不同部署目标所需的格式
从 C++ 中运行导出的模型以及原生实现的模型

在本书第一部分中，我们学习了大量关于模型的内容；而第二部分则为某一特定问题如何构建出优秀模型留下了一条非常细致的路径。既然现在我们已经拥有了这些很棒的模型，下一步就需要把它们带到真正能发挥作用的地方去。为大规模执行深度学习模型推理而维护基础设施，从架构设计和成本控制两个角度看，都可能是非常划算的。虽然 PyTorch 最初是一个以研究为导向的框架，但它后来已经经历了显著演进，加入了大量面向生产环境的特性，使其成为一个既适用于研究、又适用于大规模生产的理想端到端平台。

至于“部署到生产环境”到底意味着什么，会因具体用例而异：

对于我们的演示来说，我们会先用 Gradio（www.gradio.app/）做原型。Gradio 是一个框架，专门用于让用户能够快速为自己的机器学习模型搭建演示程序或 Web 应用。Gradio 提供了一个简单界面，便于用户与模型交互并分享给他人，因此无需前端开发技能，也能轻松测试和展示模型能力。
将我们开发的模型部署出去，最直观的方式之一，就是搭建一个网络服务来提供这些模型的访问能力。这可以借助任意一种流行的 Python Web 框架来实现，比如 FastAPI（fastapi.tiangolo.com/）。FastAPI 因其易用性和高性能而广受青睐，它还能自动生成交互式 API 文档，并对数据校验与序列化提供强大支持，因此非常适合高效构建可靠且可扩展的 API。
我们也可以把模型导出为一种标准化程度很高的格式，这样就能够借助经过优化的模型处理器、专用硬件或云服务来部署它。对于 PyTorch 模型而言，这个角色通常由 ONNX（Open Neural Network Exchange） 格式来承担。
我们可能还希望把模型集成进更大型的应用，或者集成到那些不使用 Python 的应用中。为此，如果我们的能力不被限制在 Python 内部，就会非常方便。因此，我们还会探索如何在 C++ 中使用 PyTorch 模型。
最后，在某些任务中，让模型跑在移动设备上会很有吸引力。用户可能更愿意让模型在手机本地运行，而不是把自己的私密数据（例如图像）发送到云端服务。幸运的是，PyTorch 近来已经构建了支持移动端的工具，我们也会对这部分进行讨论。

注意：Streamlit（streamlit.io/）与 Gradio 很相似，也非常值得一提。

在整章中，我们会结合自定义模型和流行的预训练模型来实现一些实际示例，展示那些在性能、易用性和资源效率之间取得平衡的部署策略。

17.1 为 PyTorch 模型提供服务

好，现在我们手里已经有了一个非常酷的模型，并且希望尽快把它分享给同事或朋友。我们会先从最简单、最容易原型化和公开分享模型的方式之一开始。按照我们一贯的动手风格，我们会先从一个能工作的最简单服务器开始。等到这个最基础版本跑通之后，再去审视它的局限性，并逐步解决这些问题。现在就开始搭建一个在网络上监听请求的服务器吧。（为了安全起见，请不要在不可信网络中这么做。）

17.1.1 用 Gradio 对外提供我们的模型

Gradio 是一个 Python 库，能够让我们很快为机器学习模型创建一个基于 Web 的交互界面。它让模型的展示与测试变得非常容易，而且不需要掌握大量 Web 开发知识。Gradio 对于制作交互式 demo 并分享给别人尤其有用。可以通过 pip 安装：

pip install --upgrade gradio

我们只需要定义一个函数，指定它接收什么输入并返回什么输出，Gradio 就能自动帮我们生成对应的 UI 组件。如果函数签名表明它接收一个字符串并返回一个字符串，那么 Gradio 就会自动创建一个文本输入框和一个文本输出框。例如下面这个例子（gradio_hello_world.py）：

代码清单 17.1 创建一个带文本输入和输出的简单 Gradio 界面

import gradio as gr

def hello_world(name):
    return "Hello, " + name + "!"

demo = gr.Interface(
    fn=hello_world,
    inputs=["text"],
    outputs=["text"],
    flagging_mode="auto",
)

demo.launch()

你的应用可以通过如下方式启动：

gradio gradio_hello_world.py

这样会自动监测这个 Python 文件的变化，并在修改后自动重载服务器。应用会运行在 7860 端口，并可通过浏览器访问 http://localhost:7860。

如果想让它更有意思一些，我们还可以接入一个预训练模型来生成图像。这里我们会用到 diffusers 库，它为扩散模型提供了非常简单的使用接口。我们将使用第 10 章中见过的某个 Stable Diffusion 模型版本，根据文本提示生成图像。为此，我们需要把模型封装进一个函数：该函数接收文本 prompt 作为输入，并返回生成出来的图像（gradio_server.py）。

代码清单 17.2 加载一个 Stable Diffusion 模型用于图像生成

image_generator = None
def load_image_generator():
    global image_generator
    if image_generator is None:
        model_id = "stable-diffusion-v1-5/stable-diffusion-v1-5"
        pipeline = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16)
        device = "cuda" if torch.cuda.is_available() else "cpu"
        pipeline.to(device)

        pipeline.enable_vae_tiling()
        pipeline.enable_sequential_cpu_offload()

        image_generator = pipeline
    return image_generator

def generate_image(prompt):
    image_generator = load_image_generator()
    image = image_generator(prompt).images[0]
    return image

demo = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(label="Prompt"),
    outputs=gr.Image(type="pil", label="Generated Image"),
    flagging_mode="auto",
)

demo.launch()

启动服务器的方式和之前类似：执行 gradio gradio_server.py 即可。我们甚至还可以把 demo.launch() 改成 demo.launch(share=True)，这样就会生成一个公共链接，任何人都可以通过这个链接访问应用。这是一个非常适合快速分享模型的办法，而且不需要额外搭建任何基础设施。图 17.1 展示了一个示例。

图 17.1 使用 Gradio 实现的图像生成示例

不过，虽然它非常适合快速原型验证和模型分享，但它并不适合真正的生产环境。Gradio 对于 demo 和测试非常出色，但它并不是为高性能或高可扩展性场景设计的。接下来我们就来学习如何搭建一个更健壮的服务器，它能够处理多个请求，并提供更好的性能。

17.1.2 让模型跑在 FastAPI 服务器后面

FastAPI 是一个非常流行的 Python 框架，专门用于构建服务端 API。它可以通过 pip 安装：

pip install "fastapi[standard]"

提示：你可能会希望在 Python 虚拟环境中运行 FastAPI。

创建 API 的方式，是根据你希望这个端点是 GET 还是 POST，在函数上加上 @app.get 或 @app.post 装饰器。例如，如果我们想提供一个 GPT 风格的文本生成模型，就可以这样写（fastapi_hello_world.py）：

from fastapi import FastAPI

app = FastAPI()

@app.get("/")
async def read_root():
    return "Hello world"

服务器可以通过如下命令启动：

fastapi dev fastapi_hello_world.py

启动后，应用会运行在 http://127.0.0.1:8000，并暴露出一个 / 路由，该路由返回 "Hello World" 字符串。这个 URL 是一个本地地址（http），用于让我们的电脑访问运行在本机上的 API；127.0.0.1 就是 localhost 的 IP 地址，而端口号 8000 则表示我们想连接到本机上的哪一个服务。

FastAPI 还会自动提供一个非常有帮助的交互式文档页面，位于 http://127.0.0.1:8000/docs，我们可以在这里手工测试 API。另外，因为我们是以开发模式启动服务器的，所以只要修改代码，FastAPI 就会自动重载服务。

到这里，我们就可以开始强化这个 Web 服务器：加载一个已经保存好的模型，并通过一个 POST 路由把它暴露出去。这个示例里，我们会使用 Hugging Face transformers 库中的一个 text-to-text 生成模型：

pip install -U transformers

我们还会利用 FastAPI 内建的能力来定义输入的数据结构。具体来说，我们会定义一个 Pydantic 模型，用来表示用户提交的输入数据（fastapi_server.py）。

注意：Pydantic 是一个 Python 库，用于在运行时对数据模型进行类型检查和校验（docs.pydantic.dev/latest/）。](docs.pydantic.dev/latest/%EF%…)

代码清单 17.3 定义一个 FastAPI 应用

from fastapi import FastAPI
from pydantic import BaseModel
# Define a Pydantic model for the input data
class TextInput(BaseModel):
    text: str

app = FastAPI()

接下来，我们不再使用 /hello 路由，而是定义一个 /generate 路由。这个路由会通过 POST 请求接收一个字符串输入（句子的开头），服务器对它进行处理，然后返回一个 JSON 响应，其中包含句子的预测续写结果。

为了获取输入数据，FastAPI 会自动把传入的 JSON 请求体解析成一个 Pydantic 模型实例，只要我们事先定义好它期望的结构即可。模型处理流程与在 Gradio 中的思路类似：从库中加载模型架构及其预训练权重。如果有 GPU，就可以通过对模型调用 .to('cuda') 把它迁移到 GPU 上，以获得更快的处理速度。另外，在推理时，我们会先调用 model.eval() 把模型切到评估模式，并在 torch.no_grad() 上下文中执行推理代码，以关闭梯度计算，从而提升推理效率并减少内存占用（fastapi_server.py）。

代码清单 17.4 使用 SmolLM2 提供文本生成服务

...
from transformers import AutoModelForCausalLM, AutoTokenizer

model = None
tokenizer = None
def get_model_and_tokenizer(device):
    global model, tokenizer
    checkpoint = "HuggingFaceTB/SmolLM2-360M-Instruct"    #1
    if model is None:
        model = AutoModelForCausalLM.from_pretrained(checkpoint)   #2
        model = model.to(device)
        model.eval()
    if tokenizer is None:
        tokenizer = AutoTokenizer.from_pretrained(checkpoint)
    print(f"Memory footprint: {model.get_memory_footprint() / 1e6:.2f} MB")
    return model, tokenizer

@app.post("/generate")
def generate_text(input_txt: TextInput = \
 TextInput(text="What is PyTorch and why is it cool?")):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model, tokenizer = get_model_and_tokenizer(device)
    messages = [{"role": "user", "content": input_txt.text}]
    input_text=tokenizer.apply_chat_template(messages, tokenize=False)
    print(f"Input text: {input_text}")
    inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
    with torch.no_grad():        #3
        outputs = model.generate(inputs, max_new_tokens=256, do_sample=True, temperature=0.7)
    generated_text = tokenizer.decode(outputs[0])
    return generated_text        #4

#1 这里使用这个模型是因为它相对较小，但也可以替换成 Hugging Face 支持的其他模型。
#2 初始化模型、加载权重，并切换到评估模式。
#3 虽然 model.generate() 内部会处理这个问题，但在推理时显式关闭 autograd 仍然是一个好习惯。
#4 返回的是完整文本（包括输入部分）作为响应。

运行服务器的方法如下：

fastapi dev fastapi_server.py

作为客户端，我们可以打开 FastAPI 的文档页面 http://127.0.0.1:8000/docs，手工与 API 交互。也可以在终端中使用 curl 发送一个 POST 请求。POST 请求是一种向服务器提交数据的 HTTP 请求方式。这里我们把请求发往 /generate 端点，也就是这个 API 中专门负责文本生成的 URL 路径。JSON payload 则是随请求一并发送的数据，它包含了我们希望 API 处理的输入文本。例如，可以在终端运行下面这条命令：

curl -X 'POST' \
  'http://127.0.0.1:8000/generate' \
  -H 'accept: application/json' \
  -H 'Content-Type: application/json' \
  -d '{
  "text": "What is PyTorch and why is it cool?"
}'

这个 API 应该会返回一段关于 PyTorch 的回答，而且希望它提到的内容，此时你都已经能看懂了。很明显，我们的服务器现在已经可以接收输入，把它们送进模型，再把输出结果返回回来。那是不是就算完成了？还没有。接下来我们来看看，还可以在哪些方面继续改进。

17.1.3 我们希望部署系统具备什么特性

现在来整理一下，我们希望一个模型服务系统具备哪些能力。首先，我们希望它支持现代协议及其特性。老式 HTTP 的一个问题是它在本质上非常串行：如果客户端在同一个连接上想连续发送多个请求，那么后续请求往往得等前一个请求响应完之后才能发出——如果你想同时处理多个请求，这种方式效率并不高。FastAPI 允许你通过 async def 定义异步端点，从而让服务器能够并发处理多个请求。这意味着，当一个请求正在等待某个 I/O 操作（例如数据库查询或外部 API 调用）时，服务器仍然可以继续处理其他请求。

异步编程听起来可能有点吓人，也常常被一些复杂术语包围着。但说到底，它的本质就是：让函数在等待计算结果或等待事件发生时，不会阻塞其他任务继续执行。这类异步函数有时被称为生成器，或者更广义地说，被称为 coroutine（协程） 。

如果要把我们的 FastAPI 端点改造成异步，只需要在函数定义里加上 async：

async def generate_text(...)

使用 FastAPI 的好处就在于，它能帮我们处理很多底层细节，而我们只需要专注于自己的代码逻辑。

不过，我们这个 API 端点的主要瓶颈，其实是模型的推理时间。尤其是在 GPU 上运行时，通常将多个请求攒成一个 batch 一起推理，要比逐个处理或者简单并发触发推理高效得多。这是由 GPU 擅长并行计算的特性决定的。因此，接下来我们就要面对这样一个问题：如何从多个连接里收集请求，把它们拼成一个 batch，送进 GPU 推理，然后再把结果分别返回给对应的请求方。听起来确实有些复杂，而且（至少在我们写这本书的时候）在很多简单教程里并不常见。也正因为如此，我们更有必要在这里把这件事做对。

我们希望并行地服务多件事情。即便使用了异步服务，我们仍然希望模型本身能在另一个线程中高效运行。这意味着，我们希望在模型执行这一块绕开 Python 那个“赫赫有名”的 GIL（global interpreter lock，全局解释器锁） 。

为了降低延迟，我们还希望能够把生成结果逐 token 地流式返回给客户端。这在现代 API 中已经很常见，尤其是文本生成任务。通过流式返回，用户不必等整个响应全部生成完才看到结果，而是能像在聊天应用里那样，边生成边看到内容，从而获得更强的交互感。这对聊天机器人、交互式写作助手之类的应用尤其有价值。

最后还有一点：为 API 提供服务时，安全性也非常关键。我们需要防止溢出、资源耗尽等问题。由于这里的输入文本长度是固定上限的，因此整体上我们处于比较有利的位置，因为在固定输入大小下，想在 PyTorch 中制造严重问题并不那么容易。相比之下，图像解码等任务会复杂得多；而我们这里只处理文本，因此简单不少。至于互联网安全，那是一个非常庞大的领域，超出了本书范围。顺便一提，神经网络本身也可能受到输入操纵的影响，从而产生错误或意外输出，这类现象被称为 对抗样本（adversarial examples） 。这些内容同样不在本书讨论范围内，所以这里先放下不谈。

说了这么多，下面就来真的改进我们的服务器。

17.1.4 请求批处理与流式响应

要实现请求批处理，我们必须把“处理请求”和“运行模型”这两件事解耦。图 17.2 展示了整个数据流。

图 17.2 带有请求批处理和流式响应的数据流

在图 17.2 顶部，客户端发起请求，这些请求会把待处理项放入一个队列中。一旦队列里的请求数量达到一个完整 batch，或者最早进入队列的那个请求已经等待超过了预设的最大时间，一个模型工作线程（model worker）就会把这批请求取出，并一起处理。

这个模型 worker 运行在后台线程中，因此 API 本身仍然能够在模型执行期间继续并发接收新请求。这种设计需要额外逻辑来管理队列的生命周期以及 worker 的生命周期。

此外，API 还会为每个 work item 提供流式响应：一旦生成出了新 token，就立即把它发回对应客户端。这是通过一个异步生成器来完成的，它会一个接一个地 yield 出这些 token。

你可以把它类比成一家非常忙碌的餐厅：顾客下单后，服务员会把订单写在单子上，放进后厨队列。等订单积攒够一批，或者某张单子已经等得太久，厨师就会开始按批次一起做菜。厨师在后台工作，这样前台服务员仍然可以不停地接新订单，而不会被阻塞。每当一道菜做好，就会立即被端给顾客；这就类似于 API 一旦有部分响应内容可用，就立刻把它发给客户端。这样顾客是一道一道菜地收到餐品，会感觉更加及时，也更有参与感。

实现

在实现上，首先我们会像之前一样搭建 FastAPI 服务器，但这一次会利用 lifespan 事件去启动一个模型 worker 线程。服务器启动时，会创建一个后台 worker 线程来运行 model_worker 函数。这个 worker 线程会被设置为 daemon，因此当主线程退出时它也会自动终止。shutdown_event 则用于在服务器关闭时，优雅地终止这个 worker 线程（fastapi_enhanced_server.py）。

代码清单 17.5 设置带后台 worker 线程的 FastAPI

shutdown_event = Event()
@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
  worker_thread = Thread(target=model_worker, daemon=True)
  worker_thread.start()
  yield
  shutdown_event.set()

class TextInput(BaseModel):
  text: str

app = FastAPI(lifespan=lifespan)

我们还要定义一些用于管理请求与响应的数据结构，它们将会被 worker 线程和服务器共享：

inference_queue = queue.Queue()
results = {}
results_lock = Lock()

其中，inference_queue 是一个线程安全的队列，用来保存进入的请求；results 字典则保存已经处理好的结果；results_lock 用来同步对这个字典的访问。

model_worker 函数负责不断从 inference_queue 中取出请求并处理。它会在一个持续循环中运行，检查是否有新请求到来，并把它们按 batch 批量处理。很多逻辑和我们之前做单请求处理时是相同的，只不过现在它要一次处理多个请求：

def model_worker() -> None:
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  model, tokenizer = get_model_and_tokenizer(device)

  while not shutdown_event.is_set():
    try:
      batch = get_batch_from_queue()
      if shutdown_event.is_set() or not batch:
          continue

      formatted_text = []
      for request_id, prompt in batch:
          text = [{"role": "user", "content": prompt}]
          results[request_id] = queue.Queue()
          formatted_text.append(tokenizer.apply_chat_template(text, tokenize=False))
    ...

逐 token 生成仍然会通过模型的 generate 方法来完成。这个方法会接收前面 tokenizer 已经处理好的 batch_tokens 和 batch_attention_mask。生成出的 token 会在每一步产生后立即被依次取出：

for _ in range(MAX_TOKENS):
    if not active_requests:
        break  # All requests are done

    with torch.no_grad():
        outputs = model.generate(
            batch_tokens,
            attention_mask=batch_attention_mask,
            max_new_tokens=1,
            pad_token_id=tokenizer.pad_token_id,
        )
    generated_token_ids = [output[-1:] for output in outputs]

    for idx in list(active_requests):
        request_id, _ = batch[idx]
        new_token_id = generated_token_ids[idx]
        new_token = tokenizer.decode(new_token_id)

    ...

    with results_lock:
        if request_id in results:
            results[request_id].put(new_token)

通过把 max_new_tokens 设成 1，我们就能确保模型一次只生成一个 token。然后我们把这个 token 放进 results 字典中对应 request ID 的结果队列里。这个过程会不断重复，直到达到最大 token 数，或者所有请求都完成。

注意：results_lock 用于同步对 results 字典的访问，保证多个线程读写它时不会造成数据损坏。

最后，当我们把结果流式返回给客户端时，可以借助 FastAPI 的 StreamingResponse 类。它允许我们在数据一可用时就往客户端发送，而不必等整个响应完全生成好之后再一次性发出：

@app.post("/generate")
async def generate_text(input_txt: TextInput) -> StreamingResponse:
    request_id = str(uuid.uuid4())
    inference_queue.put((request_id, input_txt.text))
    print(f"Received - Request ID: {request_id}, Prompt: {input_txt.text}")
    return StreamingResponse(stream_results(request_id), media_type="text/plain")

这段实现依赖于 stream_results 函数，它是一个异步生成器，会在 token 产生后一个一个地 yield 出来。该函数会从结果队列中不断读取 token，并在有内容可用时立即往客户端发送。这样我们就实现了实时流式返回。

我们可以通过下面的命令来启动这个服务器：

fastapi dev fastapi_enhanced_server.py

然后用这条 curl 命令来测试它：

curl -X 'POST' \
  'http://127.0.0.1:8000/generate' \
  -H 'accept: application/json' \
  -H 'Content-Type: application/json' \
  -N \
  -d '{
  "text": "What is PyTorch and why is it cool?"
}'

你会注意到，返回结果看起来像是在“实时打字”一样，而不是一次性完整返回——就像我们在 ChatGPT 或其他 AI 助手中看到的那样。这会给用户（或者基于你这个 API 再构建产品的开发者）带来强得多的交互体验。

同样的机制也适用于多请求批处理。我们可以同时向服务器发多个请求，而服务器会自动把它们凑成 batch 来处理。例如，我们可以创建多个 prompt 文件，然后用 curl 一起发出去：

echo "What is the capital of US?" > input1.txt
echo "What is the capital of France?" > input2.txt
echo "What is the capital of UK?" > input3.txt
for i in {1..3}; do
  cat input$i.txt | xargs -I {} curl -s -X POST \
    http://localhost:8000/generate \
    -H 'accept: application/json' \
    -H 'Content-Type: application/json' \
    -N \
    -d '{"text": "{}"}' >> input$i.txt &
done
wait

当然，你也可以直接开多个终端窗口，在每个窗口里各自执行那条 curl 命令，效果是一样的。你会发现这些 curl 请求会被并行处理，而且会在不同时间结束。

我们已经走得很远了！不妨先停下来欣赏一下自己的成果：我们已经构建了一个 Web 服务器，它能同时处理多个请求，把请求拼成 batch，在后台执行模型推理，并把结果一边生成一边流式返回给客户端。这已经是一个非常好的方式，可以把模型通过一个方便的 API 分享出去，让别人使用。

但这就结束了吗？当然没有。我们的服务器还可以继续优化，让它更高效。下一节里，我们会继续探索如何借助 PyTorch 原生技术，让模型跑得更快。

17.1.5 怎样让 PyTorch 模型跑得更快

深度学习社区如今已经越来越多地把 PyTorch 作为首选框架，不仅仅用于研究，也用于生产。作为回应，PyTorch 团队也不断把各种性能优化能力直接整合进框架本身。

使用 `torch.compile`

torch.compile 是在 PyTorch 2.0 中发布的，一经推出就引起了很大轰动。它之所以是一次重要的技术跃迁，是因为它解决了这样一个难题：如何在几乎不改动代码的前提下，让 PyTorch 代码跑得更快。

Python 是一种解释型语言，这意味着代码是由解释器在运行时逐条执行的，而不是预先编译成机器码。PyTorch 作为一个 Python 库，也天然继承了这一限制。然而，编译器通常非常有用，因为它能够把高级代码转换为更优化的机器级代码，从而显著提升性能。

torch.compile 的办法是采用 即时编译（JIT compilation） 。它会在程序运行时分析 PyTorch 代码，并自动生成各种优化，例如：

Kernel fusion（内核融合） ——把多个操作融合成一个内核，以减少调度开销并提升性能。比如 torch.cos()、relu() 这类逐元素操作，本质上都是一连串读取与比较，通常受限于内存带宽。如果每个操作都单独发起一个 kernel，开销会比较大；如果把它们融合起来，就能显著提高效率。
Operator reordering（算子重排） ——重新排列操作顺序，以改善内存访问模式、减少 cache miss。比如先做矩阵乘法再做逐元素加法，如果能在 kernel 层面把加法并进矩阵乘法阶段，那么就能减少中间结果的内存占用，并提高缓存利用率。

使用 torch.compile 非常简单，你只需要在现有模型外面包一层即可。比如我们在服务器里加载模型时：

model = AutoModelForCausalLM.from_pretrained(checkpoint)
model = model.to(device)         #1
model = torch.compile(model)   #2

#1 最佳实践是先把模型移动到目标设备上，再进行 compile。这样编译出来的模型才能针对实际运行的硬件进行优化。
#2 新增一行，对模型进行编译。

就这样， depending on 模型本身的特点，我们通常可以在推理中看到 20% 到 200% 不等的加速。

注意：加速效果取决于底层硬件以及模型架构。在某些少数情况下，甚至可能会变慢，因此它并不是一颗“万能银弹”。

本章后面，我们还会更深入地拆解一些 torch.compile 背后的“魔法”，并看看这整套技术栈是如何支撑起 PyTorch 其他特性的。

加入量化

如果我们想进一步降低模型的内存与计算开销，首先想到的通常是瘦身模型本身——也就是用更少的参数和更少的操作，去近似计算原来从输入到输出的映射。这类方法通常被称为 蒸馏（distillation） 。蒸馏的具体形式有很多：有时我们会删除一些很小或不重要的权重，从而缩小每个层；有时会把多个层压缩成一个（比如 DistilBERT）；甚至也可以训练一个完全不同、更简单的模型，去模仿大模型的输出（比如 OpenNMT 早期的 CTranslate）。之所以提这一点，是因为这些改造通常会是让模型推理更快的第一步。

另一种思路则是缩小每个参数和每次运算的表示精度：与其像平常那样每个参数使用 32 位浮点数，不如让模型改用整数（常见选择是 8 位）。这就叫做 量化（quantization） 。

注意：与量化不同，把训练过程（全部或部分）从 32 位浮点切换到 16 位浮点，通常称为 reduced precision 或 mixed precision training（如果仍有一部分保持 32 位）。

PyTorch 为此提供了量化张量（quantized tensors） 。它们以一组类似于 torch.float、torch.double 和 torch.long 的 scalar type 暴露出来（我们在 3.5 节见过类似概念）。最常见的量化张量 scalar type 是 torch.quint8 和 torch.qint8，分别表示无符号和有符号 8 位整数。PyTorch 在这里使用独立的 scalar type，是为了能利用我们在 3.11 节简要提到过的 dispatch 机制。

量化通常分为两类：量化感知训练（quantization-aware training） 和 训练后量化（post-training quantization） 。前者把量化过程直接纳入训练环节；后者则是在模型训练完成后，再对现有模型做量化。我们的演示会使用训练后量化，因为这里我们拿的是一个已经预训练好的模型，希望在不重新训练的前提下把它优化得更适合部署。首先，我们先看看量化之前模型有多大。我们准备了一个很方便的工具函数 get_serialized_model_size_in_mb，它会把模型序列化到字节缓冲区中，然后测量这个缓冲区的大小，以此估算模型的内存占用（quantization_example.ipynb）。

代码清单 17.6 测量模型大小

# In[]
model = AutoModelForCausalLM.from_pretrained(checkpoint)
size_before_quantization = get_serialized_model_size_in_mb(model)   #1
print(f"Serialized model size before quantization: {size_before_quantization:.2f} MB")

# Out[]
Serialized model size before quantization: 6528.52 MB

#1 获取模型序列化后的大小，单位是 MB。

我们可以通过把模型转换为更低精度的数据类型（比如 bfloat16）来应用量化。这个过程非常直接——只要对模型调用 .to(torch.bfloat16)，模型中的所有参数和 buffer 就会从原本的 32 位浮点格式自动转换为更节省内存的 16 位 brain floating-point 格式（quantization_example.ipynb）：

model = model.to(torch.bfloat16)
size_after_quantization = get_serialized_model_size_in_mb(model)
print(f"Serialized model size after quantization1: {size_after_quantization:.2f} MB")

PyTorch 支持多种可用于量化的数据类型，这些都可以在 torch.dtype 模块中找到（mng.bz/Rw9O），包括torch.int8 这样的 8 位有符号整数表示。乍一看，使用 8 位整数来替代 32 位浮点竟然还能工作，似乎有点反直觉；而且通常来说，模型效果确实会有一点轻微下降，但并不会很大。看起来主要有两个原因：第一，如果把舍入误差看成大致随机的，那么它们在大量操作中往往会互相抵消；第二，卷积层和线性层本质上都像是某种加权平均，因此量化误差不一定会被显著放大。

不过，需要注意的是：虽然量化确实能显著减小模型体积并加快推理，但我们不能简单地对整个模型直接来一发 .to(torch.int8) 。不同类型的模块，对量化的支持程度并不一样——线性层通常很适合 int8 量化，而归一化层或某些复杂激活函数则可能仍然需要浮点精度，或者需要特殊处理。另外，量化本身也会引入精度权衡，而有些操作对这种精度损失更敏感，可能会进而影响模型准确率。下面这个示例里，我们会实现一个自定义量化函数：它专门针对模型中的线性层，把它们的权重量化成 int8；而模型其余部分则保留在 bfloat16（quantization_example.ipynb）。

代码清单 17.7 针对线性层的自定义 int8 量化函数

...

def quantize_linear_layers_to_int8(model):
    for _, module in model.named_modules():
        if isinstance(module, torch.nn.Linear):
            weight = module.weight.data
            w_int8, scale = int8_symmetric_quantize(weight)
            module.register_buffer('weight_int8', w_int8)     #1
            module.register_buffer('weight_scale', scale)
            delattr(module, 'weight')      #2

            def new_forward(self, x):     #3
                dequantized_weight = self.weight_int8.to(x.dtype) * self.weight_scale
                return torch.nn.functional.linear(x, dequantized_weight, self.bias)
            import types
            module.forward = types.MethodType(new_forward, module)
    return model

model = AutoModelForCausalLM.from_pretrained(checkpoint)
model = model.to(torch.bfloat16)
model = quantize_linear_layers_to_int8(model)
size_after_quantization2 = get_serialized_model_size_in_mb(model)
print(f"Serialized model size after quantization2: {size_after_quantization2:.2f} MB")

# Out[]
Serialized model size after quantization2: 1825.69 MB

#1 把 int8 权重和 scale 直接存到模块内部。
#2 删除原始权重，以节省空间。
#3 创建一个新的 forward 方法，改为使用量化后的权重。

非常酷。我们成功地把模型大小从 6.5 GB 降到了 1.8 GB，这意味着体积减少了 70% 以上！效果很惊人。当然，在实际工作里，我们一般不会为每个模型手写这些逻辑。正如你所预期的，PyTorch 也内置了对量化的支持，能自动完成这类工作。torch.quantization 模块提供了一整套量化框架，支持多种量化方案和技术（quantization_example.ipynb）。

代码清单 17.8 使用 PyTorch API 应用动态量化

model = AutoModelForCausalLM.from_pretrained(checkpoint)
quantized_model = torch.quantization.quantize_dynamic(
    model,
    {torch.nn.Linear},    #1
    dtype=torch.qint8
)

#1 这里只对线性层进行量化。

量化依然是 PyTorch 中一个非常活跃的研究方向，各个团队都在不断推进其边界。想到研究者们如今甚至已经在尝试对大语言模型做 1-bit 量化——也就是权重几乎只用 0 和 1 来表示（参见 arxiv.org/html/2402.1…）——确实让人觉得不可思议！我们可以预期，未来的 PyTorch 版本还会继续在量化 API 上做出改进和变化。如果你想更深入了解 PyTorch 的量化能力、技术路线和最佳实践，可以访问官方文档：pytorch.org/docs/stable…。

现有技术生态

我们现在已经构建了一个服务器，并且让它对 LLM 支持了多种特性和优化。虽然我们当然还可以自己继续往上加更多功能，但在很多情况下，更好的做法是利用已有技术，而不是重复造轮子。下面有几个值得重点关注的项目：

vLLM——一个专门为高效运行 LLM 而设计的工具，尤其擅长在同时服务大量用户时提升效率（github.com/vllm-projec…）。
LitServe——一个基于 FastAPI 构建的灵活 AI 模型服务引擎。像 batching、streaming、GPU 自动扩缩容这类特性都已经内置，不需要你为每个模型重新手写一套 FastAPI 服务器（github.com/Lightning-A…）。

17.2 导出模型

到目前为止，我们一直是通过 PyTorch 提供的抽象，直接在 PyTorch 内部执行模型。对于快速迭代与开发来说，这种方式非常好；但当我们真正要把模型部署到生产环境中时，往往需要把模型导出到另一种格式，或者交给其他框架来执行。尤其当我们希望在不同硬件上运行模型，或者希望在Python 之外的编程语言中使用模型时，这一点就变得尤为重要。

在 Python 里，我们会遇到 GIL，也就是全局解释器锁，它会阻止多个线程同时执行 Python 字节码。这对单线程应用问题不大，但对多线程应用来说可能会成为瓶颈。（当然，随着无 GIL Python 的出现，参见 peps.python.org/pep-0703/，这件事未来也许会成为历史。）另一种场景则是嵌入式系统：在这类环境中，完整安装整套 Python 运行栈要么过于昂贵，要么根本不可行。这时候，就需要把模型导出出来。

所谓导出模型，就是把模型转换成一种能够在不同环境与平台中使用的格式。这个过程通常会得到某种标准化表示，使得模型能够被各种不同系统解释和执行。它并不是直接把模型翻译成面向某种硬件的底层指令，而是先把模型准备成一个便于后续继续优化并在指定硬件上执行的中间表示。这有点像编译器：先把高级代码翻译成某种中间形式，再进一步编译成机器码。

在这里，我们主要有两条路线可选：

迁移到 PyTorch 之外的专用框架
使用 PyTorch 自带的模型导出工具

下面分别来看。

17.2.1 借助 ONNX 实现 PyTorch 之外的互操作性

有些场景下，我们需要把模型从 PyTorch 生态中“带出去”——例如，运行在某些使用专用部署流水线的嵌入式硬件上。为此，Open Neural Network Exchange（ONNX） 提供了一种神经网络与机器学习模型的互操作格式（onnx.ai）。一旦模型被导出成 ONNX，就可以使用任何兼容 ONNX 的 runtime 来执行，比如 ONNX Runtime，只要模型中所用的操作被 ONNX 标准以及目标 runtime 所支持即可。比如，在 Raspberry Pi 上，用 ONNX Runtime 执行通常就会比直接跑 PyTorch 更快得多。除了常规硬件外，还有很多专用 AI 加速器硬件同样支持 ONNX（onnx.ai/supported-t…）。

注意：ONNX Runtime 的代码位于 github.com/microsoft/o…。不过要记得查看其隐私声明。如果你从源码自行构建 ONNX Runtime，那么最终得到的包中将不会包含任何遥测或数据共享功能。

从某种角度看，一个深度学习模型其实就是一个拥有特定指令集的程序，而这些指令由诸如矩阵乘法、卷积、relu、tanh 等底层操作构成。既然如此，只要我们能把这个计算过程序列化下来，就能在另一个理解这些底层操作的 runtime 中重新执行它。ONNX 本质上就是把这些操作及其参数的描述方式做了标准化。

如今，大多数现代深度学习框架都支持把自己的计算过程序列化为 ONNX，有些框架甚至还能反过来加载 ONNX 文件并执行它（不过 PyTorch 本身并不支持直接加载 ONNX 执行）。某些低资源边缘设备可以直接接收 ONNX 文件，并为具体硬件生成底层指令。还有一些云计算平台现在也支持上传 ONNX 文件，并把它直接暴露成一个 REST 服务端点。

要把模型导出成 ONNX，我们需要用一个dummy input（伪输入） 跑一遍模型。输入 tensor 的具体值并不重要，真正重要的是它们的形状和类型必须正确。调用 torch.onnx.export 后，PyTorch 会追踪模型执行时发生的计算，并把这些计算序列化到一个指定名称的 ONNX 文件中（onnx_example.ipynb）：

onnx_model_path = "simple_model.onnx"
dummy_input = torch.randn(1, 10)
torch.onnx.export(       #1
    model,
    dummy_input,
    onnx_model_path
)

#1 将模型导出为 ONNX 格式。

得到的 ONNX 文件现在就可以被 runtime 运行、被编译到某个边缘设备上，或者被上传到云服务中。若要在 Python 中使用它，只需安装 onnxruntime（如果有 GPU，则安装 onnxruntime-gpu），然后把输入 batch 准备成一个 NumPy 数组（onnx_example.ipynb）：

import onnxruntime
sess = onnxruntime.InferenceSession("simple_model.onnx")   #1
batch = dummy_input.numpy()
input_name = sess.get_inputs()[0].name
pred_onnx, = sess.run(None, {input_name: batch})

#1 ONNX Runtime 的 API 通过 session 来加载模型，并通过 run 方法接收一组带名字的输入。

17.2.2 PyTorch 自己的导出方式：`torch.export`

正如前面提到的，PyTorch 也有自己的一套模型导出工具。把 PyTorch 模型导出为 ONNX，主要目的是为了和其他框架以及其他 runtime 实现互操作。如果你需要跨框架兼容性，那么 ONNX 是最合适的选择；但如果你仍然希望留在 PyTorch 生态中，只是想把模型优化得更适合部署，那么 torch.export 就是一个非常好的选择。

PyTorch 的 torch.export 会接收一个 torch.nn.Module，并产出一个被追踪出来的图（traced graph） 。这个图就是模型计算图的一种表示，它刻画了模型内部的操作顺序和数据流动方式，描述输入是如何经过层层变换最终得到输出的。这个 traced graph 是以一种 AOT（ahead-of-time）编译友好 的形式被描述出来的。

你可以把 PyTorch 中的 traced graph 想象成一张流程图：它描绘出模型内部逻辑操作的执行顺序。正如流程图中每一个符号代表一个动作或决策点一样，traced graph 中的每一个节点都对应某个具体操作或变换。它告诉你：数据是怎样穿过模型各层、经过哪些操作、最终得到结果的。

追踪一个模型的方式和 ONNX 导出看起来其实非常相似。因此，也就不奇怪了：你甚至可以利用 PyTorch 的 export 技术来把模型导出成 ONNX。所需要额外加上的，只是给 export 调用加入 dynamo=True（onnx_example.ipynb）：

torch.onnx.export(
    model,
    dummy_input,
    onnx_model_path,
    dynamo=True,      #1
)

#1 加上这一行，以使用 TorchDynamo 来导出模型。

如果直接使用 torch.export，同样需要向模型传入一个示例输入，以便追踪模型（torch_export_example.py）：

from torch.export import export
dummy_input = torch.randn(2, 10)
exported_program = export(model, args=(dummy_input,))

无论是这里的 torch.export，还是上面 ONNX 导出时加上 dynamo=True 的方式，底层都利用了 TorchDynamo。TorchDynamo 是 PyTorch 中一个在运行时捕获和优化计算图的工具。它在 PyTorch 2.0 中首次登场，可以说是 PyTorch 编译器生态的一次重大跃迁，因为它让“动态追踪模型”成为可能，而这类用户体验在以往的静态追踪方法中是很难实现的。TorchDynamo（也常简称为 Dynamo）通过拦截 Python 字节码的执行来完成这一点，从而在不要求用户修改底层代码的前提下，捕获程序执行过程中发生的操作。

TorchDynamo 本质上是一个 tracer（追踪器） 。也就是说，给它一个函数和一个示例输入，它就会把这段执行过程记录成一条线性的指令序列，并保存成一个图。这种图实际上是一个 FX graph（参见 torch.fx 文档：docs.pytorch.org/docs/main/f…）。我们不会深入所有细节，但你可以把它理解成一个容器：它装着函数以及它们的参数。后续编译优化或导出成其他格式时，就会继续利用这些容器。下面我们来看看这种表示长什么样（torch_export_example.ipynb）。

代码清单 17.9 查看导出后的程序图表示

# In[]
print(exported_program)

# Out[]
ExportedProgram:
  class GraphModule(torch.nn.Module):
    def forward(self, p_fc1_weight: "f32[5, 10]", p_fc1_bias: "f32[5]",
    ↪ p_fc2_weight: "f32[2, 5]", p_fc2_bias: "f32[2]", x: "f32[2, 10]"):
      # File: .py:10 in forward, code: x = self.fc1(x)
      linear: "f32[2, 5]" = torch.ops.aten.linear.default(x, p_fc1_weight, p_fc1_bias);
      ↪ x = p_fc1_weight = p_fc1_bias = None

      # File: .py:11 in forward, code: x = self.relu(x)
      relu: "f32[2, 5]" = torch.ops.aten.relu.default(linear);  linear = None

      # File: .py:12 in forward, code: x = self.fc2(x)
      linear_1: "f32[2, 2]" = torch.ops.aten.linear.default(relu, p_fc2_weight, p_fc2_bias);
      ↪ relu = p_fc2_weight = p_fc2_bias = None
      return (linear_1,)

你会注意到，这个模型仍然有一个 forward 方法，只不过原来的参数现在都变成了这个方法的输入参数。这是因为参数已经不再“嵌在模型对象内部”，而是成为了计算图的一部分。原本的 nn.Linear 和 nn.ReLU 层，也被转换成了更底层的 torch.ops.aten 操作。此外，像 f32[2, 5] 这样的标注，表示参与计算的 tensor 的数据类型和形状。

不过，这个底层 graph module 本身依然还是一个 torch.nn.Module，这就意味着它甚至依然可以像普通 PyTorch 模型一样被直接调用（torch_export_example.ipynb）：

exported_program.module()(dummy_input)    #1

#1 调用导出程序中的 module 依然可用！

接下来，我们就可以使用 AOTInductor（AOTI） ，也就是 PyTorch 的 ahead-of-time 编译器，来执行推理：

import os
output_path = torch._inductor.aoti_compile_and_package(
    exported_program,
    package_path=os.path.join(os.getcwd(), "model.pt2"),
)

torch._inductor.aoti_compile_and_package 会把模型编译好，并打包成一个文件，之后可以再重新加载。它有点类似于我们用 torch.save 保存 PyTorch 模型，但不同之处在于，这里面还额外包含了编译后的代码：

model = torch._inductor.aoti_load_package(os.path.join(os.getcwd(), "model.pt2"))
model(dummy_input)

这样就可以调用它了。通过 torch._inductor.aoti_load_package 加载模型，本质上就是把这个编译后的模型重新加载进来，并允许我们继续用之前那样的输入去运行它。你可以把它理解成编译模型版本的 torch.load。需要注意的一点是：输入 tensor 的形状必须和导出时使用的形状一致，因为这个编译后的模型是针对那些特定输入形状做了优化的，换了形状之后可能无法正确运行。

提示：你确实可以在没有保留原始源代码的情况下，直接运行导出的 PyTorch 模型。但在实践中，我们始终应当建立这样一条自动化工作流：从“源模型代码”自动生成并部署“已安装好的 JIT/AOT 模型”。否则，未来当你想微调一点模型行为时，很可能会发现自己已经失去了重新修改与生成它的能力。

到这里为止，我们已经给你介绍了不少新术语，所以此时很适合停下来回顾一下刚才到底完成了什么。正如图 17.3 所示，我们起点是用户自己写的模型代码，也就是一个继承自 torch.nn.Module 的 PyTorch 模型。接着，我们通过 torch.onnx.export 或 torch.export 把这个模型导出成一个 ONNX 文件或 PyTorch export 文件。这个导出过程会捕获模型的计算图，并把它保存成一种特定格式。这个导出后的 program 可以像普通 PyTorch 模型一样被保存和加载，但它现在已经变成了一种更适合被进一步优化和高效执行的表示形式。然后，我们把它交给一个编译器，编译器会对图进行优化，并为目标硬件生成底层代码。最终，这段编译好的代码会在目标硬件上执行；目标硬件可以是 CPU、GPU 或专用加速器。执行过程包括加载编译好的代码，并用输入数据驱动它，从而得到最终输出。

图 17.3 从用户模型代码，到导出模型，再到执行的全过程

17.3 进一步理解 `torch.compile`

我们刚刚讨论的整个流程——从用户模型代码，到导出模型，再到最终执行——在逻辑上其实和 torch.compile 背后的事情非常相似。前面在讨论如何提升批处理服务器性能时，我们已经简要提到过 torch.compile。考虑到 PyTorch 团队对编译器这部分投入很大，而且还在持续活跃开发，我们觉得值得再深入看一看。

那么，torch.export 和 torch.compile 到底有什么区别？在什么情况下该选哪个？要回答这个问题，我们得先弄清楚什么叫做 完整图捕获（full graph capture） ，以及什么叫 离散图（disjoint graphs） 。

17.3.1 完整图捕获 vs. 离散图

torch.export 依赖的是完整图捕获。也就是说，它会捕获模型的整个计算图，包括其中所有操作及其依赖关系。这种方式能得到对模型行为非常完整的表示；但缺点是，它没法处理某些行为，例如动态控制流，或者某些依赖输入数据本身的操作。因此，如果你的模型里存在动态行为——比如条件分支、循环，而且这些分支或循环取决于输入内容——那么 torch.export 可能就无法正确捕获。例如在下面这个例子中，模型会根据输入值不同，走不同的前向路径（compile_example.ipynb）。

代码清单 17.10 带有动态条件控制流的模型

import torch
import torch.nn as nn
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
class ConditionalModel(nn.Module):
    def __init__(self):
        super(ConditionalModel, self).__init__()
        self.fc1 = nn.Linear(10, 5)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(10, 3)
    def forward(self, x):
        if x.sum() > 0:
            x = self.fc1(x)
        else:
            x = self.fc2(x)
        x = self.relu(x)
        return x
model = ConditionalModel()
model = model.to(device)
model.eval()

接下来，当我们调用 torch.export 去追踪这个模型时，就会报错：

# In[]
example_input = torch.randn(5, 10)
torch.export.export(model, (example_input,))

# Out[]
from user code:
   File "/var/...py", line 13, in forward
    if x.sum() > 0:                           #1
UserError: Dynamic control flow is not supported at the moment.

#1 这部分条件控制流当前不被 torch.export 支持。

之所以出错，是因为 torch.export 期待的是一张静态计算图，而这个 if 条件带来了动态行为，无法被表达进静态图中。如果它强行根据这个示例输入去推断，那就只能捕获图中的一条执行路径，而那显然无法代表模型的全部行为。

相比之下，torch.compile 采用的是一种更灵活的策略：当它遇到无法编译的操作时，会把图拆成更小的片段。这使得它能够处理动态行为，并且仍然尽可能高效地执行。虽然图被打断（graph breaks）之后会有一定性能损失，但这种折衷换来了对动态控制流模型的支持——例如我们刚刚定义的这个模型。使用 torch.compile 时，它会尽量捕获完整计算过程中的所有路径，并为这些路径生成优化后的执行逻辑。

`torch.compile` 的历史背景与动机

TensorFlow 曾经是深度学习框架中的绝对主角，它依赖的是一种“先定义整张图、再统一执行”的计算图模式。这种方式在性能优化上非常有利，但调试和开发体验却比较痛苦。

后来，PyTorch 作为一个更易用的替代方案出现了。它提供了按步骤执行模型的机制，也就是我们熟悉的 eager mode。这种模式下，调试变得简单得多：你可以随手 print，也可以直接用 Python 调试器单步跟踪。因此，PyTorch 很快变得更容易上手，也逐渐成为深度学习领域的首选框架。虽然 TensorFlow 2.0 后来也引入了 eager mode，但那个时候 PyTorch 已经基本坐稳了主流位置。

不过，性能优化和生产部署对“图执行”模式的需求始终存在。于是，PyTorch 最初通过 JIT（just-in-time）编译器 来应对这一需求：模型先按 eager mode 写出来，再被编译成静态图，以便获得更高性能。最开始的方式是 torch.jit.trace，它通过跟踪示例输入上的操作来捕获计算图，并把这些图表示成一种中间表示，叫做 TorchScript。但 torch.jit.trace 有很明显的局限，尤其在面对动态控制流或依赖输入数据的操作时，表现很差。

之后，PyTorch 又推出了 torch.jit.script，试图让这种图执行模式更方便一些。它允许用户用更“Pythonic”的方式编写模型，比如照常写循环和条件分支，再由 JIT 编译器去分析代码并生成静态图表示。这比 trace 前进了一大步，但它本质上是在尝试把整个 Python 重新实现成一种静态语言。结果就是，一旦你在代码里用到某个它还没实现的 Python 特性，就会在运行时报错，这对用户来说体验并不好。实践中，这就导致 torch.jit.script 只能支持一部分模型；否则，就得为了适配它而对原有模型做大规模重构。

而 torch.compile 引入之后，则通过 TorchDynamo 提供了动态图捕获的能力。这样一来，用户就可以继续像以前一样用 eager mode 写模型，同时又能享受编译优化。如果 Dynamo 遇到自己无法识别或无法编译的控制路径，它就把图切开，把剩下的部分退回到普通 Python 执行。这避免了报错，也让“离散图”成为可能。这种灵活性，是对早期深度学习框架的一次重大突破；过去很多框架要求用户手工构造图、使用刚性很强的模型定义方式，而现在 torch.compile 提供了更灵活也更友好的编译流程，让研究者可以把更多精力放在创新本身，而不是去对抗底层基础设施。

如今，你仍然可能在资料中看到 torch.jit.trace、torch.jit.script 和 TorchScript 这些名字，但它们已经不再是活跃开发对象，而更多处于维护状态。现在，真正应该重点关注的是 torch.export 和 torch.compile，它们才是 PyTorch 当前主推的模型优化与部署工具。

我们刚才拿来导出的那个模型，也可以毫无问题地直接用 torch.compile 来编译（compile.ipynb）：

compiled_model = torch.compile(model)
compiled_model

调用编译后的模型，与调用普通 PyTorch 模型几乎没有区别，而且它能够正确处理刚才那个动态控制流。我们可以对比原始模型与编译后模型的输出，看看它们是否一致：

input_data = torch.randn(5, 10).to(device)
output_original = model(input_data)
print("Output from original model:", output_original)
output_compiled = compiled_model(input_data)
print("Output from compiled model:", output_compiled)

当你使用 torch.compile 并执行编译模型的 forward 时，内部其实会经历多个阶段。首先，系统会检查当前函数是否应该跳过编译，比如当里面包含非 PyTorch 操作时。如果这段代码以前已经编译过了，它会尝试复用缓存版本；如果没有，它就会分析模型中的操作，生成一张 torch.fx (FX) graph，也就是一份可以被优化的计算表示。之后，这张 FX graph 会交给指定 backend 去编译，以提升性能。如果分析没有覆盖到整个函数，那么它还会额外生成若干函数，去处理剩余那些未被纳入图中的部分。最后，它会生成新的字节码，用来执行这个优化后的 FX graph，并同时管理必要的状态与副作用，从而确保模型能够高效运行。

torch.compile 默认使用的 backend 是 TorchInductor。TorchInductor 会把计算图下沉成高效的底层代码，这些代码可以运行在 CPU、GPU 等不同硬件架构上。它会利用诸如 OpenMP（用于 CPU）和 Triton（用于 GPU，github.com/triton-lang…）这样的技术，来生成新的高效 kernel，或者调用已有的优化 kernel。

当然，如果你希望采用别的优化策略，或者针对特定硬件架构进行编译，也可以显式指定其他 backend。很多硬件厂商和框架都已经为 torch.compile 提供了自己的 backend。如果你想尝试自己开发一个 backend，可以参考这个示例：mng.bz/269m。

归根结底，相比走 torch.export 这条导出路线，我们也可以在 PyTorch 框架内部直接利用相似的技术来把模型编译掉。这样既保留了原有 PyTorch 用户代码的形式，又能享受到编译优化带来的性能收益。

17.4 使用 `torch.profiler` 理解执行过程

虽然 torch.export 和 torch.compile 提供了非常强大的优化能力，但它们有时也会让人觉得像黑盒。要真正理解模型执行过程中到底发生了什么、性能瓶颈到底卡在哪里，我们就需要对运行中的代码有更强的可见性。而这正是 torch.profiler 非常有用的地方。

PyTorch 内置了一个 profiler，可以帮助我们分析模型在训练和推理过程中的性能。它能采集多种指标，例如执行时间、内存使用情况，以及不同操作执行的次数。这些信息可以帮助我们发现瓶颈，并进一步优化模型。

在接下来的例子中，我们会使用一个非常简单的模型，它包含多个前馈层。我们会用一个 dummy input 跑这个模型，并利用 profiler 抓取它的执行细节（profiler_example.ipynb）。

代码清单 17.11 使用 PyTorch profiler 采集模型性能指标

from torch.profiler import profile, ProfilerActivity

def run_model(model, input_data, device="cpu", warmup_iters=5):
    ....
    activities = [ProfilerActivity.CPU]
    if device == "cuda":
        activities.append(ProfilerActivity.CUDA)
    with profile(
        activities=activities,
    ) as prof:
        output = model(input_data)

    display = prof.key_averages().table(sort_by="cpu_time_total", row_limit=10)
    ...

print(run_model(model, input_data))

这段代码创建了一个 profiling 上下文，会根据设备参数去捕获 CPU 或 CPU+GPU 上的计算活动。在执行若干次 warmup 迭代（代码中未展示）之后，profiler 会开始监控模型的前向传播，采集推理期间各类操作的详细耗时信息。等 profiling 结束后，我们就能生成一张性能汇总表，列出最耗时的前 10 个操作，并按 CPU 总耗时排序：

----------------------  ------------  ...  ------------
                  Name    Self CPU %  ...    # of Calls
----------------------  ------------  ...  ------------
          aten::linear         0.79%  ...            20
           aten::addmm        65.43%  ...            20
           aten::copy_        18.10%  ...            20
            aten::relu         1.27%  ...            10
       aten::clamp_min        11.68%  ...            10
               aten::t         1.45%  ...            20
       aten::transpose         0.38%  ...            20
                   ...                ...           ...
----------------------  ------------  ...  ------------
Self CPU time total: 19.152ms

这对于从高层次了解模型性能很有帮助，但它还不足以告诉我们完整的执行流，也看不到底层操作之间的关系。如果想进一步查看这些内容，我们可以让 torch.profiler 导出一份执行 trace：

prof.export_chrome_trace("trace.json")

这会生成一个 JSON 格式的 trace 文件，用来可视化执行流程，包括每个操作耗时多久、它们之间有哪些依赖关系。你可以在 Chrome 浏览器中打开 chrome://tracing（或者新版的 ui.perfetto.dev/）来查看这些 trace。

在 trace 可视化图中，如图 17.4 所示，界面大致分成上下两部分。上半部分展示的是 CPU 的执行时间线，其中每一行表示一个操作，按时间先后排列。你可以看到 CPU 上的操作是如何触发 GPU stream 上相应任务的。下半部分则展示 GPU 的执行时间线，反映这些操作是如何在 GPU 硬件上被实际处理的。通过在模型的 forward 方法中加入 record_function() 装饰器，我们还能在 trace 中插入自定义标签，例如 “Layer 1 Forward” 和 “Layer 2 Forward”，从而更容易看出哪些操作对应网络中的哪些层。

图 17.4 模型执行过程的 trace

而当我们把模型编译之后，情况会明显不同，如图 17.5 所示：多个前向过程会被融合在一起，操作会归于一个 “Compiled Region” 的范围下。与此同时，GPU 上也会出现融合后的自定义 kernel，例如针对线性层的优化 kernel，这些都能带来更好的性能。

图 17.5 编译后模型执行过程的 trace

17.5 在 Python 之外使用 PyTorch

到目前为止，我们已经展示了如何在 Python 中使用 PyTorch，以及如何对模型做效率优化、为模型提供服务，并把它导出到其他格式中。不过，有些时候我们希望在 Python 之外使用 PyTorch 模型，比如把它集成进 C++ 应用，或者部署到移动设备上。这时候，就可以用到 PyTorch 生态系统中的其他组件，它们允许我们在不同语言和不同环境中运行 PyTorch 模型。

这里我们会重点介绍两个方向：

用于 C++ 的 LibTorch（docs.pytorch.org/cppdocs/ins…）
用于移动端的 ExecuTorch（mng.bz/15Kn）

17.5.1 LibTorch：在 C++ 中使用 PyTorch

LibTorch 是 PyTorch 的 C++ 发行版，它提供了对 PyTorch 功能的直接访问，而无需依赖 Python。实际上，PyTorch 的很多核心组件本身就是用 C++ 编写的，Python 层只是给它们包了一层用户熟悉的接口。LibTorch 则让你可以直接使用这些底层的 C++ 组件。

C++ API 在设计上是刻意对齐 Python API 的，因此对已经熟悉 PyTorch 的开发者来说，迁移会轻松很多：

// Python
model = torch.nn.Linear(10, 5)
input = torch.randn(1, 10)
output = model(input)

// Equivalent C++
auto model = torch::nn::Linear(10, 5);
auto input = torch::randn({1, 10});
auto output = model->forward(input);

不过，既然 Python 已经成为深度学习社区的事实标准语言，那么 C++ API 的功能完备度自然没有 Python API 那么高；有些特性只在 Python 中提供（例如某些分布式训练 API）。所以，对大多数 PyTorch 开发来说，Python 仍然是首选；LibTorch 更适合用在两类场景中：一是要把 PyTorch 集成进现有 C++ 应用，二是部署到那些不适合运行 Python 的环境。

使用 C++ 的一个痛点就是：入门门槛没那么低，所以接下来这部分配置过程会稍微辛苦一点。好消息是，一旦搭起来，在 C++ 里运行 PyTorch 模型本身其实并不复杂。

我们将使用 CMake 来构建 C++ 应用，所以请先确保你已经安装好 CMake。CMake 是一个构建系统，用于管理 C++ 项目的编译过程，包括处理依赖以及为不同平台生成构建文件。你可以通过系统包管理器安装它，或者从 CMake 官网（cmake.org/download/）下载安装：

brew install cmake   # mac
apt install cmake    # ubuntu

接下来，你需要有 LibTorch。如果你已经通过 pip 安装了 PyTorch，那么 LibTorch 的二进制文件其实已经包含在当前 PyTorch 安装目录里了。当然，如果你希望单独下载，也可以直接从 PyTorch 官网获取最新版 LibTorch（pytorch.org/get-started…）。

之后，我们需要告诉 cmake 到哪里去找 LibTorch 所需的 CMake 配置文件。具体做法是设置 CMAKE_PREFIX_PATH 环境变量，指向 LibTorch 所在路径。然后我们可以生成一个单独的 build 目录，并运行 CMake 进行项目配置。把构建文件和源文件分开，是一种很好的项目组织习惯——因为这样一来，想清理编译产物时，只要直接删掉整个 build 目录即可：

cmake -B build -S . -DCMAKE_PREFIX_PATH=\
 $(python3 -c 'import torch;print(torch.utils.cmake_prefix_path)')

接着，就可以用 CMake 来真正编译项目了。这个过程会编译 C++ 代码，并把它和 LibTorch 链接起来：

cmake --build build --config Release

然后我们就可以运行编译出来的 C++ 可执行程序。这个可执行文件会位于 build 目录中，直接在那里运行即可：

./build/example

我们示例的第一部分，会展示最基本的 LibTorch 用法：创建一个简单的线性模型，并在上面做推理。这部分 C++ 代码和我们之前的 Python 代码几乎是对应的，从中也可以看出两种 API 之间有多相似。

而示例的第二部分，则展示如何创建并训练一个简单神经网络。这段代码定义了一个 Net 类，它继承自 torch::nn::Module，也就是 Python 中 torch.nn.Module 的 C++ 版本。这个网络是一个简单的前馈神经网络，用于 MNIST 数字分类，结构包括：

一个输入层，接收展平后的 28 × 28 图像（共 784 个特征）
一个带 128 个神经元和 ReLU 激活的隐藏层
一个输出层，包含 10 个神经元（分别对应 10 个数字），并使用 log-softmax 激活

训练过程看起来也和我们在 Python 里做过的事情非常相似（example.cpp）。

代码清单 17.12 在 C++ 中使用 LibTorch 训练模型

void model_train_example() {
  auto model = std::make_shared<Net>();    #1
  torch::optim::SGD optimizer(model->parameters(), /*lr=*/0.01);

  // Create fake data
  auto x = torch::randn({64, 784});
  auto y = torch::randint(0, 10, {64});

  model->train();
  for (size_t epoch = 0; epoch < 3; ++epoch) {    #2
    auto prediction = model->forward(x);
    auto loss = torch::nll_loss(prediction, y);
    optimizer.zero_grad();
    loss.backward();
    optimizer.step();

    std::cout << "Epoch: " << epoch << " | Loss: " << loss.item<float>()
              << std::endl;
  }

  // Inference
  model->eval();
  torch::NoGradGuard no_grad;      #3
  auto test_input = torch::randn({1, 784});
  auto output = model->forward(test_input);
  auto predicted = output.argmax(1);

  std::cout << "Prediction: " << predicted.item<int64_t>() << std::endl;
}

#1 创建模型实例和一个 SGD 优化器。
#2 使用标准训练循环训练模型 3 个 epoch。
#3 使用训练好的模型执行推理。

最后，在最终示例中，我们会使用 ahead-of-time (AOT) inductor 功能，把 PyTorch 模型编译并保存成文件。这样，这个模型之后就能在 C++ 应用中被加载和执行，而不需要保留原始 Python 代码：

torch::inductor::AOTIModelPackageLoader loader("./llm.pt2");
std::vector<torch::Tensor> outputs = loader.run(inputs);

17.6 走向移动端：ExecuTorch

作为部署模型的最后一种形态，我们来讨论一下部署到移动设备。在快速发展的机器学习世界里，能把模型部署到各种资源受限设备上——包括手机和嵌入式设备——变得越来越重要。ExecuTorch 正是 PyTorch 生态中针对这一挑战给出的解决方案。它的目标，是让 PyTorch 的运算能够在那些传统上不支持大量 PyTorch 操作的设备上高效运行。

这类设备，比如手机或嵌入式系统，通常没有足够的算力和资源去支持 PyTorch 全量那几千种操作。ExecuTorch 所做的，就是把这条鸿沟填上：它通过针对受限设备优化 PyTorch 模型的执行方式，让这些设备也能够有效执行复杂计算。

ExecuTorch 的工作流大致分成两个阶段：

首先，用 torch.export 捕获模型的计算图。
图捕获完成后，ExecuTorch 会提供一个针对受限设备量身定制的 runtime 环境。在服务器级硬件上，这个角色通常由推理服务来承担；而 ExecuTorch 则把同样的思路适配到了移动端和嵌入式设备的限制与需求上。

捕获这张图的过程，与我们之前用 torch.export 时很相似，只不过需要额外再做一步：把导出的 program 转换成 ExecuTorch 的 Edge IR 格式。这是一种专为移动设备高效执行而设计的表示形式。

代码清单 17.13 将模型导出为 ExecuTorch

import torch
from torch.export import export
from executorch.exir import to_edge

exported_program = export(model, example_input)    #1
edge_program = to_edge(exported_program)    #2
executorch_program = edge_program.to_executorch()    #3
with open(output_path, "wb") as f:
    f.write(executorch_program.buffer)
return executorch_program

#1 使用 torch.export 导出模型。
#2 将导出的 program 下沉为 ExecuTorch 的 Edge IR 格式。
#3 把 Edge IR 序列化成一个可加载的 ExecuTorch 包，其字节内容可供移动端 runtime 加载。

接下来，我们就可以在移动应用中加载这个 package 并使用它。为了演示方便，我们仍然用 Python 展示这一过程，但其基本原理在移动环境中也是一样的。

代码清单 17.14 使用 ExecuTorch 模型执行推理

from executorch.extension.pybindings.portable_lib import _load_for_executorch

executorch_module = _load_for_executorch(model_path)
outputs = executorch_module.forward([input_data])

理论上，LibTorch 和 C++ 版本的 PyTorch 都可以被编译到 Android 和 iOS 平台上。在 Android 上，我们可以通过 Java Native Interface（JNI）从 Java 层调用这些组件；而在 iOS 上，则可以利用 Objective-C++ 在 Swift / Objective-C 和 C++ 之间做桥接。

但在很多时候，我们真正需要的其实只是 PyTorch 的极少几个能力：比如加载一个已经 JIT/AOT 编译好的模型，把输入转成 tensor 或 IValue，送入模型执行，再取出结果。为了避免在 Android 上折腾 JNI，或者在 iOS 上编写 Objective-C++ 桥接代码，ExecuTorch 让在移动设备上运行 PyTorch 模型这件事变得更简单。

移动端开发的环境配置，比普通 C++ 开发还要多出一些额外步骤，而且这些步骤本身与深度学习无关。因此，我们就把这部分具体配置细节留作读者练习。若你想查看官方说明，可以参考 PyTorch 文档中关于 Android 和 iOS 设置 ExecuTorch 的说明：mng.bz/Pw0v。

17.7 结论

至此，我们完成了对 PyTorch 部署选项的一次整体巡礼：从使用 Gradio 快速做 demo，到使用 FastAPI 构建生产 API；从模型优化技术，到 ONNX 带来的跨平台兼容性；再到通过 LibTorch 集成进 C++，以及通过 ExecuTorch 部署到移动端。PyTorch 生态仍在快速发展，像 torch.compile 和 torch.export 这样的工具，也让部署这件事变得越来越可及。

虽然 PyTorch 的某些部署能力还在不断成熟之中，但其底层基础技术已经足够强大，而且还在持续增强。现在，你已经拥有了一套很扎实的框架，可以把自己的模型部署到它真正被需要的地方——无论是作为 Web 服务、作为 C++ 应用的一部分、跑在移动设备上，还是部署到专用硬件上。

我们希望，这本书已经实现了它最初的承诺：让你真正掌握深度学习基础，并熟悉 PyTorch 这一套工具链。也希望你在阅读这本书时，获得了和我们写作时一样多的乐趣。（其实希望你更开心一点——毕竟写书真的很累！）

17.8 练习

随着《Deep Learning with PyTorch》即将收尾，我们留给你最后一个练习：

挑一个真正让你感到兴奋的项目开始动手做。Kaggle 是一个非常好的起点。去做吧。

你已经具备了成功所需要的技能，也掌握了必要的工具。我们真的很期待看到你接下来会做出什么！

小结

我们可以使用 Gradio 为 PyTorch 模型快速创建交互式 demo，它为展示模型提供了一种非常简单的方式，而不需要大量 Web 开发知识。
FastAPI 让我们能够构建高性能的 REST API 来服务 PyTorch 模型，并且内置了数据校验和异步支持。
请求批处理和流式响应有助于更高效地利用 GPU 资源，并为诸如文本生成这样的任务提供更好的用户体验。
torch.export 允许我们在 PyTorch 生态中把模型优化为更适合部署的形式，并支持 AOT 编译来提升性能。
如果想把模型部署到 PyTorch 生态之外，ONNX 提供了一种标准化格式，用于实现跨框架、跨硬件兼容。
PyTorch 模型中的量化是一种缩小模型体积并降低计算需求的技术，它通过把权重和激活从高精度格式（如 32 位浮点）转换为低精度格式（如 8 位整数）来实现。
使用 torch.compile，往往可以在无需改动代码的前提下显著提升推理性能，这得益于诸如 kernel fusion 之类的优化。
PyTorch profiler 能通过可视化执行流、操作耗时和资源使用情况，帮助我们找出性能瓶颈。
LibTorch 允许我们直接在 C++ 应用中使用 PyTorch 模型，其 API 在设计上尽量贴近 Python 接口，便于迁移。
ExecuTorch 使得在计算资源受限的移动端和嵌入式设备上高效部署 PyTorch 模型成为可能。

PyTorch 深度学习——部署到生产环境