匯出 TorchRL 模組

作者：Vincent Moens

注意

要在 notebook 中執行本教程，請在開頭新增一個安裝單元格，其中包含：

!pip install tensordict
!pip install torchrl
!pip install "gymnasium[atari]"

介紹

學習到的策略如果無法在真實環境中部署，其價值就微乎其微。正如其他教程所示，TorchRL 極其注重模組化和可組合性：得益於 tensordict，該庫的元件可以以最通用的方式編寫，透過將它們的簽名抽象為對輸入 TensorDict 的一系列操作。這可能會給人一種錯覺，認為該庫僅限於訓練使用，因為典型的底層執行硬體（邊緣裝置、機器人、Arduino、Raspberry Pi）不執行 Python 程式碼，更不用說安裝了 pytorch、tensordict 或 torchrl。

幸運的是，PyTorch 提供了一整套解決方案，用於將程式碼和訓練好的模型匯出到裝置和硬體，TorchRL 完全具備與這些解決方案互動的能力。可以選擇各種後端，包括本教程中示例化的 ONNX 或 AOTInductor。本教程簡要概述瞭如何將訓練好的模型隔離並打包成獨立的、可執行的檔案，以便匯出到硬體上。

主要學習內容

• 在訓練後匯出任何 TorchRL 模組；

• 使用各種後端；

• 測試匯出的模型。

快速回顧：一個簡單的 TorchRL 訓練迴圈

在本節中，我們將重現上一篇入門教程中的訓練迴圈，並稍作調整，以便與 gymnasium 庫渲染的 Atari 遊戲一起使用。我們將繼續使用 DQN 示例，並展示稍後如何使用輸出值分佈的策略。

import time
from pathlib import Path

import numpy as np

import torch

from tensordict.nn import (
    TensorDictModule as Mod,
    TensorDictSequential,
    TensorDictSequential as Seq,
)

from torch.optim import Adam

from torchrl._utils import timeit
from torchrl.collectors import SyncDataCollector
from torchrl.data import LazyTensorStorage, ReplayBuffer

from torchrl.envs import (
    Compose,
    GrayScale,
    GymEnv,
    Resize,
    set_exploration_type,
    StepCounter,
    ToTensorImage,
    TransformedEnv,
)

from torchrl.modules import ConvNet, EGreedyModule, QValueModule

from torchrl.objectives import DQNLoss, SoftUpdate

torch.manual_seed(0)

env = TransformedEnv(
    GymEnv("ALE/Pong-v5", categorical_action_encoding=True),
    Compose(
        ToTensorImage(), Resize(84, interpolation="nearest"), GrayScale(), StepCounter()
    ),
)
env.set_seed(0)

value_mlp = ConvNet.default_atari_dqn(num_actions=env.action_spec.space.n)
value_net = Mod(value_mlp, in_keys=["pixels"], out_keys=["action_value"])
policy = Seq(value_net, QValueModule(spec=env.action_spec))
exploration_module = EGreedyModule(
    env.action_spec, annealing_num_steps=100_000, eps_init=0.5
)
policy_explore = Seq(policy, exploration_module)

init_rand_steps = 5000
frames_per_batch = 100
optim_steps = 10
collector = SyncDataCollector(
    env,
    policy_explore,
    frames_per_batch=frames_per_batch,
    total_frames=-1,
    init_random_frames=init_rand_steps,
)
rb = ReplayBuffer(storage=LazyTensorStorage(100_000))

loss = DQNLoss(value_network=policy, action_space=env.action_spec, delay_value=True)
optim = Adam(loss.parameters())
updater = SoftUpdate(loss, eps=0.99)

total_count = 0
total_episodes = 0
t0 = time.time()
for data in collector:
    # Write data in replay buffer
    rb.extend(data)
    max_length = rb[:]["next", "step_count"].max()
    if len(rb) > init_rand_steps:
        # Optim loop (we do several optim steps
        # per batch collected for efficiency)
        for _ in range(optim_steps):
            sample = rb.sample(128)
            loss_vals = loss(sample)
            loss_vals["loss"].backward()
            optim.step()
            optim.zero_grad()
            # Update exploration factor
            exploration_module.step(data.numel())
            # Update target params
            updater.step()
            total_count += data.numel()
            total_episodes += data["next", "done"].sum()
    if max_length > 200:
        break

匯出基於 TensorDictModule 的策略

TensorDict 使我們能夠構建一個具有高度靈活性的策略：從一個常規的 Module（該模組從觀察值輸出動作值），我們添加了一個 QValueModule 模組，該模組讀取這些值並使用某種啟發式方法（例如，argmax 呼叫）計算動作。

然而，在我們的例子中有一個小技術難點：環境（實際的 Atari 遊戲）返回的不是灰度、84x84 的影像，而是原始的全屏彩色影像。我們附加到環境中的轉換確保模型可以讀取影像。我們可以看到，從訓練的角度來看，環境和模型之間的界限是模糊的，但在執行時，情況就清晰多了：模型應該負責將輸入資料（影像）轉換為我們的 CNN 可以處理的格式。

在這裡，tensordict 的魔力將再次幫助我們：事實證明，TorchRL 的大多數本地（非遞迴）轉換都可以用作環境轉換或 Module 例項內的預處理塊。讓我們看看如何將它們前置到我們的策略中。

policy_transform = TensorDictSequential(
    env.transform[
        :-1
    ],  # the last transform is a step counter which we don't need for preproc
    policy_explore.requires_grad_(
        False
    ),  # Using the explorative version of the policy for didactic purposes, see below.
)

我們建立一個假的輸入，並將其傳遞給 export() 和策略。這將產生一個“原始”的 Python 函式，該函式將讀取我們的輸入張量並輸出一個動作，而無需引用 TorchRL 或 tensordict 模組。

一個好的做法是呼叫 select_out_keys()，讓模型知道我們只需要一組特定的輸出（以防策略返回多個張量）。

fake_td = env.base_env.fake_tensordict()
pixels = fake_td["pixels"]
with set_exploration_type("DETERMINISTIC"):
    exported_policy = torch.export.export(
        # Select only the "action" output key
        policy_transform.select_out_keys("action"),
        args=(),
        kwargs={"pixels": pixels},
        strict=False,
    )

表示策略可能很有啟發性：我們可以看到第一個操作是 permute、div、unsqueeze、resize，然後是卷積層和 MLP 層。

print("Deterministic policy")
exported_policy.graph_module.print_readable()

作為最後的檢查，我們可以使用一個 dummy 輸入來執行策略。輸出（對於單個影像）應該是 0 到 6 之間的整數，代表要在遊戲中執行的動作。

output = exported_policy.module()(pixels=pixels)
print("Exported module output", output)

有關匯出 TensorDictModule 例項的更多詳細資訊，請參閱 tensordict 的 文件。

注意

匯出接受和輸出巢狀鍵的模組是完全可以的。相應的 kwargs 將是鍵的 “_”.join(key) 版本，即 (“group0”, “agent0”, “obs”) 鍵將對應於 “group0_agent0_obs” 關鍵字引數。衝突的鍵（例如，(“group0_agent0”, “obs”) 和 (“group0”, “agent0_obs”)）可能會導致未定義行為，應不惜一切代價避免。顯然，鍵名也應始終生成有效的關鍵字引數，即它們不應包含特殊字元，如空格或逗號。

torch.export 還有許多其他功能，我們將在下面進一步探討。在此之前，讓我們對測試時推理、以及遞迴策略的探索和隨機策略做一個小小的離題。

處理隨機策略

您可能已經注意到，上面我們使用了 set_exploration_type 上下文管理器來控制策略的行為。如果策略是隨機的（例如，策略輸出動作空間的分佈，就像 PPO 或其他類似的線上策略演算法一樣）或具有探索性（帶有附加的探索模組，如 E-Greedy、加性高斯或 Ornstein-Uhlenbeck），我們可能希望或不希望在其匯出的版本中使用該探索策略。幸運的是，匯出工具可以理解該上下文管理器，只要匯出發生在正確的上下文管理器內，策略的行為就應該與指示的一致。為了證明這一點，讓我們嘗試另一種探索型別。

with set_exploration_type("RANDOM"):
    exported_stochastic_policy = torch.export.export(
        policy_transform.select_out_keys("action"),
        args=(),
        kwargs={"pixels": pixels},
        strict=False,
    )

我們的匯出策略現在應該在呼叫堆疊的末尾有一個隨機模組，這與之前的版本不同。事實上，最後三個操作是：生成一個 0 到 6 之間的隨機整數，使用一個隨機掩碼，並根據掩碼中的值選擇網路輸出或隨機動作。

print("Stochastic policy")
exported_stochastic_policy.graph_module.print_readable()

處理遞迴策略

另一個典型用例是遞迴策略，它將輸出一個動作以及一個或多個遞迴狀態。LSTM 和 GRU 是基於 CuDNN 的模組，這意味著它們的行為與常規 Module 例項不同（匯出工具可能無法很好地跟蹤它們）。幸運的是，TorchRL 提供了這些模組的 Python 實現，可以在需要時替換 CuDNN 版本。

為了展示這一點，讓我們編寫一個依賴於 RNN 的原型策略。

from tensordict.nn import TensorDictModule
from torchrl.envs import BatchSizeTransform
from torchrl.modules import LSTMModule, MLP

lstm = LSTMModule(
    input_size=32,
    num_layers=2,
    hidden_size=256,
    in_keys=["observation", "hidden0", "hidden1"],
    out_keys=["intermediate", "hidden0", "hidden1"],
)

如果 LSTM 模組不是基於 Python 而是基於 CuDNN（LSTM），則可以使用 make_python_based() 方法來使用 Python 版本。

lstm = lstm.make_python_based()

現在讓我們建立策略。我們將兩個修改輸入形狀的層（unsqueeze/squeeze 操作）與 LSTM 和 MLP 結合起來。

recurrent_policy = TensorDictSequential(
    # Unsqueeze the first dim of all tensors to make LSTMCell happy
    BatchSizeTransform(reshape_fn=lambda x: x.unsqueeze(0)),
    lstm,
    TensorDictModule(
        MLP(in_features=256, out_features=5, num_cells=[64, 64]),
        in_keys=["intermediate"],
        out_keys=["action"],
    ),
    # Squeeze the first dim of all tensors to get the original shape back
    BatchSizeTransform(reshape_fn=lambda x: x.squeeze(0)),
)

與之前一樣，我們選擇相關的鍵。

recurrent_policy.select_out_keys("action", "hidden0", "hidden1")
print("recurrent policy input keys:", recurrent_policy.in_keys)
print("recurrent policy output keys:", recurrent_policy.out_keys)

我們現在準備匯出。為此，我們構建假的輸入並將它們傳遞給 export()。

fake_obs = torch.randn(32)
fake_hidden0 = torch.randn(2, 256)
fake_hidden1 = torch.randn(2, 256)

# Tensor indicating whether the state is the first of a sequence
fake_is_init = torch.zeros((), dtype=torch.bool)

exported_recurrent_policy = torch.export.export(
    recurrent_policy,
    args=(),
    kwargs={
        "observation": fake_obs,
        "hidden0": fake_hidden0,
        "hidden1": fake_hidden1,
        "is_init": fake_is_init,
    },
    strict=False,
)
print("Recurrent policy graph:")
exported_recurrent_policy.graph_module.print_readable()

AOTInductor：將您的策略匯出為不依賴 PyTorch 的 C++ 二進位制檔案

AOTInductor 是一個 PyTorch 模組，允許您將模型（策略或其他）匯出為不依賴 PyTorch 的 C++ 二進位制檔案。當您需要在 PyTorch 不可用的裝置或平臺上部署模型時，這尤其有用。

以下是如何使用 AOTInductor 匯出策略的示例，靈感來自 AOTI 文件。

from tempfile import TemporaryDirectory

from torch._inductor import aoti_compile_and_package, aoti_load_package

with TemporaryDirectory() as tmpdir:
    path = str(Path(tmpdir) / "model.pt2")
    with torch.no_grad():
        pkg_path = aoti_compile_and_package(
            exported_policy,
            # Specify the generated shared library path
            package_path=path,
        )
    print("pkg_path", pkg_path)

    compiled_module = aoti_load_package(pkg_path)

print(compiled_module(pixels=pixels))

使用 ONNX 匯出 TorchRL 模型

注意

要執行此指令碼部分，請確保已安裝 pytorch onnx。

!pip install onnx-pytorch
!pip install onnxruntime

您還可以透過 此處 找到更多關於在 PyTorch 生態系統中使用 ONNX 的資訊。以下示例基於此文件。

在本節中，我們將展示如何以一種可以在不依賴 PyTorch 的情況下執行模型的方式匯出模型。

網上有很多資源解釋了 ONNX 如何用於在各種硬體和裝置上部署 PyTorch 模型，包括 Raspberry Pi、NVIDIA TensorRT、iOS 和 Android。

我們訓練的 Atari 遊戲可以使用 ALE 庫 在沒有 TorchRL 或 gymnasium 的情況下進行隔離，因此為我們提供了可以使用 ONNX 實現的良好示例。

讓我們看看這個 API 是什麼樣的。

from ale_py import ALEInterface, roms

# Create the interface
ale = ALEInterface()
# Load the pong environment
ale.loadROM(roms.Pong)
ale.reset_game()

# Make a step in the simulator
action = 0
reward = ale.act(action)
screen_obs = ale.getScreenRGB()
print("Observation from ALE simulator:", type(screen_obs), screen_obs.shape)

from matplotlib import pyplot as plt

plt.tick_params(left=False, bottom=False, labelleft=False, labelbottom=False)
plt.imshow(screen_obs)
plt.title("Screen rendering of Pong game.")

匯出到 ONNX 與上面的 Export/AOTI 非常相似。

import onnxruntime

with set_exploration_type("DETERMINISTIC"):
    # We use torch.onnx.dynamo_export to capture the computation graph from our policy_explore model
    pixels = torch.as_tensor(screen_obs)
    onnx_policy_export = torch.onnx.dynamo_export(policy_transform, pixels=pixels)

我們現在可以將程式儲存在磁碟上並載入它。

with TemporaryDirectory() as tmpdir:
    onnx_file_path = str(Path(tmpdir) / "policy.onnx")
    onnx_policy_export.save(onnx_file_path)

    ort_session = onnxruntime.InferenceSession(
        onnx_file_path, providers=["CPUExecutionProvider"]
    )

onnxruntime_input = {ort_session.get_inputs()[0].name: screen_obs}
onnx_policy = ort_session.run(None, onnxruntime_input)

使用 ONNX 執行 rollout

我們現在有了一個執行我們策略的 ONNX 模型。讓我們將其與原始 TorchRL 例項進行比較：由於其輕量級，ONNX 版本應該比 TorchRL 版本更快。

def onnx_policy(screen_obs: np.ndarray) -> int:  # noqa: F811
    onnxruntime_input = {ort_session.get_inputs()[0].name: screen_obs}
    onnxruntime_outputs = ort_session.run(None, onnxruntime_input)
    action = int(onnxruntime_outputs[0])
    return action


with timeit("ONNX rollout"):
    num_steps = 1000
    ale.reset_game()
    for _ in range(num_steps):
        screen_obs = ale.getScreenRGB()
        action = onnx_policy(screen_obs)
        reward = ale.act(action)

with timeit("TorchRL version"), torch.no_grad(), set_exploration_type("DETERMINISTIC"):
    env.rollout(num_steps, policy_explore)

print(timeit.print())

請注意，ONNX 還提供了直接最佳化模型的可能性，但這超出了本教程的範圍。

結論

在本教程中，我們學習瞭如何使用各種後端匯出 TorchRL 模組，例如 PyTorch 的內建匯出功能、AOTInductor 和 ONNX。我們演示瞭如何匯出在 Atari 遊戲上訓練的模型，並使用 ALE 庫在不依賴 PyTorch 的環境中執行它。我們還比較了原始 TorchRL 例項與匯出的 ONNX 模型的效能。

關鍵要點

• 匯出 TorchRL 模組允許部署在未安裝 PyTorch 的裝置上。

• AOTInductor 和 ONNX 提供了匯出模型的替代後端。

• 最佳化 ONNX 模型可以提高效能。

進一步閱讀和學習步驟

• 有關更多資訊，請檢視 PyTorch 的 匯出功能、AOTInductor 和 ONNX 的官方文件。

• 嘗試將匯出的模型部署到不同的裝置上。

• 探索 ONNX 模型的最佳化技術以提高效能。