import torch.nn as nn import torch from .GaussianChannel import GaussianChannel from .NormalizePower import NormalizePower from . import util class ConstellationNet(nn.Module): """ Autoencoder network to automatically shape a constellation of symbols for efficient communication over an optical fiber channel. """ def __init__( self, order=2, encoder_layers_sizes=(), decoder_layers_sizes=(), channel_model=GaussianChannel() ): """ Create an autoencoder. :param order: Order of the constellation, i.e. the number of messages that are to be transmitted, or equivalently the number of symbols whose placements in the constellation have to be learned. :param encoder_layers_sizes: Shape of the encoder’s hidden layers. The size of this sequence is the number of hidden layers, with each element being a number which specifies the number of neurons in its channel. :param decoder_layers_sizes: Shape of the decoder’s hidden layers. Uses the same convention as `encoder_layers_sizes` above. :param channel_model: Instance of the channel model to use between the encoder and decoder network. """ super().__init__() self.order = order # Build the encoder network taking a one-hot encoded message as input # and outputting an I/Q vector. The network additionally uses hidden # layers as specified in `encoder_layers_sizes` prev_layer_size = order encoder_layers = [] for layer_size in encoder_layers_sizes: encoder_layers.append(nn.Linear(prev_layer_size, layer_size)) encoder_layers.append(nn.ReLU()) encoder_layers.append(nn.BatchNorm1d(layer_size)) prev_layer_size = layer_size encoder_layers += [ nn.Linear(prev_layer_size, 2), NormalizePower(), ] self.encoder = nn.Sequential(*encoder_layers) self.channel = channel_model # Build the decoder network taking the noisy I/Q vector received from # the channel as input and outputting a probability vector for each # original message. The network additionally uses hidden layers as # specified in `decoder_layers_sizes` prev_layer_size = 2 decoder_layers = [] for layer_size in decoder_layers_sizes: decoder_layers.append(nn.Linear(prev_layer_size, layer_size)) encoder_layers.append(nn.ReLU()) decoder_layers.append(nn.BatchNorm1d(layer_size)) prev_layer_size = layer_size # Softmax is not used at the end of the network because the # CrossEntropyLoss criterion is used for training, which includes # LogSoftmax decoder_layers.append(nn.Linear(prev_layer_size, order)) self.decoder = nn.Sequential(*decoder_layers) def forward(self, x): """ Perform encoding and decoding of an input vector and compute its reconstructed vector. :param x: Original one-hot encoded data. :return: Reconstructed vector. """ symbol = self.encoder(x) noisy_symbol = self.channel(symbol) return self.decoder(noisy_symbol) def get_constellation(self): """ Extract symbol constellation out of the trained encoder. :return: Matrix containing `order` rows with the nᵗʰ one being the I/Q vector that is the result of encoding the nᵗʰ message. """ with torch.no_grad(): return self.encoder( util.messages_to_onehot( torch.arange(0, self.order), self.order ) )