PyTorch 深度学习（GPT 重译）（六）(2)https://developer.aliyun.com/article/1485254
15.1 提供 PyTorch 模型
我们将从将模型放在服务器上需要做什么开始。忠于我们的实践方法，我们将从最简单的服务器开始。一旦我们有了基本的工作内容，我们将看看它的不足之处，并尝试解决。最后，我们将看看在撰写本文时的未来。让我们创建一个监听网络的东西。¹
15.1.1 我们的模型在 Flask 服务器后面
Flask 是最广泛使用的 Python 模块之一。可以使用pip进行安装：²
pip install Flask
API 可以通过装饰函数创建。
列表 15.1 flask_hello_world.py:1
from flask import Flask app = Flask(__name__) @app.route("/hello") def hello(): return "Hello World!" if __name__ == '__main__': app.run(host='0.0.0.0', port=8000)
应用程序启动后将在端口 8000 上运行，并公开一个路由/hello，返回“Hello World”字符串。此时，我们可以通过加载先前保存的模型并通过POST路由公开它来增强我们的 Flask 服务器。我们将以第十四章的模块分类器为例。
我们将使用 Flask 的（有点奇怪地导入的）request来获取我们的数据。更准确地说，request.files 包含一个按字段名称索引的文件对象字典。我们将使用 JSON 来解析输入，并使用 flask 的jsonify助手返回一个 JSON 字符串。
现在，我们将暴露一个/predict 路由，该路由接受一个二进制块（系列的像素内容）和相关的元数据（包含一个以shape为键的字典的 JSON 对象）作为POST请求提供的输入文件，并返回一个 JSON 响应，其中包含预测的诊断。更确切地说，我们的服务器接受一个样本（而不是一批），并返回它是恶性的概率。
为了获取数据，我们首先需要将 JSON 解码为二进制，然后使用numpy.frombuffer将其解码为一维数组。我们将使用torch.from_numpy将其转换为张量，并查看其实际形状。
模型的实际处理方式就像第十四章中一样：我们将从第十四章实例化LunaModel，加载我们从训练中得到的权重，并将模型置于eval模式。由于我们不进行训练任何东西，我们会在with torch.no_grad()块中告诉 PyTorch 在运行模型时不需要梯度。
列表 15.2 flask_server.py:1
import numpy as np import sys import os import torch from flask import Flask, request, jsonify import json from p2ch13.model_cls import LunaModel app = Flask(__name__) model = LunaModel() # ❶ model.load_state_dict(torch.load(sys.argv[1], map_location='cpu')['model_state']) model.eval() def run_inference(in_tensor): with torch.no_grad(): # ❷ # LunaModel takes a batch and outputs a tuple (scores, probs) out_tensor = model(in_tensor.unsqueeze(0))[1].squeeze(0) probs = out_tensor.tolist() out = {'prob_malignant': probs[1]} return out @app.route("/predict", methods=["POST"]) # ❸ def predict(): meta = json.load(request.files['meta']) # ❹ blob = request.files['blob'].read() in_tensor = torch.from_numpy(np.frombuffer( blob, dtype=np.float32)) # ❺ in_tensor = in_tensor.view(*meta['shape']) out = run_inference(in_tensor) return jsonify(out) # ❻ if __name__ == '__main__': app.run(host='0.0.0.0', port=8000) print (sys.argv[1])
❶ 设置我们的模型，加载权重，并转换为评估模式
❷ 对我们来说没有自动求导。
❸ 我们期望在“/predict”端点进行表单提交（HTTP POST）。
❹ 我们的请求将有一个名为 meta 的文件。
❺ 将我们的数据从二进制块转换为 torch
❻ 将我们的响应内容编码为 JSON
运行服务器的方法如下：
python3 -m p3ch15.flask_server data/part2/models/cls_2019-10-19_15.48.24_final_cls.best.state
我们在 cls_client.py 中准备了一个简单的客户端，发送一个示例。从代码目录中，您可以运行它如下：
python3 p3ch15/cls_client.py
它应该告诉您结节极不可能是恶性的。显然，我们的服务器接受输入，通过我们的模型运行它们，并返回输出。那我们完成了吗？还不完全。让我们看看下一节中可以改进的地方。
15.1.2 部署的期望
让我们收集一些为提供模型服务而期望的事情。首先，我们希望支持现代协议及其特性。老式的 HTTP 是深度串行的，这意味着当客户端想要在同一连接中发送多个请求时，下一个请求只会在前一个请求得到回答后才会发送。如果您想发送一批东西，这并不是很有效。我们在这里部分交付–我们升级到 Sanic 肯定会使我们转向一个有雄心成为非常高效的框架。
在使用 GPU 时，批量请求通常比逐个处理或并行处理更有效。因此，接下来，我们的任务是从几个连接收集请求，将它们组装成一个批次在 GPU 上运行，然后将结果返回给各自的请求者。这听起来很复杂，（再次，当我们编写这篇文章时）似乎在简单的教程中并不经常做。这足以让我们在这里正确地做。但请注意，直到由模型运行持续时间引起的延迟成为问题（在等待我们自己的运行时是可以的；但在请求到达时等待正在运行的批次完成，然后等待我们的运行给出结果是禁止的），在给定时间内在一个 GPU 上运行多个批次没有太多理由。增加最大批量大小通常更有效。
我们希望并行提供几件事情。即使使用异步提供服务，我们也需要我们的模型在第二个线程上高效运行–这意味着我们希望通过我们的模型摆脱（臭名昭著的）Python 全局解释器锁（GIL）。
我们还希望尽量减少复制。无论从内存消耗还是时间的角度来看，反复复制东西都是不好的。许多 HTTP 事物都是以 Base64 编码（一种将二进制编码为更多或更少字母数字字符串的格式，每字节限制为 6 位）的形式编码的，比如，对于图像，将其解码为二进制，然后再转换为张量，然后再转换为批处理显然是相对昂贵的。我们将部分实现这一点——我们将使用流式PUT请求来避免分配 Base64 字符串，并避免通过逐渐追加到字符串来增长字符串（对于字符串和张量来说，这对性能非常糟糕）。我们说我们没有完全实现，因为我们并没有真正最小化复制。
为了提供服务，最后一个理想的事情是安全性。理想情况下，我们希望有安全的解码。我们希望防止溢出和资源耗尽。一旦我们有了固定大小的输入张量，我们应该大部分都没问题，因为从固定大小的输入开始很难使 PyTorch 崩溃。为了达到这个目标，解码图像等工作可能更令人头疼，我们不做任何保证。互联网安全是一个足够庞大的领域，我们将完全不涉及它。我们应该注意到神经网络容易受到输入操纵以生成期望但错误或意想不到的输出（称为对抗性示例），但这与我们的应用并不是非常相关，所以我们会在这里跳过它。
言归正传。让我们改进一下我们的服务器。
15.1.3 请求批处理
我们的第二个示例服务器将使用 Sanic 框架（通过同名的 Python 包安装）。这将使我们能够使用异步处理来并行处理许多请求，因此我们将在列表中勾选它。顺便说一句，我们还将实现请求批处理。
图 15.1 请求批处理的数据流
异步编程听起来可能很可怕，并且通常伴随着大量术语。但我们在这里所做的只是允许函数非阻塞地等待计算或事件的结果。
为了进行请求批处理，我们必须将请求处理与运行模型分离。图 15.1 显示了数据的流动。
在图 15.1 的顶部是客户端，发出请求。这些一个接一个地通过请求处理器的上半部分。它们导致工作项与请求信息一起入队。当已经排队了一个完整的批次或最老的请求等待了指定的最长时间时，模型运行器会从队列中取出一批，处理它，并将结果附加到工作项上。然后这些工作项一个接一个地由请求处理器的下半部分处理。
实现
我们通过编写两个函数来实现这一点。模型运行函数从头开始运行并永远运行。每当需要运行模型时，它会组装一批输入，在第二个线程中运行模型（以便其他事情可以发生），然后返回结果。
请求处理器然后解码请求，将输入加入队列，等待处理完成，并返回带有结果的输出。为了理解这里异步的含义，可以将模型运行器视为废纸篓。我们为本章所涂鸦的所有图纸都可以快速地放在桌子右侧的垃圾桶里处理掉。但是偶尔——无论是因为篮子已满还是因为到了晚上清理的时候——我们需要将所有收集的纸张拿出去扔到垃圾桶里。类似地，我们将新请求加入队列，如果需要则触发处理，并在发送结果作为请求答复之前等待结果。图 15.2 展示了我们在执行的两个函数块之前无间断执行的情况。
图 15.2 我们的异步服务器由三个模块组成：请求处理器、模型运行器和模型执行。这些模块有点像函数，但前两个在中间会让出事件循环。
相对于这个图片，一个轻微的复杂性是我们有两个需要处理事件的场合：如果我们积累了一个完整的批次，我们立即开始；当最老的请求达到最大等待时间时，我们也想运行。我们通过为后者设置一个定时器来解决这个问题。⁵
所有我们感兴趣的代码都在一个ModelRunner类中，如下列表所示。
列表 15.3 request_batching_server.py:32, ModelRunner
class ModelRunner: def __init__(self, model_name): self.model_name = model_name self.queue = [] # ❶ self.queue_lock = None # ❷ self.model = get_pretrained_model(self.model_name, map_location=device) # ❸ self.needs_processing = None # ❹ self.needs_processing_timer = None # ❺
❶ 队列
❷ 这将成为我们的锁。
❸ 加载并实例化模型。这是我们将需要更改以切换到 JIT 的（唯一）事情。目前，我们从 p3ch15/cyclegan.py 导入 CycleGAN（稍微修改为标准化为 0…1 的输入和输出）。
❹ 我们运行模型的信号
❺ 最后，定时器
ModelRunner 首先加载我们的模型并处理一些管理事务。除了模型，我们还需要一些其他要素。我们将请求输入到一个queue中。这只是一个 Python 列表，我们在后面添加工作项，然后在前面删除它们。
当我们修改queue时，我们希望防止其他任务在我们下面更改队列。为此，我们引入了一个queue_lock，它将是由asyncio模块提供的asyncio.Lock。由于我们在这里使用的所有asyncio对象都需要知道事件循环，而事件循环只有在我们初始化应用程序后才可用，因此我们在实例化时将其临时设置为None。尽管像这样锁定可能并不是绝对必要的，因为我们的方法在持有锁时不会返回事件循环，并且由于 GIL 的原因，对队列的操作是原子的，但它确实明确地编码了我们的基本假设。如果我们有多个工作进程，我们需要考虑加锁。一个警告：Python 的异步锁不是线程安全的。（叹气。）
ModelRunner 在没有任务时等待。我们需要从RequestProcessor向其发出信号，告诉它停止偷懒，开始工作。这通过名为needs_processing的asyncio.Event完成。ModelRunner使用wait()方法等待needs_processing事件。然后，RequestProcessor使用set()来发出信号，ModelRunner会被唤醒并清除事件。
最后，我们需要一个定时器来保证最大等待时间。当我们需要时，通过使用app.loop.call_at来创建此定时器。它设置needs_processing事件；我们现在只是保留一个插槽。因此，实际上，有时事件将直接被设置，因为一个批次已经完成，或者当定时器到期时。当我们在定时器到期之前处理一个批次时，我们将清除它，以便不做太多的工作。
从请求到队列
接下来，我们需要能够将请求加入队列，这是图 15.2 中RequestProcessor的第一部分的核心（不包括解码和重新编码）。我们在我们的第一个async方法process_input中完成这个操作。
列表 15.4 request_batching_server.py:54
async def process_input(self, input): our_task = {"done_event": asyncio.Event(loop=app.loop), # ❶ "input": input, "time": app.loop.time()} async with self.queue_lock: # ❷ if len(self.queue) >= MAX_QUEUE_SIZE: raise HandlingError("I'm too busy", code=503) self.queue.append(our_task) self.schedule_processing_if_needed() # ❸ await our_task["done_event"].wait() # ❹ return our_task["output"]
❶ 设置任务数据
❷ 使用锁，我们添加我们的任务和…
❸ …安排处理。处理将设置needs_processing，如果我们有一个完整的批次。如果我们没有，并且没有设置定时器，它将在最大等待时间到达时设置一个定时器。
❹ 等待（并使用 await 将控制权交还给循环）处理完成。
我们设置一个小的 Python 字典来保存我们任务的信息：当然是input，任务被排队的time，以及在任务被处理后将被设置的done_event。处理会添加一个output。
持有队列锁（方便地在async with块中完成），我们将我们的任务添加到队列中，并在需要时安排处理。作为预防措施，如果队列变得太大，我们会报错。然后，我们只需等待我们的任务被处理，并返回它。
注意 使用循环时间（通常是单调时钟）非常重要，这可能与time.time()不同。否则，我们可能会在排队之前为处理安排事件，或者根本不进行处理。
这就是我们处理请求所需的一切（除了解码和编码）。
从队列中运行批处理
接下来，让我们看一下图 15.2 右侧的model_runner函数，它执行模型调用。
列表 15.5 request_batching_server.py:71，.run_model
async def model_runner(self): self.queue_lock = asyncio.Lock(loop=app.loop) self.needs_processing = asyncio.Event(loop=app.loop) while True: await self.needs_processing.wait() # ❶ self.needs_processing.clear() if self.needs_processing_timer is not None: # ❷ self.needs_processing_timer.cancel() self.needs_processing_timer = None async with self.queue_lock: # ... line 87 to_process = self.queue[:MAX_BATCH_SIZE] # ❸ del self.queue[:len(to_process)] self.schedule_processing_if_needed() batch = torch.stack([t["input"] for t in to_process], dim=0) # we could delete inputs here... result = await app.loop.run_in_executor( None, functools.partial(self.run_model, batch) # ❹ ) for t, r in zip(to_process, result): # ❺ t["output"] = r t["done_event"].set() del to_process
❶ 等待有事情要做
❷ 如果设置了定时器，则取消定时器
❸ 获取一个批次并安排下一个批次的运行（如果需要）
❹ 在单独的线程中运行模型，将数据移动到设备，然后交给模型处理。处理完成后我们继续进行处理。
❺ 将结果添加到工作项中并设置准备事件
如图 15.2 所示，model_runner进行一些设置，然后无限循环（但在之间让出事件循环）。它在应用程序实例化时被调用，因此它可以设置我们之前讨论过的queue_lock和needs_processing事件。然后它进入循环，等待needs_processing事件。
当事件发生时，首先我们检查是否设置了时间，如果设置了，就清除它，因为我们现在要处理事情了。然后model_runner从队列中获取一个批次，如果需要的话，安排下一个批次的处理。它从各个任务中组装批次，并启动一个使用asyncio的app.loop.run_in_executor评估模型的新线程。最后，它将输出添加到任务中并设置done_event。
基本上就是这样。Web 框架–大致看起来像是带有async和await的 Flask–需要一个小包装器。我们需要在事件循环中启动model_runner函数。正如之前提到的，如果我们没有多个运行程序从队列中取出并可能相互中断，那么锁定队列就不是必要的，但是考虑到我们的代码将被适应到其他项目，我们选择保守一点，以免丢失请求。
我们通过以下方式启动我们的服务器
python3 -m p3ch15.request_batching_server data/p1ch2/horse2zebra_0.4.0.pth
现在我们可以通过上传图像数据/p1ch2/horse.jpg 进行测试并保存结果：
curl -T data/p1ch2/horse.jpg http://localhost:8000/image --output /tmp/res.jpg
请注意，这个服务器确实做了一些正确的事情–它为 GPU 批处理请求并异步运行–但我们仍然使用 Python 模式，因此 GIL 阻碍了我们在主线程中并行运行模型以响应请求。在潜在的敌对环境（如互联网）中，这是不安全的。特别是，请求数据的解码似乎既不是速度最优也不是完全安全的。
一般来说，如果我们可以进行解码，那将会更好，我们将请求流传递给一个函数，同时传递一个预分配的内存块，函数将从流中为我们解码图像。但我们不知道有哪个库是这样做的。
15.2 导出模型
到目前为止，我们已经从 Python 解释器中使用了 PyTorch。但这并不总是理想的：GIL 仍然可能阻塞我们改进的 Web 服务器。或者我们可能希望在 Python 过于昂贵或不可用的嵌入式系统上运行。这就是我们导出模型的时候。我们可以以几种方式进行操作。我们可能完全放弃 PyTorch 转向更专业的框架。或者我们可能留在 PyTorch 生态系统内部并使用 JIT，这是 PyTorch 专用 Python 子集的即时编译器。即使我们在 Python 中运行 JIT 模型，我们可能也追求其中的两个优势：有时 JIT 可以实现巧妙的优化，或者–就像我们的 Web 服务器一样–我们只是想摆脱 GIL，而 JIT 模型可以做到。最后（但我们需要一些时间才能到达那里），我们可能在libtorch下运行我们的模型，这是 PyTorch 提供的 C++ 库，或者使用衍生的 Torch Mobile。
15.2.1 与 ONNX 一起实现跨 PyTorch 的互操作性
有时，我们希望带着手头的模型离开 PyTorch 生态系统–例如，为了在具有专门模型部署流程的嵌入式硬件上运行。为此，Open Neural Network Exchange 提供了一个用于神经网络和机器学习模型的互操作格式（onnx.ai）。一旦导出，模型可以使用任何兼容 ONNX 的运行时执行，例如 ONNX Runtime，⁶前提是我们模型中使用的操作得到 ONNX 标准和目标运行时的支持。例如，在树莓派上比直接运行 PyTorch 要快得多。除了传统硬件外，许多专门的 AI 加速器硬件都支持 ONNX（onnx.ai/supported-tools .html#deployModel）。
从某种意义上说，深度学习模型是一个具有非常特定指令集的程序，由矩阵乘法、卷积、relu、tanh等粒度操作组成。因此，如果我们可以序列化计算，我们可以在另一个理解其低级操作的运行时中重新执行它。ONNX 是描述这些操作及其参数的格式的标准化。
大多数现代深度学习框架支持将它们的计算序列化为 ONNX，其中一些可以加载 ONNX 文件并执行它（尽管 PyTorch 不支持）。一些低占用量（“边缘”）设备接受 ONNX 文件作为输入，并为特定设备生成低级指令。一些云计算提供商现在可以上传 ONNX 文件并通过 REST 端点查看其暴露。
要将模型导出到 ONNX，我们需要使用虚拟输入运行模型：输入张量的值并不重要；重要的是它们具有正确的形状和类型。通过调用torch.onnx.export函数，PyTorch 将跟踪模型执行的计算，并将其序列化为一个带有提供的名称的 ONNX 文件：
torch.onnx.export(seg_model, dummy_input, "seg_model.onnx")
生成的 ONNX 文件现在可以在运行时运行，编译到边缘设备，或上传到云服务。在安装onnxruntime或onnxruntime-gpu并将batch作为 NumPy 数组获取后，可以从 Python 中使用它。
代码清单 15.6 onnx_example.py
import onnxruntime sess = onnxruntime.InferenceSession("seg_model.onnx") # ❶ input_name = sess.get_inputs()[0].name pred_onnx, = sess.run(None, {input_name: batch})
❶ ONNX 运行时 API 使用会话来定义模型，然后使用一组命名输入调用运行方法。这在处理静态图中定义的计算时是一种典型的设置。
并非所有 TorchScript 运算符都可以表示为标准化的 ONNX 运算符。如果导出与 ONNX 不兼容的操作，当我们尝试使用运行时时，将会出现有关未知aten运算符的错误。
15.2.2 PyTorch 自己的导出：跟踪
当互操作性不是关键，但我们需要摆脱 Python GIL 或以其他方式导出我们的网络时，我们可以使用 PyTorch 自己的表示，称为TorchScript 图。我们将在下一节中看到这是什么，以及生成它的 JIT 如何工作。但现在就让我们试一试。
制作 TorchScript 模型的最简单方法是对其进行跟踪。这看起来与 ONNX 导出完全相同。这并不奇怪，因为在幕后 ONNX 模型也使用了这种方法。在这里，我们只需使用torch.jit.trace函数将虚拟输入馈送到模型中。我们从第十三章导入UNetWrapper，加载训练参数，并将模型置于评估模式。
在我们追踪模型之前，有一个额外的注意事项：任何参数都不应该需要梯度，因为使用torch.no_grad()上下文管理器严格来说是一个运行时开关。即使我们在no_grad内部追踪模型，然后在外部运行，PyTorch 仍会记录梯度。如果我们提前看一眼图 15.4，我们就会明白为什么：在模型被追踪之后，我们要求 PyTorch 执行它。但是在执行记录的操作时，追踪的模型将需要梯度的参数，并且会使所有内容都需要梯度。为了避免这种情况，我们必须在torch.no_grad上下文中运行追踪的模型。为了避免这种情况–根据经验，很容易忘记然后对性能的缺乏感到惊讶–我们循环遍历模型参数并将它们全部设置为不需要梯度。
但我们只需要调用torch.jit.trace。
列出 15.7 trace_example.py
import torch from p2ch13.model_seg import UNetWrapper seg_dict = torch.load('data-unversioned/part2/models/p2ch13/seg_2019-10-20_15.57.21_none.best.state', map_location='cpu') seg_model = UNetWrapper(in_channels=8, n_classes=1, depth=4, wf=3, padding=True, batch_norm=True, up_mode='upconv') seg_model.load_state_dict(seg_dict['model_state']) seg_model.eval() for p in seg_model.parameters(): # ❶ p.requires_grad_(False) dummy_input = torch.randn(1, 8, 512, 512) traced_seg_model = torch.jit.trace(seg_model, dummy_input) # ❷
❶ 将参数设置为不需要梯度
❷ 追踪
追踪给我们一个警告：
TracerWarning: Converting a tensor to a Python index might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means the trace might not generalize to other inputs! return layer[:, :, diff_y:(diff_y + target_size[0]), diff_x:(diff_x + target_size[1])]
这源自我们在 U-Net 中进行的裁剪，但只要我们计划将大小为 512 × 512 的图像馈送到模型中，我们就没问题。在下一节中，我们将更仔细地看看是什么导致了警告，以及如何避开它突出的限制（如果需要的话）。当我们想要将比卷积网络和 U-Net 更复杂的模型转换为 TorchScript 时，这也将很重要。
我们可以保存追踪的模型
torch.jit.save(traced_seg_model, 'traced_seg_model.pt')
然后加载回来而不需要任何东西，然后我们可以调用它：
loaded_model = torch.jit.load('traced_seg_model.pt') prediction = loaded_model(batch)
PyTorch JIT 将保留我们保存模型时的状态：我们已经将其置于评估模式，并且我们的参数不需要梯度。如果我们之前没有注意到这一点，我们将需要在执行中使用with torch.no_grad():。
提示 您可以运行 JIT 编译并导出的 PyTorch 模型而不保留源代码。但是，我们总是希望建立一个工作流程，自动从源模型转换为已安装的 JIT 模型以进行部署。如果不这样做，我们将发现自己处于这样一种情况：我们想要调整模型的某些内容，但已经失去了修改和重新生成的能力。永远保留源代码，卢克！
15.2.3 带有追踪模型的服务器
现在是时候将我们的网络服务器迭代到这种情况下的最终版本了。我们可以将追踪的 CycleGAN 模型导出如下：
python3 p3ch15/cyclegan.py data/p1ch2/horse2zebra_0.4.0.pth data/p3ch15/traced_zebra_model.pt
现在我们只需要在服务器中用torch.jit.load替换对get_pretrained_model的调用（并删除现在不再需要的import get_pretrained_model）。这也意味着我们的模型独立于 GIL 运行–这正是我们希望我们的服务器在这里实现的。为了您的方便，我们已经将小的修改放在 request_batching_jit_server.py 中。我们可以用追踪的模型文件路径作为命令行参数来运行它。
现在我们已经尝试了 JIT 对我们有什么帮助，让我们深入了解细节吧！
15.3 与 PyTorch JIT 交互
在 PyTorch 1.0 中首次亮相，PyTorch JIT 处于围绕 PyTorch 的许多最新创新的中心，其中之一是提供丰富的部署选项。
15.3.1 超越经典 Python/PyTorch 时可以期待什么
经常有人说 Python 缺乏速度。虽然这有一定道理，但我们在 PyTorch 中使用的张量操作通常本身足够大，以至于它们之间的 Python 速度慢并不是一个大问题。对于像智能手机这样的小设备，Python 带来的内存开销可能更重要。因此，请记住，通常通过将 Python 排除在计算之外来加快速度的提升是 10% 或更少。
另一个不在 Python 中运行模型的即时加速仅在多线程环境中出现，但这时它可能是显著的：因为中间结果不是 Python 对象，计算不受所有 Python 并行化的威胁，即 GIL。这是我们之前考虑到的，并且当我们在服务器上使用跟踪模型时实现了这一点。
从经典的 PyTorch 执行一项操作后再查看下一项的方式转变过来，确实让 PyTorch 能够全面考虑计算：也就是说，它可以将计算作为一个整体来考虑。这为关键的优化和更高级别的转换打开了大门。其中一些主要适用于推断，而其他一些也可以在训练中提供显著的加速。
让我们通过一个快速示例来让你体会一下为什么一次查看多个操作会有益。当 PyTorch 在 GPU 上运行一系列操作时，它为每个操作调用一个子程序（在 CUDA 术语中称为内核）。每个内核从 GPU 内存中读取输入，计算结果，然后存储结果。因此，大部分时间通常不是用于计算，而是用于读取和写入内存。这可以通过仅读取一次，计算多个操作，然后在最后写入来改进。这正是 PyTorch JIT 融合器所做的。为了让你了解这是如何工作的，图 15.3 展示了长短期记忆（LSTM；en.wikipedia.org/wiki/ Long_short-term_memory）单元中进行的逐点计算，这是递归网络的流行构建块。
图 15.3 的细节对我们来说并不重要，但顶部有 5 个输入，底部有 2 个输出，中间有 7 个圆角指数表示的中间结果。通过在一个单独的 CUDA 函数中一次性计算所有这些，并将中间结果保留在寄存器中，JIT 将内存读取次数从 12 降低到 5，写入次数从 9 降低到 2。这就是 JIT 带来的巨大收益；它可以将训练 LSTM 网络的时间缩短四倍。这看似简单的技巧使得 PyTorch 能够显著缩小 LSTM 和在 PyTorch 中灵活定义的通用 LSTM 单元与像 cuDNN 这样提供的高度优化 LSTM 实现之间速度差距。
总之，使用 JIT 来避免 Python 的加速并不像我们可能天真地期望的那样大，因为我们被告知 Python 非常慢，但避免 GIL 对于多线程应用程序来说是一个重大胜利。JIT 模型的大幅加速来自 JIT 可以实现的特殊优化，但这些优化比仅仅避免 Python 开销更为复杂。
图 15.3 LSTM 单元逐点操作。从顶部的五个输入，该块计算出底部的两个输出。中间的方框是中间结果，普通的 PyTorch 会将其存储在内存中，但 JIT 融合器只会保留在寄存器中。
15.3.2 PyTorch 作为接口和后端的双重性质
要理解如何摆脱 Python 的工作原理，有益的是在头脑中将 PyTorch 分为几个部分。我们在第 1.4 节中初步看到了这一点。我们的 PyTorch torch.nn 模块–我们在第六章首次看到它们，自那以后一直是我们建模的主要工具–保存网络的参数，并使用功能接口实现：接受和返回张量的函数。这些被实现为 C++ 扩展，交给了 C++ 级别的自动求导启用层。 （然后将实际计算交给一个名为 ATen 的内部库，执行计算或依赖后端来执行，但这不重要。）
鉴于 C++ 函数已经存在，PyTorch 开发人员将它们制作成了官方 API。这就是 LibTorch 的核心，它允许我们编写几乎与其 Python 对应物相似的 C++ 张量操作。由于torch.nn模块本质上只能在 Python 中使用，C++ API 在一个名为torch::nn的命名空间中镜像它们，设计上看起来很像 Python 部分，但是独立的。
这将使我们能够在 C++ 中重新做我们在 Python 中做的事情。但这不是我们想要的：我们想要导出模型。幸运的是，PyTorch 还提供了另一个接口来访问相同的函数：PyTorch JIT。PyTorch JIT 提供了计算的“符号”表示。这个表示是TorchScript 中间表示（TorchScript IR，有时只是 TorchScript）。我们在第 15.2.2 节讨论延迟计算时提到了 TorchScript。在接下来的章节中，我们将看到如何获取我们 Python 模型的这种表示以及如何保存、加载和执行它们。与我们讨论常规 PyTorch API 时所述类似，PyTorch JIT 函数用于加载、检查和执行 TorchScript 模块也可以从 Python 和 C++ 中访问。
总结一下，我们有四种调用 PyTorch 函数的方式，如图 15.4 所示：从 C++ 和 Python 中，我们可以直接调用函数，也可以让 JIT 充当中介。所有这些最终都会调用 C++ 的 LibTorch 函数，从那里进入 ATen 和计算后端。
图 15.4 调用 PyTorch 的多种方式
15.3.3 TorchScript
TorchScript 是 PyTorch 设想的部署选项的核心。因此，值得仔细研究它的工作原理。
创建 TorchScript 模型有两种简单直接的方式：追踪和脚本化。我们将在接下来的章节中分别介绍它们。在非常高的层面上，这两种方式的工作原理如下：
在追踪中，我们在第 15.2.2 节中使用过，使用样本（随机）输入执行我们通常的 PyTorch 模型。PyTorch JIT 对每个函数都有钩子（在 C++ autograd 接口中），允许它记录计算过程。在某种程度上，这就像在说“看我如何计算输出–现在你也可以这样做。”鉴于 JIT 仅在调用 PyTorch 函数（以及nn.Module）时才起作用，你可以在追踪时运行任何 Python 代码，但 JIT 只会注意到那些部分（尤其是对控制流一无所知）。当我们使用张量形状–通常是整数元组–时，JIT 会尝试跟踪发生的情况，但可能不得不放弃。这就是在追踪 U-Net 时给我们警告的原因。
在脚本化中，PyTorch JIT 查看我们计算的实际 Python 代码，并将其编译成 TorchScript IR。这意味着，虽然我们可以确保 JIT 捕获了程序的每个方面，但我们受限于编译器理解的部分。这就像在说“我告诉你如何做–现在你也这样做。”听起来真的像编程。
我们不是来讨论理论的，所以让我们尝试使用一个非常简单的函数进行追踪和脚本化，该函数在第一维上进行低效的加法：
# In[2]: def myfn(x): y = x[0] for i in range(1, x.size(0)): y = y + x[i] return y
我们可以追踪它：
# In[3]: inp = torch.randn(5,5) traced_fn = torch.jit.trace(myfn, inp) print(traced_fn.code) # Out[3]: def myfn(x: Tensor) -> Tensor: y = torch.select(x, 0, 0) # ❶ y0 = torch.add(y, torch.select(x, 0, 1), alpha=1) # ❷ y1 = torch.add(y0, torch.select(x, 0, 2), alpha=1) y2 = torch.add(y1, torch.select(x, 0, 3), alpha=1) _0 = torch.add(y2, torch.select(x, 0, 4), alpha=1) return _0 TracerWarning: Converting a tensor to a Python index might cause the trace # ❸ to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means the trace might not generalize to other inputs!
❶ 在我们函数的第一行中进行索引
❷ 我们的循环–但完全展开并固定为 1…4，不管 x 的大小如何
❸ 令人害怕，但却如此真实！
我们看到了一个重要的警告–实际上，这段代码已经为五行修复了索引和添加，但对于四行或六行的情况并不会按预期处理。
这就是脚本化的用处所在：
# In[4]: scripted_fn = torch.jit.script(myfn) print(scripted_fn.code) # Out[4]: def myfn(x: Tensor) -> Tensor: y = torch.select(x, 0, 0) _0 = torch.__range_length(1, torch.size(x, 0), 1) # ❶ y0 = y for _1 in range(_0): # ❷ i = torch.__derive_index(_1, 1, 1) y0 = torch.add(y0, torch.select(x, 0, i), alpha=1) # ❸ return y0
❶ PyTorch 从张量大小构建范围长度。
❷ 我们的 for 循环–即使我们必须采取看起来有点奇怪的下一行来获取我们的索引 i
❸ 我们的循环体，稍微冗长一点
我们还可以打印脚本化的图，这更接近 TorchScript 的内部表示：
# In[5]: xprint(scripted_fn.graph) # end::cell_5_code[] # tag::cell_5_output[] # Out[5]: graph(%x.1 : Tensor): %10 : bool = prim::Constant[value=1]() # ❶ %2 : int = prim::Constant[value=0]() %5 : int = prim::Constant[value=1]() %y.1 : Tensor = aten::select(%x.1, %2, %2) # ❷ %7 : int = aten::size(%x.1, %2) %9 : int = aten::__range_length(%5, %7, %5) # ❸ %y : Tensor = prim::Loop(%9, %10, %y.1) # ❹ block0(%11 : int, %y.6 : Tensor): %i.1 : int = aten::__derive_index(%11, %5, %5) %18 : Tensor = aten::select(%x.1, %2, %i.1) # ❺ %y.3 : Tensor = aten::add(%y.6, %18, %5) -> (%10, %y.3) return (%y)
❶ 看起来比我们需要的要冗长得多
❷ y 的第一个赋值
❸ 在看到代码后，我们可以识别出构建范围的方法。
❹ 我们的 for 循环返回它计算的值（y）。
❺ for 循环的主体：选择一个切片，并将其添加到 y 中
在实践中，您最常使用torch.jit.script作为装饰器的形式：
@torch.jit.script def myfn(x): ...
您也可以使用自定义的trace装饰器来处理输入，但这并没有流行起来。
尽管 TorchScript（语言）看起来像 Python 的一个子集，但存在根本性差异。如果我们仔细观察，我们会发现 PyTorch 已经向代码添加了类型规范。这暗示了一个重要的区别：TorchScript 是静态类型的–程序中的每个值（变量）都有且只有一个类型。此外，这些类型限于 TorchScript IR 具有表示的类型。在程序内部，JIT 通常会自动推断类型，但我们需要用它们的类型注释脚本化函数的任何非张量参数。这与 Python 形成鲜明对比，Python 中我们可以将任何内容分配给任何变量。
到目前为止，我们已经追踪函数以获取脚本化函数。但是我们很久以前就从仅在第五章中使用函数转向使用模块了。当然，我们也可以追踪或脚本化模型。然后，这些模型将大致表现得像我们熟悉和喜爱的模块。对于追踪和脚本化，我们分别将Module的实例传递给torch.jit.trace（带有示例输入）或torch.jit.script（不带示例输入）。这将给我们带来我们习惯的forward方法。如果我们想要暴露其他方法（这仅适用于脚本化）以便从外部调用，我们在类定义中用@torch.jit.export装饰它们。
当我们说 JIT 模块的工作方式与 Python 中的工作方式相同时，这包括我们也可以用它们进行训练。另一方面，这意味着我们需要为推断设置它们（例如，使用torch.no_grad()上下文），就像我们传统的模型一样，以使它们做正确的事情。
对于算法相对简单的模型–如 CycleGAN、分类模型和基于 U-Net 的分割–我们可以像之前一样追踪模型。对于更复杂的模型，一个巧妙的特性是我们可以在构建和追踪或脚本化模块时使用来自其他脚本化或追踪代码的脚本化或追踪函数，并且我们可以在调用nn.Models时追踪函数，但是我们需要将所有参数设置为不需要梯度，因为这些参数将成为追踪模型的常数。
由于我们已经看到了追踪，让我们更详细地看一个脚本化的实际示例。
PyTorch 深度学习（GPT 重译）（六）(4)https://developer.aliyun.com/article/1485256
