使用 NVIDIA TensorRT for RTX Runtime 加速 Unreal Engine NNE 推理

神经网络技术正越来越多地用于计算机图形，以提升图像质量、改善性能并简化内容创作。超分辨率、降噪和神经渲染等方法有助于实时引擎更高效地工作，在兼顾性能的同时带来新的创作可能。

Unreal Engine 5（UE5）在这一方向上采取了多项举措，其中包括引入 Neural Network Engine（NNE）。NNE 作为一个抽象层，用于统一多个后端上的推理工作负载。开发者可以根据可用硬件在 GPU 上使用不同运行时，或回退到 CPU，从而在实时图形工作流中无缝集成神经网络功能。

本文介绍一个新插件，它将 NVIDIA TensorRT for RTX 添加为 NNE 的运行时选项（NNERuntimeTRT），用于在 NVIDIA RTX GPU 上进行高效推理。为展示其优势，我将使用一个简化的 UE 项目，该项目运行一个后处理 AI 模型，以突出其相较于 DirectML 等其他 GPU 运行时的收益。

首先，我们简要讨论该项目涉及的不同组件。

TensorRT for RTX 概览

TensorRT for RTX 使用户能够更高效地在 RTX GPU 上部署 AI 模型。它在运行时内使用即时（JIT）优化器，生成针对用户 GPU 定制的推理引擎。此编译过程在用户机器上执行一次，并针对其特定硬件优化模型。

因此，与默认执行提供程序相比，TensorRT for RTX 可以提供更高吞吐量。例如，在 NVIDIA GeForce RTX 5090 GPU 上测得的多个模型吞吐量对比显示，使用 TensorRT for RTX 相比 DirectML 有所提升。

TensorRT for RTX 仅兼容 NVIDIA RTX GPU，范围从 Turing 代（计算能力 7.5）到 NVIDIA Blackwell 代（计算能力 10.0）。

Unreal Engine 神经网络引擎概览

NNE 支持多种运行时，用于调用推理任务并在 CPU 和 GPU 之间进行选择。TensorRT for RTX 面向 GPU，本概览重点介绍 NNE GPU 运行时。

NNE 可以通过以下方式在 GPU 上运行推理：

• 由 CPU 同步调用，需要进行内存同步。
• 通过 Render Dependency Graph（RDG）异步执行，与帧渲染对齐。

同步方法适用于编辑器以及 LLM 等基于事件的推理任务，在这些场景中主机与设备之间的数据复制不是问题。相比之下，RDG 将模型评估与渲染资源绑定，因此非常适合 AI 后处理、升频或降噪。

NNE TensorRT for RTX 插件同时支持 GPU 和 RDG 方法，为渲染、动画、语言和语音等各类 AI 应用提供灵活性，同时在消费级设备上保持较强性能。

风格迁移后处理示例项目

我构建了一个基础 UE5 项目，用于测试 NNE TensorRT for RTX 插件，该插件在后处理期间应用风格迁移模型。测试时，我使用一些基本图元和固定摄像机搭建了一个简单关卡，以便在 DirectML 和 TensorRT 之间切换时保持视觉一致，从而更容易比较结果和性能。

先决条件

虽然该项目几乎可以直接使用，但建议具备 UE5、后处理材质和引擎源码编译经验。

要运行风格迁移推理并管理渲染资源，需要一个 NNE 实现。UE5 通过 Neural Post-Processing 插件提供这一能力，该插件通过 RDG 在 GPU 上执行推理。在本项目中，我们将使用 NNERuntimeTRT RDG 方法。

所需的只是训练或下载合适的风格模型。我们将使用 ONNX zoo 中的一个预训练模型。

该项目已包含一个导入的模型（candy-9-720.uasset）。

项目设置

尽管 NNE TensorRT for RTX 插件兼容可从启动器获取的 5.7 二进制引擎版本，但 neural post-processing 插件在其 neural profile 资产中包含硬编码的运行时列表。因此需要更新其代码，将 NNERuntimeTRT 纳入 neural profile 资产的可用运行时列表。

开始操作：

1. 从 GitHub 获取引擎源码。

如果这是你第一次访问引擎代码库，需要将你的 GitHub 账户与 Epic 账户关联。请查看此文档，其中说明了访问引擎代码库的流程：

1. 完成初始引擎设置后，你应会得到一个 Visual Studio 解决方案。打开该解决方案，并在以下路径下找到 neuralprofile.h/cpp 文件：Engine\Source\Runtime\Engine\Classes\Engine
2. 在 neuralprofile.h 中，将 NNERuntimeTRT 添加到 ENeuralProfileRuntimeType（第 59 行），如下所示：

UNNERuntimeRDGTensorRT UMETA(DisplayName = "NNERuntimeTRT"),

4. 完整的 enum class 应如下所示：

UNNERuntimeRDGTensorRT UMETA(DisplayName = "NNERuntimeTRT"), UENUM(BlueprintType)
enum class ENeuralProfileRuntimeType : uint8
{
NNERuntimeORTDml UMETA(DisplayName = "NNERuntimeORTDml"),

/** Does not have full operator support*/
NNERuntimeRDGHlsl UMETA(DisplayName = "NNERuntimeRDGHlsl"),
UNNERuntimeRDGTensorRT UMETA(DisplayName = "NNERuntimeTRT"),
MAX				 UMETA(Hidden)
};

5. 在 neuralprofile.cpp 中，找到 GetNeuralProfileRuntimeName，并将其添加到 kRuntimeNames 数组（第 28 行）：

,TEXT("NNERuntimeTRT")

6. 完整的 kRuntimeNames 数组应为：

static const TCHAR* const kRuntimeNames[] = {
	TEXT("NNERuntimeORTDml"),
	TEXT("NNERuntimeRDGHlsl"),
	TEXT("NNERuntimeTRT")
};

1. 从 Fab 或 NVIDIA Developer 页面获取插件。
2. 将插件解压到引擎插件文件夹的以下路径下：…\Engine\Plugins\Runtime\Nvidia 也可以将插件放在项目中，但在我们的场景中，将其保留在引擎插件下更简单。
3. 构建引擎。请使用这些详细说明来编译引擎。首次编译需要一些时间，因此可以考虑休息一下或喝杯咖啡。
4. 开始使用示例项目。克隆示例项目仓库。使用已编译的引擎加载项目，并运行测试关卡（LVL_PPStyleTest）。

性能分析

在运行测试关卡时，启用引擎统计信息，并在 TensorRT for RTX（TRT）和 DirectML（DML）之间切换，以评估性能提升。若要进行更全面的分析，可使用 Unreal Insights 捕获 DML 和 TRT 的帧信息及分解数据。

在配备 NVIDIA GeForce RTX 5090 GPU、1080p 分辨率的系统上，在 Unreal Insights 中，DML 需要 5.7 毫秒，而 TRT 用 3.8 毫秒完成，性能提升为 1.5 倍。

使用其他风格迁移模型

你可以使用 ONNX zoo 中的任何风格迁移模型，也可以自行训练。请注意，ONNX zoo 模型的输入和输出张量固定维度均为 1x3x224x224。neural process 插件支持对小模型进行平铺，以适配更大的帧缓冲区。虽然这能提供视觉上可接受的结果，但并不推荐，因为它会在每帧生成多个推理任务，导致 NVIDIA CUDA 与图形之间频繁发生上下文切换。为避免多次上下文切换带来的额外开销，可将模型维度改为 1x3x720x720，以便在不平铺的情况下运行推理，同时保持良好的视觉质量。

在示例项目仓库中，我包含了一个 Python 脚本，用于调整风格迁移 ONNX 模型的输入和输出张量维度。

有关使用 NNE 构建你自己的 AI 应用的更多信息，请参阅官方文档。