constellationnet/train.py

79 lines
2.2 KiB
Python
Raw Normal View History

2019-12-13 20:17:57 +00:00
import constellation
from constellation import util
2019-12-13 17:11:09 +00:00
import torch
# Number of symbols to learn
order = 4
# Number of training examples in an epoch
2019-12-15 05:59:10 +00:00
epoch_size_multiple = 8
2019-12-13 17:11:09 +00:00
# Number of epochs
2019-12-15 05:59:10 +00:00
num_epochs = 5000
2019-12-13 17:11:09 +00:00
# Number of epochs to skip between every loss report
2019-12-13 20:17:57 +00:00
loss_report_epoch_skip = 500
2019-12-13 17:11:09 +00:00
2019-12-13 20:17:57 +00:00
# File in which the trained model is saved
2019-12-15 04:04:35 +00:00
output_file = 'output/constellation-order-{}.pth'.format(order)
model = constellation.ConstellationNet(
order=order,
encoder_layers_sizes=(4,),
decoder_layers_sizes=(4,),
channel_model=constellation.GaussianChannel()
)
2019-12-13 22:10:40 +00:00
# Train the model with random data
model.train()
2019-12-13 20:17:57 +00:00
print('Starting training with {} epochs\n'.format(num_epochs))
2019-12-13 17:11:09 +00:00
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters())
running_loss = 0
2019-12-15 05:59:10 +00:00
classes_ordered = torch.arange(order).repeat(epoch_size_multiple)
2019-12-13 17:11:09 +00:00
for epoch in range(num_epochs):
2019-12-15 05:59:10 +00:00
classes_dataset = classes_ordered[torch.randperm(len(classes_ordered))]
2019-12-13 20:17:57 +00:00
onehot_dataset = util.messages_to_onehot(classes_dataset, order)
2019-12-13 17:11:09 +00:00
optimizer.zero_grad()
predictions = model(onehot_dataset)
loss = criterion(predictions, classes_dataset)
loss.backward()
optimizer.step()
# Report loss
running_loss += loss.item()
if epoch % loss_report_epoch_skip == loss_report_epoch_skip - 1:
print('Epoch {}/{}'.format(epoch + 1, num_epochs))
print('Loss is {}'.format(running_loss))
running_loss = 0
2019-12-13 20:17:57 +00:00
print('\nFinished training\n')
# Print some examples of reconstruction
2019-12-13 22:10:40 +00:00
model.eval()
print('Reconstruction examples:')
print('Input vector\t\t\tOutput vector after softmax')
2019-12-13 20:17:57 +00:00
2019-12-13 22:10:40 +00:00
with torch.no_grad():
onehot_example = util.messages_to_onehot(torch.arange(0, order), order)
2019-12-13 20:17:57 +00:00
raw_output = model(onehot_example)
raw_output.required_grad = False
reconstructed_example = torch.nn.functional.softmax(raw_output, dim=1)
2019-12-13 22:10:40 +00:00
for index in range(order):
2019-12-13 20:17:57 +00:00
print('{}\t\t{}'.format(
2019-12-13 22:10:40 +00:00
onehot_example[index].tolist(),
2019-12-13 20:17:57 +00:00
'[{}]'.format(', '.join(
'{:.5f}'.format(x)
2019-12-13 22:10:40 +00:00
for x in reconstructed_example[index].tolist()
2019-12-13 20:17:57 +00:00
))
))
2019-12-13 17:11:09 +00:00
2019-12-13 20:17:57 +00:00
print('\nSaving model as {}'.format(output_file))
torch.save(model.state_dict(), output_file)