trainable constellation using NN instead of SGD

[GCX-art]

I want to train constellation using NN instead of SGD.

Sionna used SGD to train constellation and the neural demapper in ''End-to-end Learning with Autoencoders'', I have removed the neural demapper and then trained constellation only,but the result was not good.
So I hope to train constellation using NN instead of SGD, I have trained DNN but it not work. It is really difficult for me ,I hope to get your help , thanks you.

[SebastianCa]

Hi @GCX-art,
It’s challenging to provide detailed advice without more specifics about your experiment.

SGD is typically used as the optimizer to find optimal weights for both the neural demapper and the constellation. Are you considering replacing it with a non-gradient-based optimizer, or could you clarify what “NN” stands for in this context?

When you remove the neural demapper, did you switch to using an APP demapper instead?

Also, are you working with an AWGN channel? How does the resulting constellation appear in your setup?

[GCX-art]

Hi@SebastianCa,thanks for your response
I used an APP demapper and an AWGN channel, the resulting constellation just like "Visualizing the Learned Constellations", but the final BER comparison is not that good，here is my code:

import os
import pickle
from keras import Model
gpu_num = 0 # Use "" to use the CPU
os.environ["CUDA_VISIBLE_DEVICES"] = f"{gpu_num}"
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
# Import Sionna
import sionna as sn
# Import TensorFlow and NumPy
import tensorflow as tf
# Avoid warnings from TensorFlow
tf.get_logger().setLevel('ERROR')
import numpy as np
# For plotting
import matplotlib.pyplot as plt

#############################################################
'''parameters'''
#############################################################
BATCH_SIZE = 128 # How many examples are processed by Sionna in parallel
NUM_BITS_PER_SYMBOL = 6
# Number of iterations used for training
NUM_TRAINING_ITERATIONS = 10000
EBN0_DB_MIN = 10.0
EBN0_DB_MAX = 20.0

################################################
'''initialization'''
###############################################3
class End2EndSystem(Model): # Inherits from Keras Model

    def __init__(self, training):

        super().__init__() # Must call the Keras model initializer
        self.binary_source = sn.utils.BinarySource()
        self.constellation = sn.mapping.Constellation("qam", NUM_BITS_PER_SYMBOL, trainable=True) # Constellation is trainable
        self.mapper = sn.mapping.Mapper(constellation=self.constellation)
        self.awgn_channel = sn.channel.AWGN()
        self.demapper = sn.mapping.Demapper('app', constellation=self.constellation) # Intantiate the NeuralDemapper custom layer as any other
        self.bce = tf.keras.losses.BinaryCrossentropy(from_logits=True) # Loss function
        self.training = training

    @tf.function# Enable graph execution to speed things up
    def __call__(self, batch_size, ebno_db):

        # no channel coding used; we set coderate=1.0
        no = sn.utils.ebnodb2no(ebno_db,
                                num_bits_per_symbol=NUM_BITS_PER_SYMBOL,
                                coderate=1.0)
        bits = self.binary_source([batch_size, 1200]) # Blocklength set to 1200 bits
        x = self.mapper(bits)
        y = self.awgn_channel([x, no])
        llr = self.demapper([y, no])  # Call the NeuralDemapper custom layer as any other
        if self.training:
            loss = self.bce(bits, llr)
            return loss
        else:
            return bits, llr

#########################################################
'''SGD'''
#########################################################
# Set a seed for reproducibility
tf.random.set_seed(1)
# Instantiating the end-to-end model for training
model_train = End2EndSystem(training=True)
# Adam optimizer (SGD variant)
optimizer = tf.keras.optimizers.Adam(0.01)
for i in range(NUM_TRAINING_ITERATIONS):
    # Forward pass
    with tf.GradientTape() as tape:
        loss = model_train(BATCH_SIZE, 11.0) # The model is assumed to return the BMD rate
    # Computing and applying gradients
    grads = tape.gradient(loss, model_train.trainable_weights)
    optimizer.apply_gradients(zip(grads, model_train.trainable_weights))

#########################################################
'''constellation'''
#########################################################
fig1 = sn.mapping.Constellation("qam", NUM_BITS_PER_SYMBOL).show(labels=False)
# fig2 = model_train.constellation.show()
plt.scatter(tf.math.real(model_train.constellation.points), tf.math.imag(model_train.constellation.points),
            color='red', label='After SGD')
plt.legend(loc='upper right',fontsize=12)
plt.show()
#########################################################
'''ber'''
#########################################################
# Save the weightsin a file
weights = model_train.get_weights()
with open('trained model', 'wb') as f:
    pickle.dump(weights, f)

model = End2EndSystem(False)

model1 = End2EndSystem(False)
model1(tf.constant(1, tf.int32), tf.constant(11.0, tf.float32))
with open('trained model', 'rb') as f:
    weights = pickle.load(f)
    model1.set_weights(weights)

ber_plots = sn.utils.PlotBER("trained vs untrained")
ber_plots.simulate(model,
                  ebno_dbs=np.linspace(EBN0_DB_MIN, EBN0_DB_MAX, 20),
                  batch_size=BATCH_SIZE,
                  num_target_block_errors=100, # simulate until 100 block errors occured
                  legend="Untrained",
                  soft_estimates=True,
                  max_mc_iter=100, # run 100 Monte-Carlo simulations (each with batch_size samples)
                  show_fig=False)
ber_plots.simulate(model1,
                  ebno_dbs=np.linspace(EBN0_DB_MIN, EBN0_DB_MAX, 20),
                  batch_size=BATCH_SIZE,
                  num_target_block_errors=100,
                  legend="Trained ",
                  soft_estimates=True,
                  max_mc_iter=100,
                  show_fig=True)

plt.show()

The following are my results:

At the same time, I read the ''Trainable Communication Systems:
Concepts and Prototype'',

I want to reproduce the process，the "NN" is “neural network”，sionna have neural demapper ,but it don't have neural mapper , I want to train constellation using NN just like the neural mapper in the paper , but I failed to reproduce it ,could you give me some help?

[SebastianCa]

As a quick comment, you can find an example of a neural demapper in the tutorial notebook. Switching to a neural demapper shouldn’t pose a problem.

Regarding your code snippet: to fully realize the geometric shaping gains, you’ll need an outer FEC decoder. The performance improvements in uncoded and hard-decided cases are not guaranteed. Below are the results from a quick experiment using an outer r=0.5 LDPC decoder for 64-QAM. Note that you need to adjust the training SNR to optimize the constellation for the waterfall region of the code (in my experiment, the training SNR was set to 5 dB).

Regarding the continuum of constellations: this setup is slightly more complex. You’ll need to replace the “classical” mapper and constellation object with the SNR-dependent NeuralMapper, as described in the paper. While Sionna doesn’t currently provide an example for this, the implementation should be fairly straightforward.

[GCX-art]

Hi@SebastianCa,thanks for your advice. I aleady tried using an outer FEC decoder，it have a nice gain. And I have also tried using the CDL-C channel model to replace AWGN，its performance has also improved.

In the beginning,I was trying to see what would happen if I only trained the constellation , so I used an APP demapper to replace the neural demapper,but the result was not good,your advice is very very inspiring to me!
However, about the NeuralMapper,here is my code:

import tensorflow as tf
import os
gpu_num = 0
os.environ["CUDA_VISIBLE_DEVICES"] = f"{gpu_num}"
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

n = 1296
coderate = 0.5
k = int(n*coderate)
num_bits_per_symbol = 6
ebno_db=5.0


def r2c(x):
    num_points = 2 ** num_bits_per_symbol
    assert tf.shape(x)[0] == 2 * num_points
    real_part = x[:num_points]
    imaginary_part = x[num_points:]
    constellation_points = tf.complex(real_part, imaginary_part)
    return constellation_points


def normalization(x):
    energy = tf.reduce_mean(tf.square(tf.abs(x)))
    energy_sqrt = tf.sqrt(energy)
    x = x / energy_sqrt
    return x


def bit2onehot(x):
    x = tf.cast(x, dtype=tf.int32)
    num_bits = tf.shape(x)[-1]
    depth = 2 ** num_bits
    binary_base = 2 ** tf.range(num_bits - 1, -1, -1)
    indices = tf.reduce_sum(x * binary_base, axis=-1)
    one_hot_encoding = tf.one_hot(indices, depth)
    return one_hot_encoding


class NeuralMapper(Layer):

    def __init__(self):
        super().__init__()

        self._dense_1 = Dense(2 ** (num_bits_per_symbol + 1), 'relu')
        self._dense_2 = Dense(2 ** (num_bits_per_symbol + 1), 'linear')

    def call(self, ebno_db):
        ebno_db = expand_to_rank(ebno_db, 1)
        cons = self._dense_1(ebno_db)
        cons = self._dense_2(cons)
        cons = r2c(cons)
        cons = normalization(cons)
        s = bit2onehot(c)
        x = s*cons
        return x

I have tried to use it in the system just like the example of neural demapper,but it can't work ,could you give me some help? I think it may be the problem of code,I've tried modifying it several times, but it still doesn't work.

[SebastianCa]

Please understand that we cannot provide such a detailed support and debug individual user code. Just a view quick comments:

• c is never defined, however, it should be an input of the mapper?
• It seems like there is an issue with the input shape of the dense layers?
• the output of the dense2 layer is 2^(m+1), however the r2c function has an assertion that tests for 2^m; actuallly it should be 2*2^m such that the constellation has 2^m points after r2c
• For training, c needs to support a batch-dimension