import constellation from constellation import util import torch from matplotlib import pyplot import matplotlib from mpl_toolkits.axisartist.axislines import SubplotZero # Number learned symbols order = 16 # File in which the trained model is saved input_file = 'output/constellation-order-{}.pth'.format(order) # Color map used for decision regions color_map = matplotlib.cm.Dark2 # Restore model from file model = constellation.ConstellationNet( order=order, encoder_layers_sizes=(8, 4), decoder_layers_sizes=(4, 8), channel_model=constellation.GaussianChannel() ) model.load_state_dict(torch.load(input_file)) model.eval() # Setup plot fig = pyplot.figure() ax = SubplotZero(fig, 111) fig.add_subplot(ax) constellation = model.get_constellation() util.plot_constellation( ax, constellation, model.channel, model.decoder, grid_step=0.001, noise_samples=5000 ) pyplot.show()