用于 PUSCH 信道估计的 AI/ML 组件#

PUSCH 信道估计可以加载 TrTEngine 模型并将其作为 PUSCH-RX 管道的一部分运行。此信道估计模式可用于 1 个 UE 和 1 个小区。

如何启用和运行 TrTEngine 模型#

按照以下步骤启用 TrTEngine 模型作为 PUSCH-RX 管道的一部分

  1. 在 cuphycontroller 顶层 YAML 中启用 TrTEng。例如，您可以在 cuphycontroller_nrSim_SCF_CG1.yaml 中使用以下设置

puschrx_chest_factory_settings_filename: /opt/nvidia/cuBB/cuPHY/examples/ch_est/chest_trt.yaml

  1. 将 test_mac.yaml 中的 testMAC 验证值更改为 1

validate_enable: 1

TrTEngine YAML 文件示例#

以下是 TrTEngine YAML 文件的示例。输入/输出张量的名称可能因 AI 模型而异。根据所选位置，路径和引擎文件名也可能不同。

输入和输出的维度显示在下面的 YAML 文件中。

---

chest_factory: trtengine

file: /opt/nvidia/cuBB/cuPHY/examples/ch_est/channel_estimator_fp16-True_tf32-True.engine

max_batch_size: 1

inputs:
  - name: z
    dataType: 0 # CUPHY_R_32F
    dims: [1638, 4, 4, 2, 2] # subcarriers, layers, antennas, symbols, real&imag

outputs:
  - name: zout
    dataType: 0 # CUPHY_R_32F
    dims: [4, 4, 3276, 2, 2] # antennas, layers, subcarriers, symbols, real&imag
...

Python 脚本示例#

torch_to_trt_chest_example.py 示例脚本采用基于 PyAerial 的模型，并将其编译为可在 cuPHY 中运行的 TrTEngine。

先决条件#

在使用此脚本之前，请确保要编译的模型接受来自 PyAerial 最小二乘信道估计器 (algo=3) 的输入，并输出与 PyAerial MMSE 估计器 (algo=1) 相同形状的估计。

  • 输入: [batch, subcarriers, layers, rx antennas, symbols, 2]

  • 输出: [batch, rx antennas, layers, subcarriers, symbols, 2]

    注意

    输出应具有比输入多 2 倍的子载波。

如何执行脚本#

  1. 将 PyTorch 或 TensorFlow 中的模型定义复制到此脚本。此示例使用 PyTorch 模型；有关使用 TensorFlow 的更多信息，请参阅下面的注释部分。

  2. 加载训练后的模型权重。

  3. 选择所需的输入类型

  • “<use_tvs> = True”：用于 cuPHY 验证的测试向量

  • “<use_tvs> = False”：使用 Sionna 生成并使用 PyAerial 估计的信道

  • 可选地，只要维度与输入和输出匹配，任何其他输入和输出都可以保存并与模型一起使用。

  1. 将模型作为参数提供给 check_results() 函数。

    注意

    不应修改此函数，因为它模拟 cuPHY 中的步骤。

  2. 将提供的模型的 MSE 与 LS 的 MSE 进行比较。这应该指示模型是否正常工作。示例模型的结果（取决于输入类型）如下

  1. 使用 TV，获得的 LS 和模型的 MSE 如下

    • LS = -7.6 dB；ML：-14.1 dB (tv = 7201)

    • LS = -7.6 dB；ML：-13.4 dB (tv = 7217)

  2. 使用 Sionna 信道，获得的 MSE 如下

    • LS = -20.0 dB；ML = -24.8 dB

  1. 稍后在脚本中，应将模型导出到 ONNX 并再次评估。结果应与之前的数字匹配。

  2. 该模型被编译为使用 polygraphy 的 TrT。结果应再次与之前获得的结果匹配。

注释#

  • 此脚本可以端到端执行而无需修改，并且提供的值应出现在相应的步骤中。

  • 应用于输入数据的任何预处理或后处理都应包含在模型内部。这限制了模型中允许的操作。

    • 有关使用 PyTorch 导出到 ONNX 时支持的操作，请参阅 PyTorch ONNX 文档。

    • 有关使用 TensorFlow 导出到 ONNX 时支持的操作，请参阅 TensorFlow ONNX 文档。

      • 另请注意，“当前的实现仅适用于不包含任何控制流或嵌入相关操作的图”，如 TensorFlow GitHub 文档中所述。

    • 有关编译为 TRT 时支持的操作，请参阅 ONNX-TensorRT 文档。

    • 如果无法调整模型以在没有给定禁止操作的情况下工作，则可能的解决方法是使用插件（例如 TensorRT DFT 插件）。

  • 批次维度用于堆叠多个用户的估计。

  • 如果您使用的是 TensorFlow 模型而不是 PyTorch，则需要确定如何将模型导出到 ONNX。一种选择是 tf2onnx，它在 PyPI 中可用。

  • MD5 校验和

    • channel_estimator.onnx: 64af9b805c9c7e7d831a93dbb4a646ad (可重复)

    • channel_estimator_fp16-True_tf32-True.engine: 2170dd84c2e64470b3f221ca6a310ef3 (不可重复)

  • 依赖信息

    • mamba install numpy pytorch torchvision torchaudio pytorch-cuda=12.4 -c pytorch -c nvidia

    • pip install ipykernel polygraphy onnx nvidia-pyindex nvidia-tensorrt

    • 确保 TRT 的容器版本与 Python 版本匹配：pip install tensorrt==10.3 onnx2torch

本页内容