constellationnet/train.py

import constellation
from constellation import util
import torch
from matplotlib import pyplot
from mpl_toolkits.axisartist.axislines import SubplotZero

torch.manual_seed(42)

# Number of symbols to learn
order = 16

# Number of batches to skip between every loss report
loss_report_batch_skip = 50

# Size of batches
batch_size = 32

# File in which the trained model is saved
output_file = 'output/constellation-order-{}.pth'.format(order)

###

# Setup plot for showing training progress
fig = pyplot.figure()
ax = SubplotZero(fig, 111)
fig.add_subplot(ax)

pyplot.show(block=False)

# Train the model with random data
model = constellation.ConstellationNet(
    order=order,
    encoder_layers_sizes=(8, 4),
    decoder_layers_sizes=(4, 8),
    channel_model=constellation.GaussianChannel()
)

print('Starting training\n')

# Current batch index
batch = 0

# Accumulated loss for last batches
running_loss = 0

# True in the first training phase where small batches are used, and false in
# the second phase where point positions are refined using large batches
is_coarse_optim = True

# List of training examples (not shuffled)
classes_ordered = torch.arange(order).repeat(batch_size)

# Constellation from the previous training batch
prev_constellation = model.get_constellation()
total_change = float('inf')

# Optimizer settings
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
    optimizer, verbose=True,
    factor=0.25,
    patience=100,
    cooldown=50,
    threshold=1e-8
)

while total_change >= 1e-4:
    # Shuffle training data and convert to one-hot encoding
    classes_dataset = classes_ordered[torch.randperm(len(classes_ordered))]
    onehot_dataset = util.messages_to_onehot(classes_dataset, order)

    # Perform training step for current batch
    model.train()
    optimizer.zero_grad()
    predictions = model(onehot_dataset)
    loss = criterion(predictions, classes_dataset)
    loss.backward()
    optimizer.step()

    # Update learning rate scheduler
    scheduler.step(loss)

    # Check for convergence
    model.eval()
    constellation = model.get_constellation()
    total_change = (constellation - prev_constellation).abs().sum()
    prev_constellation = constellation

    # Report loss
    running_loss += loss.item()

    if batch % loss_report_batch_skip == loss_report_batch_skip - 1:
        print('Batch #{}'.format(batch + 1))
        print('\tLoss is {}'.format(running_loss / loss_report_batch_skip))
        print('\tChange is {}\n'.format(total_change))

        running_loss = 0

    # Update figure with current encoding
    ax.clear()
    util.plot_constellation(
        ax, constellation,
        model.channel, model.decoder,
        noise_samples=0
    )
    fig.canvas.draw()
    fig.canvas.flush_events()

    batch += 1

print('\nFinished training\n')

# Print some examples of reconstruction
model.eval()
print('Reconstruction examples:')
print('Input vector\t\t\tOutput vector after softmax')

with torch.no_grad():
    onehot_example = util.messages_to_onehot(torch.arange(0, order), order)
    raw_output = model(onehot_example)
    raw_output.required_grad = False
    reconstructed_example = torch.nn.functional.softmax(raw_output, dim=1)

    for index in range(order):
        print('{}\t\t{}'.format(
            onehot_example[index].tolist(),
            '[{}]'.format(', '.join(
                '{:.5f}'.format(x)
                for x in reconstructed_example[index].tolist()
            ))
        ))

print('\nSaving model as {}'.format(output_file))
torch.save(model.state_dict(), output_file)