NVlabs / neural_rx

Real-Time Inference of 5G NR Multi-user MIMO Neural Receivers
Other
32 stars 6 forks source link

How do you handle multiple streams if the receiver is an AI receiver? #1

Open modestlyh opened 2 weeks ago

modestlyh commented 2 weeks ago

[batch size, num ofdm symbols, num subcarriers, num_bits_per_symbol]

        llr = self._neural_receiver([y, no])
        # [batch_size, num_rx, num_streams_per_rx, num_ofdm_symbols, fft_size, num_bits_per_symbol]
        llr = insert_dims(llr, 2, 1)
        # [batch_size, num_rx, num_streams_per_rx, num_data_symbols, num_bits_per_symbol]
        llr = self._rg_demapper(llr)
        # [batch_size, num_tx, 1, num_data_symbols*num_bit_per_symbols]
        llr = flatten_last_dims(llr, 2)
        # Remove stream dimension, NOTE: does not support
        # multiple-streams
        # per user; conceptually the neural receiver does, but would
        # require modified reshapes
        llr = tf.squeeze(llr, axis=-2)

How to change the code if num_streams_per_rx is not equal to 1

SebastianCa commented 2 weeks ago

Hi @modestlyh, the easiest option is to define multiple independent users and treat each stream as an individual user. This boils down to defining multiple active DMRS ports in the config file.

modestlyh commented 2 weeks ago

@SebastianCa , Thanks for your answer,it's a good idea, but I would like to consider how to implement it under single user, please tell me some ideas and I will try it out!

SebastianCa commented 2 weeks ago

On the transmitter side, you would need to configure a proper stream managment meaning that the 5G NR / DMRS config must be adjusted such that the single user sends multiple streams assigned to diffferent DMRS ports (see DMRS configuration).

On the receiver side, the intitial LS channel estimation needs to be adjusted. Further, the CGNNOFDM class must be modified such that it handles the additional stream-dimension (e.g., positional encoding must be modified). For the NRX core, you can either reshape the input (as suggested above), or you would need to add the stream_per_tx dimension to the network architecture.

modestlyh commented 2 weeks ago

OK,thanks for your reply,I will try it!

modestlyh commented 2 weeks ago

@SebastianCa ,This is my stream management: bs_ut_association = np.zeros([NUM_BS, _pusch_transmitter.resource_grid.num_tx]) bs_ut_association[0, :] = 1 self.sm = StreamManagement(bs_ut_association, pusch_config.num_layers)

I set num_streams_per_tx = pusch_config.num_layers=2,I found that after TBEncoder, num_coded_bits becomes twice as much, can I just split num_coded_bits into two parts and see it as two single streams, I want to know if two single streams are fed into the mapper, can the AI receiver learn the two sets of mapping rules?

SebastianCa commented 2 weeks ago

As mentioned above, you need to modify the entire receiver such that it handles the stream_dimension properly (i.e., the tensor shapes must fit to the problem).

I would still suggest to configure the NRX for 2 active users/DMRS ports and avoid the need to rewrite the NRX architecture.

modestlyh commented 2 weeks ago

@SebastianCa ,Modifying the NRX core is a real pain, so I'm thinking of doing a single-user, dual-stream 2x4MIMO using the Resnet that sionna used for SIMO as the AI receiver before. the Resnet is fine for single-stream, but with dual-stream, after some modifications, it won't drop the loss at around 0.69 during training。 here is the code to modify it to dual stream:

[batch size, num ofdm symbols, num subcarriers, 2num_bits_per_symbol],The channel at the last layer of the network is 2num_bits_per_symbol

        llr = self._neural_receiver([y, no])  
        llr1, llr2 = tf.split(llr,num_or_size_splits=2, axis=-1)
        # [batch size, num ofdm symbols, num subcarriers, num_bits_per_symbol, 2]
        llr = tf.stack([llr1, llr2], axis=-1)
        # [batch_size, num_rx, num_ofdm_symbols, fft_size, num_bits_per_symbol, 2]
        llr = insert_dims(llr, 1, 1)
        llr = tf.transpose(llr, [0, 1, 5, 2, 3, 4])
        # [batch_size, num_rx, num_streams_per_rx, num_data_symbols, num_bits_per_symbol]
        llr = self._rg_demapper(llr)
        # [batch_size, num_tx, num_layers, num_data_symbols*num_bit_per_symbols]
        llr = flatten_last_dims(llr, 2)
        # [batch_size, num_tx, num_layers*num_data_symbols*num_bit_per_symbols]
        llr = flatten_last_dims(llr,2)

Resnet's inputs are y and no as before,I don't know if the reason the network doesn't work is because the input is missing the positional encoding of the pilots, or because there is a problem with the design of the network structure,BTW, I don't think single user dual stream is the same thing as dual user single stream.

SebastianCa commented 2 weeks ago

I’m not sure I fully understand the question. It seems like you’re running the Sionna single-user receiver in ‘multi-stream’ mode, which, unsurprisingly, doesn’t perform well in that configuration. The NRX extensions in this repository are specifically designed for multi-user scenarios, including features like position-encoded pilots, initial LS estimates, and the GCNN NRX architecture.

modestlyh commented 1 week ago

@SebastianCa ,OK, I thought earlier that modifying the NRX seemed like it would be a pain, so I thought I'd try the previous sionna's single-user receiver as a replacement how would it work, didn't realise it couldn't handle multiple streams. But my core problem is really still the single user transmitting dual streams, and I was hoping I could split it into two single streams feeding into different mappers, wondering if the neural network could learn different mapping rules. In terms of code implementation, my idea is to take the c I get after TBEncoder, split it into two parts c1, c2, and feed them into two mapper functions, and then splice them together and input them into the layer mapping, which I suspect would be problematic with a high probability, but then I can't think of a suitable way to implement the core problem I'm talking about.