AI/ML Components for PUSCH Channel Estimation#

PUSCH Channel Estimation can load the TrTEngine model and run it as part of the PUSCH-RX pipeline. This channel estimation mode can work for 1UE and 1 Cell.

How to Enable and Run the TrTEngine Model#

Follow these steps to enable the TrTEngine model as part of the PUSCH-RX pipeline:

  1. Enable TrTEng in the cuphycontroller top-level YAML. For example, you use the following setting in cuphycontroller_nrSim_SCF_CG1.yaml:

puschrx_chest_factory_settings_filename: /opt/nvidia/cuBB/cuPHY/examples/ch_est/chest_trt.yaml
  1. Change the testMAC validation value to 1 in test_mac.yaml:

validate_enable: 1

TrTEngine YAML File Example#

The following is an example of the TrTEngine YAML file. The names of the input/output tensors might differ depending on the AI model. Depending on the chosen location, the path and engine filename might also be different.

The dimensions of the input and outputs are shown in the YAML file below.

---

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
...

Example Python Script#

The torch_to_trt_chest_example.py example script takes a PyAerial-based model and compiles it into a TrTEngine that can run in cuPHY.

Prerequisites#

Before using this script, ensure the model to be compiled accepts inputs from the PyAerial Least Squares Channel Estimator (algo=3) and outputs an estimate with the same shape as the PyAerial MMSE Estimator (algo=1).

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

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

    Note

    The output should have 2x more subcarriers than the input.

How to Execute the Script#

  1. Copy the model definition in PyTorch or TensorFlow to this script. This example uses a PyTorch model; refer to the Notes section below for more information on using TensorFlow.

  2. Load the trained model weights.

  3. Select the desired input type:

  • “<use_tvs> = True”: Test Vectors for cuPHY validation

  • “<use_tvs> = False”: A channel generated with Sionna and estimated with PyAerial

  • Optionally, any other input and output can be saved and used with the model as long as the dimensions match the input and output.

  1. Provide the model as an argument to the check_results() function.

    Note

    This function should not be modified because it emulates the steps in cuPHY.

  2. Compare the MSE of the model provided with the MSE of LS. This should give an indication of whether the model is working properly. The results for the example model (depending on the input type) are the following:

  1. Using TVs, the MSEs obtained for LS and the model are as follows:

    • LS = -7.6 dB; ML: -14.1 dB (tv = 7201)

    • LS = -7.6 dB; ML: -13.4 dB (tv = 7217)

  2. Using Sionna channels, the MSEs obtained are as follows:

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

  1. Later in the script, the model should be exported to ONNX and evaluated again. Results should match the previous numbers.

  2. The model is compiled to TrT using polygraphy. The results should again match the ones previously obtained.

Notes#

  • This script can be executed end-to-end without modification, and the provided values should appear in the respective steps.

  • Any pre- or post-processing applied to the input data should be included inside the model. This limits the operations allowed in the model.

    • Refer to the PyTorch ONNX documentation for supported operations when exporting to ONNX with PyTorch.

    • Refer to the TensorFlow ONNX documentation for supported operations when exporting to ONNX with TensorFlow.

      • Note further that “The current implementation only works for graphs that do not contain any control flow or embedding related ops,” as described in the TensorFlow GitHub documentation.

    • Refer to the ONNX-TensorRT documentation for supported operations when compiling to TRT.

    • If the model cannot be adjusted to work without a given forbidden operation, then a possible workaround is to use a plugin (e.g. a TensorRT DFT plugin).

  • The batch dimension is used to stack the estimations of multiple users.

  • If you’re using a TensorFlow model instead of PyTorch, you will need to determine how to export the model to ONNX. One option is tf2onnx, which is available in PyPI.

  • MD5 sums:

    • channel_estimator.onnx: 64af9b805c9c7e7d831a93dbb4a646ad (repeatable)

    • channel_estimator_fp16-True_tf32-True.engine: 2170dd84c2e64470b3f221ca6a310ef3 (not repeatable)

  • Dependencies information:

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

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

    • Ensure that the container version of TRT matches the Python version: pip install tensorrt==10.3 onnx2torch