LLRNet: Model training and testing#

The wireless ML design flow using Aerial is depicted in the figure below.

ML design flow

In this notebook, we use the generated LLRNet data for training and validating LLRNet as part of the PUSCH receiver chain, implemented using pyAerial, with Aerial SDK cuPHY library working as the backend. The LLRNet is plugged in the PUSCH receiver chain in place of the conventional soft demapper. So this notebook works as an example of using pyAerial for model validation.

Finally, the model is exported into a format consumed by the TensorRT inference engine that is used for integrating the model into Aerial SDK for testing the model with real hardware in an over the air environment.

Note 1: This notebook requires that the Aerial test vectors have been generated. The test vector directory is set below in AERIAL_TEST_VECTOR_DIR variable. Note 2: This notebook also requires that the notebook example on LLRNet dataset generation has been run first.

Imports#

[1]:
%matplotlib widget
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
os.environ["CUDA_MODULE_LOADING"] = "LAZY"

import cupy as cp
import h5py as h5
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
from tqdm.auto import tqdm

# Enable TensorFloat32 for faster float32 matmul
torch.set_float32_matmul_precision('high')
import onnx
from IPython.display import Markdown
from IPython.display import display

# PyAerial components
from aerial.phy5g.algorithms import ChannelEstimator
from aerial.phy5g.algorithms import ChannelEqualizer
from aerial.phy5g.algorithms import NoiseIntfEstimator
from aerial.phy5g.algorithms import Demapper
from aerial.phy5g.algorithms import TrtEngine
from aerial.phy5g.algorithms import TrtTensorPrms
from aerial.phy5g.ldpc import LdpcDeRateMatch
from aerial.phy5g.ldpc import LdpcDecoder
from aerial.phy5g.ldpc import CrcChecker
from aerial.phy5g.config import PuschConfig
from aerial.phy5g.config import PuschUeConfig
from aerial.util.cuda import CudaStream
from aerial.util.data import load_pickle
from aerial.util.fapi import dmrs_fapi_to_bit_array
[2]:
tb_errors = dict(aerial=dict(), llrnet=dict(), logmap=dict())
tb_count = dict(aerial=dict(), llrnet=dict(), logmap=dict())
[3]:
# Dataset root directory.
DATA_DIR = "data/"

# Aerial test vector directory.
AERIAL_TEST_VECTOR_DIR = "/mnt/cicd_tvs/develop/GPU_test_input/"

# LLRNet dataset directory.
dataset_dir = DATA_DIR + "example_llrnet_dataset/QPSK/"

# LLRNet model target path
llrnet_onnx_file = f"../models/llrnet.onnx"
llrnet_trt_file = f"../models/llrnet.trt"

# Training vs. testing SNR. Assume these exist in the dataset.
train_snr = list(np.arange(-7.5, 2.5, 1.0))
test_snr = list(np.arange(-7.5, 2.5, 1.0))

# Training, validation and test split in percentages if the same SNR is used for
# training and testing.
train_split = 15
val_split = 5
test_split = 80

# Training hyperparameters.
batch_size = 512
epochs = 5
learning_rate = 5e-4
weight_decay = 1e-4
# Learning rate schedule milestones (in steps)
lr_milestones = [21875, 28125]
lr_gamma = 0.2  # Multiply LR by this factor at each milestone

# Modulation order. LLRNet needs to be trained separately for each modulation order.
mod_order = 2

Define the LLRNet model#

The LLRNet model follows the original paper

  1. Shental, J. Hoydis, “’Machine LLRning’: Learning to Softly Demodulate”, https://arxiv.org/abs/1907.01512

and is a very simple MLP model with a single hidden layer. It takes the equalized symbols in its input with the real and imaginary parts separated, and outputs soft bits (log-likelihood ratios) that can be further fed into LDPC (de)rate matching and decoding.

[4]:
class LLRNet(nn.Module):
    """LLRNet model for soft demapping.

    Simple MLP with a single hidden layer that takes equalized symbols
    (real and imaginary parts separated) and outputs soft bits (LLRs).
    """
    def __init__(self, mod_order=2):
        super().__init__()
        self.fc1 = nn.Linear(2, 16)
        self.fc2 = nn.Linear(16, 8)
        self.mod_order = mod_order

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model = LLRNet(mod_order=mod_order).cuda()

def loss_fn(predictions, llr, mod_order):
    """MSE loss between predicted and target LLRs."""
    mse = torch.mean((predictions[:, :mod_order] - llr) ** 2)
    return mse

Training, validation and testing datasets#

Here, the dataset gets loaded and spåålit into training, validation and testing datasets, as well as put in the right format for the model.

[5]:
# Load the main data file
try:
    df = pd.read_parquet(dataset_dir + "l2_metadata.parquet", engine="pyarrow")
except FileNotFoundError:
    display(Markdown("**Data not found - has llrnet_dataset_generation.ipynb been run?**"))

# Query the entries for the selected modulation order.
df = df[df["qamModOrder"] == mod_order]

# Collect the dataset by SNR.
llrs = dict()
eq_syms = dict()
indices = dict()
for pusch_record in df.itertuples():
    user_data_filename = dataset_dir + pusch_record.user_data_filename
    user_data = load_pickle(user_data_filename)

    if user_data["snr"] not in llrs.keys():
        llrs[user_data["snr"]] = []
        eq_syms[user_data["snr"]] = []
        indices[user_data["snr"]] = []

    llrs[user_data["snr"]].append(user_data["map_llrs"])
    eq_syms[user_data["snr"]].append(user_data["eq_syms"])

    indices[user_data["snr"]].append(pusch_record.Index)

llr_train, llr_val = [], []
sym_train, sym_val = [], []
test_indices = []
for key in llrs.keys():
    llrs[key] = np.stack(llrs[key])
    eq_syms[key] = np.stack(eq_syms[key])

    # Randomize the order.
    permutation = np.arange(llrs[key].shape[0])
    np.random.shuffle(permutation)
    llrs[key] = llrs[key][permutation, ...]
    eq_syms[key] = eq_syms[key][permutation, ...]
    indices[key] = list(np.array(indices[key])[permutation])

    # Separate real and imaginary parts of the symbols.
    eq_syms[key] = np.stack((np.real(eq_syms[key]), np.imag(eq_syms[key])))

    num_slots = llrs[key].shape[0]
    if key in train_snr and key in test_snr:
        num_train_slots = int(np.round(train_split / 100 * num_slots))
        num_val_slots = int(np.round(val_split / 100 * num_slots))
        num_test_slots = int(np.round(test_split / 100 * num_slots))
    elif key in train_snr:
        num_train_slots = int(np.round(train_split / (train_split + val_split) * num_slots))
        num_val_slots = int(np.round(val_split / (train_split + val_split) * num_slots))
        num_test_slots = 0
    elif key in test_snr:
        num_train_slots = 0
        num_val_slots = 0
        num_test_slots = num_slots
    else:
        num_train_slots = 0
        num_val_slots = 0
        num_test_slots = 0

    # Collect training/validation/testing sets.
    llr_train.append(llrs[key][:num_train_slots, ...])
    llr_val.append(llrs[key][num_train_slots:num_train_slots+num_val_slots, ...])
    sym_train.append(eq_syms[key][:, :num_train_slots, ...])
    sym_val.append(eq_syms[key][:, num_train_slots:num_train_slots+num_val_slots, ...])
    # Just indices for the test set.
    test_indices += indices[key][num_train_slots+num_val_slots:num_train_slots+num_val_slots+num_test_slots]

llr_train = np.transpose(np.concatenate(llr_train, axis=0), (1, 0, 2))
llr_val = np.transpose(np.concatenate(llr_val, axis=0), (1, 0, 2))
sym_train = np.concatenate(sym_train, axis=1)
sym_val = np.concatenate(sym_val, axis=1)

# Fetch the total number of slots in each set.
num_train_slots = llr_train.shape[1]
num_val_slots = llr_val.shape[1]
num_test_slots = len(test_indices)

normalizer = 1.0 #np.sqrt(np.var(llr_train))
llr_train = llr_train / normalizer
llr_val = llr_val / normalizer

# Reshape into samples x mod_order array.
llr_train = llr_train.reshape(mod_order, -1).T
llr_val = llr_val.reshape(mod_order, -1).T
# Reshape into samples x 2 array.
sym_train = sym_train.reshape(2, -1).T
sym_val = sym_val.reshape(2, -1).T

print(f"Total number of slots in the training set: {num_train_slots}")
print(f"Total number of slots in the validation set: {num_val_slots}")
print(f"Total number of slots in the test set: {num_test_slots}")
Total number of slots in the training set: 3000
Total number of slots in the validation set: 1000
Total number of slots in the test set: 16000

Model training and validation#

Model training is done using PyTorch here.

[6]:
# Create PyTorch DataLoaders with performance optimizations
train_dataset = TensorDataset(
    torch.tensor(sym_train, dtype=torch.float32),
    torch.tensor(llr_train, dtype=torch.float32)
)
val_dataset = TensorDataset(
    torch.tensor(sym_val, dtype=torch.float32),
    torch.tensor(llr_val, dtype=torch.float32)
)

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True,
                          pin_memory=True, num_workers=0)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False,
                        pin_memory=True, num_workers=0)

# Compile model for faster execution (PyTorch 2.0+)
# Keep reference to original model for ONNX export (compiled models can't be traced)
model_for_export = model
if hasattr(torch, 'compile'):
    model = torch.compile(model)

# Create optimizer and learning rate scheduler
optimizer = optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=weight_decay)
scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=lr_milestones, gamma=lr_gamma)

# Training loop
for epoch in range(epochs):
    model.train()
    train_loss = 0.0
    pbar = tqdm(train_loader, desc=f'Training: Epoch {epoch+1}/{epochs}')
    for batch_x, batch_y in pbar:
        batch_x = batch_x.cuda(non_blocking=True)
        batch_y = batch_y.cuda(non_blocking=True)

        optimizer.zero_grad(set_to_none=True)  # Slightly faster than zero_grad()
        predictions = model(batch_x)
        loss = loss_fn(predictions, batch_y, mod_order)
        loss.backward()
        optimizer.step()
        scheduler.step()

        train_loss += loss.item()
        pbar.set_postfix(loss=loss.item())

    # Validation
    model.eval()
    val_loss = 0.0
    with torch.no_grad():
        for batch_x, batch_y in val_loader:
            batch_x = batch_x.cuda(non_blocking=True)
            batch_y = batch_y.cuda(non_blocking=True)
            predictions = model(batch_x)
            val_loss += loss_fn(predictions, batch_y, mod_order).item()

    avg_train_loss = train_loss / len(train_loader)
    avg_val_loss = val_loss / len(val_loader)
    print(f"Epoch {epoch+1}/{epochs} - loss: {avg_train_loss:.6f} - val_loss: {avg_val_loss:.6f}")

print("Training complete.")
Epoch 1/5 - loss: 10.292629 - val_loss: 9.153076
Epoch 2/5 - loss: 9.232538 - val_loss: 9.128629
Epoch 3/5 - loss: 9.227299 - val_loss: 9.125972
Epoch 4/5 - loss: 9.224715 - val_loss: 9.123435
Epoch 5/5 - loss: 9.222145 - val_loss: 9.120977
Training complete.

Export to TensorRT#

Finally, the model gets exported to ONNX format. The ONNX format needs to be converted to TRT engine format to be consumed by the TensorRT inference engine, this is done here using the command line tool trtexec.

[7]:
# Export PyTorch model to ONNX format
# Use model_for_export (original model) since compiled models can't be traced
model_for_export.eval()
dummy_input = torch.randn(1, 2).cuda()
import warnings
with warnings.catch_warnings():
    warnings.simplefilter("ignore", DeprecationWarning)
    torch.onnx.export(
        model_for_export,
        dummy_input,
        llrnet_onnx_file,
        input_names=["input"],
        output_names=["output"],
        dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}},
        opset_version=18,
        dynamo=False,
    )
print("ONNX model created. Converting to TRT engine file...")
command = f"trtexec " + \
    f"--onnx={llrnet_onnx_file} " + \
    f"--saveEngine={llrnet_trt_file} " + \
    f"--skipInference " + \
    f"--minShapes=input:1x2 " + \
    f"--optShapes=input:42588x2 " + \
    f"--maxShapes=input:85176x2 " + \
    f"--inputIOFormats=fp32:chw " + \
    f"--outputIOFormats=fp32:chw" + \
    f"> /dev/null"
return_val = os.system(command)
if return_val == 0:
    print("TRT engine model created.")
else:
    raise SystemExit("Failed to create the TRT engine file!")
ONNX model created. Converting to TRT engine file...
TRT engine model created.

Define a PUSCH receiver chain using pyAerial#

This class encapsulates the whole PUSCH receiver chain. The components include channel estimation, noise and interference estimation, channel equalization and soft demapping, LDPC (de)rate matching and LDPC decoding. The receiver outputs the received transport block in bytes.

The soft demapping part can be replaced by LLRNet.

[8]:
class Receiver:
    """PUSCH receiver class.

    This class encapsulates the whole PUSCH receiver chain built using
    pyAerial components.
    """

    def __init__(self,
                 llrnet_model_file,
                 num_rx_ant,
                 enable_pusch_tdi,
                 eq_coeff_algo,
                 mod_order):
        """Initialize the PUSCH receiver."""
        self._stream = CudaStream()
        self.cuda_stream = self._stream

        # Build the components of the receiver.
        self.channel_estimator = ChannelEstimator(
            num_rx_ant=num_rx_ant,
            cuda_stream=self.cuda_stream
        )
        self.channel_equalizer = ChannelEqualizer(
            num_rx_ant=num_rx_ant,
            enable_pusch_tdi=enable_pusch_tdi,
            eq_coeff_algo=eq_coeff_algo,
            cuda_stream=self.cuda_stream
        )
        self.noise_intf_estimator = NoiseIntfEstimator(
            num_rx_ant=num_rx_ant,
            eq_coeff_algo=eq_coeff_algo,
            cuda_stream=self.cuda_stream
        )
        self.demapper = Demapper(mod_order=mod_order)
        self.trt_engine = TrtEngine(
            trt_model_file=llrnet_model_file,
            max_batch_size=85176,
            input_tensors=[TrtTensorPrms('input', (2,), np.float32)],
            output_tensors=[TrtTensorPrms('output', (8,), np.float32)],
            cuda_stream=self.cuda_stream
        )

        self.derate_match = LdpcDeRateMatch(
            enable_scrambling=True,
            cuda_stream=self.cuda_stream
        )
        self.decoder = LdpcDecoder(cuda_stream=self.cuda_stream)
        self.crc_checker = CrcChecker(cuda_stream=self.cuda_stream)
        self.llr_method = "llrnet"


    def set_llr_method(self, method):
        """Set the used LLR computation method.

        Args:
            method (str): Either "aerial" meaning the conventional log-likelihood
                ratio computation, or "llrnet" for using LLRNet instead.
        """
        if method not in ["aerial", "logmap", "llrnet"]:
            raise ValueError("Invalid LLR computation method!")
        self.llr_method = method

    def run(
        self,
        rx_slot,
        slot,
        pusch_configs):
        """Run the receiver."""
        # Ensure processing happens on GPU.
        rx_slot = cp.array(rx_slot, order='F', dtype=cp.complex64)

        # Channel estimation.
        ch_est = self.channel_estimator.estimate(
            rx_slot=rx_slot,
            slot=slot,
            pusch_configs=pusch_configs
        )

        # Noise and interference estimation.
        lw_inv, noise_var_pre_eq = self.noise_intf_estimator.estimate(
            rx_slot=rx_slot,
            channel_est=ch_est,
            slot=slot,
            pusch_configs=pusch_configs
        )

        # Channel equalization and soft demapping. Note that the cuPHY kernel actually computes both
        # the equalized symbols and the LLRs.
        llr, eq_sym = self.channel_equalizer.equalize(
            rx_slot=rx_slot,
            channel_est=ch_est,
            lw_inv=lw_inv,
            noise_var_pre_eq=noise_var_pre_eq,
            pusch_configs=pusch_configs
        )

        # Use the LLRNet model here to get the log-likelihood ratios.
        dmrs_syms = pusch_configs[0].dmrs_syms
        start_sym = pusch_configs[0].start_sym
        num_symbols = pusch_configs[0].num_symbols
        num_prbs = pusch_configs[0].num_prbs
        mod_order = pusch_configs[0].ue_configs[0].mod_order
        layers = pusch_configs[0].ue_configs[0].layers
        num_data_sym = (np.array(dmrs_syms[start_sym:start_sym + num_symbols]) == 0).sum()
        if self.llr_method == "llrnet":
            # Put the input in the right format (CuPy arrays stay on GPU).
            eq_sym_input = cp.stack((cp.real(eq_sym[0]), cp.imag(eq_sym[0]))).reshape(2, -1).T
            # Run the model.
            llr_output = self.trt_engine.run({"input": eq_sym_input})["output"]

            # Reshape the output in the right format for the LDPC decoding process.
            llr_output = llr_output[:, :mod_order].T.reshape(mod_order, layers, num_prbs * 12, num_data_sym)
            llr_output *= normalizer
        elif self.llr_method == "aerial":
            llr_output = llr[0]
        elif self.llr_method == "logmap":
            inv_noise_var_lin = self.channel_equalizer.ree_diag_inv[0]
            inv_noise_var_lin = cp.transpose(inv_noise_var_lin[..., 0], (1, 2, 0)).reshape(inv_noise_var_lin.shape[1], -1)
            llr_output = self.demapper.demap(eq_sym[0][0, :, :], inv_noise_var_lin[0, :, None])
            llr_output = llr_output[:, None, :, :]

        # De-rate matching and descrambling.
        coded_blocks = self.derate_match.derate_match(
            input_llrs=[llr_output],
            pusch_configs=pusch_configs
        )

        # LDPC decoding of the derate matched blocks.
        code_blocks = self.decoder.decode(
            input_llrs=coded_blocks,
            pusch_configs=pusch_configs
        )

        # Combine the code blocks into a transport block.
        tb, _ = self.crc_checker.check_crc(
            input_bits=code_blocks,
            pusch_configs=pusch_configs
        )

        return tb[0]

Model testing on Aerial test vectors#

[9]:
if mod_order == 2:
    test_vector_filename = "TVnr_7201_PUSCH_gNB_CUPHY_s0p0.h5"
elif mod_order == 4:
    test_vector_filename = "TVnr_7916_PUSCH_gNB_CUPHY_s0p0.h5"
elif mod_order == 6:
    test_vector_filename = "TVnr_7203_PUSCH_gNB_CUPHY_s0p0.h5"
filename = AERIAL_TEST_VECTOR_DIR + test_vector_filename
input_file = h5.File(filename, "r")

num_rx_ant = input_file["gnb_pars"]["nRx"][0]
enable_pusch_tdi = input_file["gnb_pars"]["TdiMode"][0]
eq_coeff_algo = input_file["gnb_pars"]["eqCoeffAlgoIdx"][0]

receiver = Receiver(
    llrnet_trt_file,
    num_rx_ant=num_rx_ant,
    enable_pusch_tdi=enable_pusch_tdi,
    eq_coeff_algo=eq_coeff_algo,
    mod_order=mod_order
)

# Extract the test vector data and parameters.
rx_slot = np.array(input_file["DataRx"])["re"] + 1j * np.array(input_file["DataRx"])["im"]
rx_slot = rx_slot.transpose(2, 1, 0)

slot = np.array(input_file["gnb_pars"]["slotNumber"])[0]

# Wrap the parameters in a PuschConfig structure.
pusch_ue_config = PuschUeConfig(
    scid=input_file["tb_pars"]["nSCID"][0],
    layers=input_file["tb_pars"]["numLayers"][0],
    dmrs_ports=input_file["tb_pars"]["dmrsPortBmsk"][0],
    rnti=input_file["tb_pars"]["nRnti"][0],
    data_scid=input_file["tb_pars"]["dataScramId"][0],
    mcs_table=input_file["tb_pars"]["mcsTableIndex"][0],
    mcs_index=input_file["tb_pars"]["mcsIndex"][0],
    code_rate=input_file["tb_pars"]["targetCodeRate"][0],
    mod_order=input_file["tb_pars"]["qamModOrder"][0],
    tb_size=input_file["tb_pars"]["nTbByte"][0],
    rv=input_file["tb_pars"]["rv"][0],
    ndi=input_file["tb_pars"]["ndi"][0]
)
# Note that this is a list. One UE group only in this case.
pusch_configs = [PuschConfig(
    ue_configs=[pusch_ue_config],
    num_dmrs_cdm_grps_no_data=input_file["tb_pars"]["numDmrsCdmGrpsNoData"][0],
    dmrs_scrm_id=input_file["tb_pars"]["dmrsScramId"][0],
    start_prb=input_file["ueGrp_pars"]["startPrb"][0],
    num_prbs=input_file["ueGrp_pars"]["nPrb"][0],
    dmrs_syms=dmrs_fapi_to_bit_array(input_file["ueGrp_pars"]["dmrsSymLocBmsk"][0]),
    dmrs_max_len=input_file["tb_pars"]["dmrsMaxLength"][0],
    dmrs_add_ln_pos=input_file["tb_pars"]["dmrsAddlPosition"][0],
    start_sym=input_file["ueGrp_pars"]["StartSymbolIndex"][0],
    num_symbols=input_file["ueGrp_pars"]["NrOfSymbols"][0]
)]


# Run the receiver with the test vector parameters.
receiver.set_llr_method("llrnet")
tb = receiver.run(
    rx_slot=rx_slot,
    slot=slot,
    pusch_configs=pusch_configs
)

# Check that the received TB matches with the transmitted one.
tb_size = pusch_configs[0].ue_configs[0].tb_size
if np.array_equal(np.array(input_file["tb_data"])[:tb_size, 0], tb[:tb_size].get()):
    print("CRC check passed!")
else:
    print("CRC check failed!")
CRC check passed!

Model testing on synthetic/simulated data#

[10]:
# Create a receiver matching the simulated dataset parameters (num_rx_ant=2, TDI, MMSE).
sim_receiver = Receiver(
    llrnet_trt_file,
    num_rx_ant=2,
    enable_pusch_tdi=1,
    eq_coeff_algo=1,
    mod_order=mod_order
)

# Preload all test data into memory to avoid repeated disk I/O during the test loop.
test_records = list(df.take(test_indices).itertuples(index=False))
rx_slot_cache = {}
user_data_cache = {}
total_bytes = 0

for pusch_record in tqdm(test_records, desc="Preloading test data"):
    rx_fn = pusch_record.rx_iq_data_filename
    ud_fn = pusch_record.user_data_filename

    if rx_fn not in rx_slot_cache:
        rx_data = load_pickle(dataset_dir + rx_fn)
        rx_slot_cache[rx_fn] = rx_data
        total_bytes += rx_data.nbytes

    if ud_fn not in user_data_cache:
        user_data_cache[ud_fn] = load_pickle(dataset_dir + ud_fn)

print(f"Preloaded {len(rx_slot_cache)} IQ files, ~{total_bytes / 1e9:.2f} GB in memory.")

for pusch_record in tqdm(test_records, desc="Testing"):

    user_data = user_data_cache[pusch_record.user_data_filename]
    snr = user_data["snr"]

    rx_slot = rx_slot_cache[pusch_record.rx_iq_data_filename]

    ref_tb = pusch_record.macPdu
    tb_size = len(pusch_record.macPdu)
    slot = pusch_record.Slot

    # Wrap the parameters in a PuschConfig structure.
    pusch_ue_config = PuschUeConfig(
        scid=pusch_record.SCID,
        layers=pusch_record.nrOfLayers,
        dmrs_ports=pusch_record.dmrsPorts,
        rnti=pusch_record.RNTI,
        data_scid=pusch_record.dataScramblingId,
        mcs_table=pusch_record.mcsTable,
        mcs_index=pusch_record.mcsIndex,
        code_rate=pusch_record.targetCodeRate,
        mod_order=pusch_record.qamModOrder,
        tb_size=tb_size
    )
    # Note that this is a list. One UE group only in this case.
    pusch_configs = [PuschConfig(
        ue_configs=[pusch_ue_config],
        num_dmrs_cdm_grps_no_data=pusch_record.numDmrsCdmGrpsNoData,
        dmrs_scrm_id=pusch_record.ulDmrsScramblingId,
        start_prb=pusch_record.rbStart,
        num_prbs=pusch_record.rbSize,
        dmrs_syms=dmrs_fapi_to_bit_array(pusch_record.ulDmrsSymbPos),
        dmrs_max_len=1,
        dmrs_add_ln_pos=1,
        start_sym=pusch_record.StartSymbolIndex,
        num_symbols=pusch_record.NrOfSymbols
    )]

    for llr_method in ["aerial", "llrnet", "logmap"]:

        if snr not in tb_errors[llr_method].keys():
            tb_errors[llr_method][snr] = 0
            tb_count[llr_method][snr] = 0

        sim_receiver.set_llr_method(llr_method)
        tb = sim_receiver.run(
            rx_slot=rx_slot,
            slot=slot,
            pusch_configs=pusch_configs
        )

        tb_count[llr_method][snr] += 1
        tb_errors[llr_method][snr] += (not np.array_equal(tb[:tb_size].get(), ref_tb[:tb_size]))
WARNING The logger passed into createInferRuntime differs from one already registered for an existing builder, runtime, or refitter. So the current new logger is ignored, and TensorRT will use the existing one which is returned by nvinfer1::getLogger() instead.
Preloaded 16000 IQ files, ~11.74 GB in memory.
[11]:
esno_dbs = tb_count["aerial"].keys()
bler = dict(aerial=[], llrnet=[], logmap=[])
for esno_db in esno_dbs:
    bler["aerial"].append(tb_errors["aerial"][esno_db] / tb_count["aerial"][esno_db])
    bler["llrnet"].append(tb_errors["llrnet"][esno_db] / tb_count["llrnet"][esno_db])
    bler["logmap"].append(tb_errors["logmap"][esno_db] / tb_count["logmap"][esno_db])
[12]:
esno_dbs = np.array(list(esno_dbs))
fig = plt.figure(figsize=(10, 10))
plt.yscale('log')
plt.ylim(0.01, 1)
plt.xlim(np.min(esno_dbs), np.max(esno_dbs))
plt.title("BLER Performance vs. Es/No")
plt.ylabel("BLER")
plt.xlabel("Es/No [dB]")
plt.grid()
plt.plot(esno_dbs, bler["aerial"], marker="d", linestyle="-", color="blue", markersize=8)
plt.plot(esno_dbs, bler["llrnet"], marker="s", linestyle="-", color="black", markersize=8)
plt.plot(esno_dbs, bler["logmap"], marker="o", linestyle="-", color="red", markersize=8)
plt.legend(["Aerial", "LLRNet", "Log-MAP"])
[12]:
<matplotlib.legend.Legend at 0x7f5de588ee90>