Improve plot

This commit is contained in:
Mattéo Delabre 2019-12-13 17:10:40 -05:00
parent 49e63775dd
commit a01d83f339
Signed by: matteo
GPG Key ID: AE3FBD02DC583ABB
2 changed files with 87 additions and 17 deletions

77
plot.py
View File

@ -2,12 +2,15 @@ import constellation
from constellation import util from constellation import util
import torch import torch
from matplotlib import pyplot from matplotlib import pyplot
from mpl_toolkits.axisartist.axislines import SubplotZero
import math
import numpy
# Number learned symbols # Number learned symbols
order = 4 order = 4
# File in which the trained model is saved # File in which the trained model is saved
input_file = 'output/constellation-net.tc' input_file = 'output/constellation-order-{}.tc'.format(order)
model = constellation.ConstellationNet(order=order) model = constellation.ConstellationNet(order=order)
model.load_state_dict(torch.load(input_file)) model.load_state_dict(torch.load(input_file))
@ -19,8 +22,74 @@ with torch.no_grad():
torch.arange(0, order), torch.arange(0, order),
order order
) )
).tolist() )
fig = pyplot.figure()
ax = SubplotZero(fig, 111)
fig.add_subplot(ax)
# Extend axes symmetrically around zero so that they fit data
extent = max(
abs(encoded_vectors.min()),
abs(encoded_vectors.max())
) * 1.05
ax.set_xlim(-extent, extent)
ax.set_ylim(-extent, extent)
# Hide borders
for direction in ['left', 'bottom', 'right', 'top']:
ax.axis[direction].set_visible(False)
# Show zero-centered axes
for direction in ['xzero', 'yzero']:
axis = ax.axis[direction]
axis.set_visible(True)
axis.set_axisline_style("-|>")
# Configure axes ticks and labels
ax.annotate(
'I', (1, 0.5), xycoords='axes fraction',
xytext=(25, 0), textcoords='offset points',
va='center', ha='right'
)
ax.axis['xzero'].major_ticklabels.set_backgroundcolor('white')
ax.axis['xzero'].major_ticklabels.set_ha('center')
ax.axis['xzero'].major_ticklabels.set_va('top')
ax.annotate(
'Q', (0.5, 1), xycoords='axes fraction',
xytext=(0, 25), textcoords='offset points',
va='center', ha='center'
)
ax.axis['yzero'].major_ticklabels.set_rotation(-90)
ax.axis['yzero'].major_ticklabels.set_backgroundcolor('white')
ax.axis['yzero'].major_ticklabels.set_ha('left')
ax.axis['yzero'].major_ticklabels.set_va('center')
# Add a single tick on 0
ax.set_xticks(ax.get_xticks()[ax.get_xticks() != 0])
ax.set_yticks(ax.get_yticks()[ax.get_yticks() != 0])
ax.annotate(
'0', (0, 0),
xytext=(15, -10), textcoords='offset points',
va='center', ha='center'
)
ax.grid()
# Plot encoded vectors
ax.scatter(*zip(*encoded_vectors.tolist()), zorder=10)
# Add index label for each vector
for row in range(order):
ax.annotate(
row + 1, encoded_vectors[row],
xytext=(5, 5), textcoords='offset points',
backgroundcolor='white', zorder=9
)
fig, axis = pyplot.subplots()
axis.scatter(*zip(*encoded_vectors))
pyplot.show() pyplot.show()

View File

@ -15,11 +15,14 @@ num_epochs = 20000
loss_report_epoch_skip = 500 loss_report_epoch_skip = 500
# File in which the trained model is saved # File in which the trained model is saved
output_file = 'output/constellation-net.tc' output_file = 'output/constellation-order-{}.tc'.format(order)
print('Starting training with {} epochs\n'.format(num_epochs))
model = constellation.ConstellationNet(order=order) model = constellation.ConstellationNet(order=order)
# Train the model with random data
model.train()
print('Starting training with {} epochs\n'.format(num_epochs))
criterion = torch.nn.CrossEntropyLoss() criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters()) optimizer = torch.optim.Adam(model.parameters())
@ -46,24 +49,22 @@ for epoch in range(num_epochs):
print('\nFinished training\n') print('\nFinished training\n')
# Print some examples of reconstruction # Print some examples of reconstruction
with torch.no_grad(): model.eval()
num_examples = 5 print('Reconstruction examples:')
print('Input vector\t\t\tOutput vector after softmax')
classes_example = util.get_random_messages(num_examples, order) with torch.no_grad():
onehot_example = util.messages_to_onehot(classes_example, order) onehot_example = util.messages_to_onehot(torch.arange(0, order), order)
raw_output = model(onehot_example) raw_output = model(onehot_example)
raw_output.required_grad = False raw_output.required_grad = False
reconstructed_example = torch.nn.functional.softmax(raw_output, dim=1) reconstructed_example = torch.nn.functional.softmax(raw_output, dim=1)
print('Reconstruction examples:') for index in range(order):
print('Input vector\t\t\tOutput vector after softmax')
for example_index in range(num_examples):
print('{}\t\t{}'.format( print('{}\t\t{}'.format(
onehot_example[example_index].tolist(), onehot_example[index].tolist(),
'[{}]'.format(', '.join( '[{}]'.format(', '.join(
'{:.5f}'.format(x) '{:.5f}'.format(x)
for x in reconstructed_example[example_index].tolist() for x in reconstructed_example[index].tolist()
)) ))
)) ))