def test_snapshoting():
CHECKPOINT_PATH = Path("./test_checkpoints")
core.init()
sender = core.GumbelSoftmaxWrapper(ToyAgent(), temperature=1)
receiver = Receiver()
loss = lambda sender_input, message, receiver_input, receiver_output, labels, aux_input: (
F.cross_entropy(receiver_output, labels),
{},
)
game = core.SymbolGameGS(sender, receiver, loss)
optimizer = torch.optim.Adam(game.parameters())
data = Dataset()
trainer = core.Trainer(
game,
optimizer,
train_data=data,
validation_data=None,
callbacks=[core.CheckpointSaver(checkpoint_path=CHECKPOINT_PATH)],
)
trainer.train(2)
assert (CHECKPOINT_PATH / Path("1.tar")).exists()
assert (CHECKPOINT_PATH / Path("2.tar")).exists()
assert (CHECKPOINT_PATH / Path("final.tar")).exists()
shutil.rmtree(CHECKPOINT_PATH) # Clean-up
"""
def test_snapshoting():
CHECKPOINT_PATH = Path('./test_checkpoints')
core.init()
sender = core.GumbelSoftmaxWrapper(ToyAgent(), temperature=1)
receiver = Receiver()
loss = lambda sender_input, message, receiver_input, receiver_output, labels: \
(F.cross_entropy(receiver_output, labels), {})
game = core.SymbolGameGS(sender, receiver, loss)
optimizer = torch.optim.Adam(game.parameters())
data = Dataset()
trainer = core.Trainer(
game,
optimizer,
train_data=data,
validation_data=None,
callbacks=[core.CheckpointSaver(checkpoint_path=CHECKPOINT_PATH)])
trainer.train(2)
assert (CHECKPOINT_PATH / Path('1.tar')).exists()
assert (CHECKPOINT_PATH / Path('2.tar')).exists()
assert (CHECKPOINT_PATH / Path('final.tar')).exists()
del trainer
trainer = core.Trainer(game, optimizer,
train_data=data) # Re-instantiate trainer
trainer.load_from_latest(CHECKPOINT_PATH)
assert trainer.start_epoch == 2
trainer.train(3)
shutil.rmtree(CHECKPOINT_PATH) # Clean-up
def test_max_snapshoting():
CHECKPOINT_PATH = Path("./test_checkpoints")
core.init()
sender = core.GumbelSoftmaxWrapper(ToyAgent(), temperature=1)
receiver = Receiver()
loss = lambda sender_input, message, receiver_input, receiver_output, labels: (
F.cross_entropy(receiver_output, labels),
{},
)
game = core.SymbolGameGS(sender, receiver, loss)
optimizer = torch.optim.Adam(game.parameters())
data = Dataset()
trainer = core.Trainer(
game,
optimizer,
train_data=data,
validation_data=None,
callbacks=[
core.CheckpointSaver(checkpoint_path=CHECKPOINT_PATH,
max_checkpoints=2)
],
)
trainer.train(n_epochs=6)
assert (CHECKPOINT_PATH / Path("5.tar")).exists()
assert (CHECKPOINT_PATH / Path("6.tar")).exists()
assert (CHECKPOINT_PATH / Path("final.tar")).exists()
assert len([x for x in CHECKPOINT_PATH.glob("**/*") if x.is_file()]) == 3
del trainer
trainer = core.Trainer(game, optimizer,
train_data=data) # Re-instantiate trainer
trainer.load_from_latest(CHECKPOINT_PATH)
assert trainer.start_epoch == 6
trainer.train(3)
shutil.rmtree(CHECKPOINT_PATH) # Clean-up
def main(params):
opts = get_params(params)
print(opts, flush=True)
device = opts.device
force_eos = opts.force_eos == 1
if opts.probs == 'uniform':
probs = np.ones(opts.n_features)
elif opts.probs == 'powerlaw':
probs = 1 / np.arange(1, opts.n_features + 1, dtype=np.float32)
else:
probs = np.array([float(x) for x in opts.probs.split(',')],
dtype=np.float32)
probs /= probs.sum()
print('the probs are: ', probs, flush=True)
train_loader = OneHotLoader(n_features=opts.n_features,
batch_size=opts.batch_size,
batches_per_epoch=opts.batches_per_epoch,
probs=probs)
# single batches with 1s on the diag
test_loader = UniformLoader(opts.n_features)
#################################
# define sender (speaker) agent #
#################################
sender = Sender(n_features=opts.n_features, n_hidden=opts.sender_hidden)
sender = RnnSenderReinforce(sender,
opts.vocab_size,
opts.sender_embedding,
opts.sender_hidden,
cell=opts.sender_cell,
max_len=opts.max_len,
num_layers=opts.sender_num_layers,
force_eos=force_eos,
noise_loc=opts.sender_noise_loc,
noise_scale=opts.sender_noise_scale)
####################################
# define receiver (listener) agent #
####################################
receiver = Receiver(n_features=opts.n_features,
n_hidden=opts.receiver_hidden)
receiver = RnnReceiverDeterministic(receiver,
opts.vocab_size,
opts.receiver_embedding,
opts.receiver_hidden,
cell=opts.receiver_cell,
num_layers=opts.receiver_num_layers,
noise_loc=opts.receiver_noise_loc,
noise_scale=opts.receiver_noise_scale)
###################
# define channel #
###################
channel = Channel(vocab_size=opts.vocab_size, p=opts.channel_repl_prob)
game = SenderReceiverRnnReinforce(
sender,
receiver,
loss,
sender_entropy_coeff=opts.sender_entropy_coeff,
receiver_entropy_coeff=opts.receiver_entropy_coeff,
length_cost=opts.length_cost,
effective_max_len=opts.effective_max_len,
channel=channel,
sender_entropy_common_ratio=opts.sender_entropy_common_ratio)
optimizer = core.build_optimizer(game.parameters())
callbacks = [
EarlyStopperAccuracy(opts.early_stopping_thr),
core.ConsoleLogger(as_json=True, print_train_loss=True)
]
if opts.checkpoint_dir:
'''
info in checkpoint_name:
- n_features as f
- vocab_size as vocab
- random_seed as rs
- lr as lr
- sender_hidden as shid
- receiver_hidden as rhid
- sender_entropy_coeff as sentr
- length_cost as reg
- max_len as max_len
- sender_noise_scale as sscl
- receiver_noise_scale as rscl
- channel_repl_prob as crp
- sender_entropy_common_ratio as scr
'''
checkpoint_name = (
f'{opts.name}' + '_aer' +
('_uniform' if opts.probs == 'uniform' else '') +
f'_f{opts.n_features}' + f'_vocab{opts.vocab_size}' +
f'_rs{opts.random_seed}' + f'_lr{opts.lr}' +
f'_shid{opts.sender_hidden}' + f'_rhid{opts.receiver_hidden}' +
f'_sentr{opts.sender_entropy_coeff}' + f'_reg{opts.length_cost}' +
f'_max_len{opts.max_len}' + f'_sscl{opts.sender_noise_scale}' +
f'_rscl{opts.receiver_noise_scale}' +
f'_crp{opts.channel_repl_prob}' +
f'_scr{opts.sender_entropy_common_ratio}')
callbacks.append(
core.CheckpointSaver(checkpoint_path=opts.checkpoint_dir,
checkpoint_freq=opts.checkpoint_freq,
prefix=checkpoint_name))
trainer = core.Trainer(game=game,
optimizer=optimizer,
train_data=train_loader,
validation_data=test_loader,
callbacks=callbacks)
trainer.train(n_epochs=opts.n_epochs)
print('<div id="prefix test without eos">')
prefix_test(trainer.game, opts.n_features, device, add_eos=False)
print('</div>')
print('<div id="prefix test with eos">')
prefix_test(trainer.game, opts.n_features, device, add_eos=True)
print('<div id="suffix test">')
suffix_test(trainer.game, opts.n_features, device)
print('</div>')
print('<div id="replacement test">')
replacement_test(trainer.game, opts.n_features, opts.vocab_size, device)
print('</div>')
print('<div id="dump">')
dump(trainer.game, opts.n_features, device, False)
print('</div>')
core.close()
def main(params):
opts = get_params(params)
print(opts, flush=True)
device = opts.device
force_eos = opts.force_eos == 1
if opts.probs == 'uniform':
probs = np.ones(opts.n_features)
elif opts.probs == 'powerlaw':
probs = 1 / np.arange(1, opts.n_features+1, dtype=np.float32)
else:
probs = np.array([float(x) for x in opts.probs.split(',')], dtype=np.float32)
probs /= probs.sum()
print('the probs are: ', probs, flush=True)
train_loader = OneHotLoader(n_features=opts.n_features, batch_size=opts.batch_size,
batches_per_epoch=opts.batches_per_epoch, probs=probs)
# single batches with 1s on the diag
test_loader = UniformLoader(opts.n_features)
if opts.sender_cell == 'transformer':
sender = Sender(n_features=opts.n_features, n_hidden=opts.sender_embedding)
sender = core.TransformerSenderReinforce(agent=sender, vocab_size=opts.vocab_size,
embed_dim=opts.sender_embedding, max_len=opts.max_len,
num_layers=opts.sender_num_layers, num_heads=opts.sender_num_heads,
hidden_size=opts.sender_hidden,
force_eos=opts.force_eos,
generate_style=opts.sender_generate_style,
causal=opts.causal_sender)
else:
sender = Sender(n_features=opts.n_features, n_hidden=opts.sender_hidden)
sender = core.RnnSenderReinforce(sender,
opts.vocab_size, opts.sender_embedding, opts.sender_hidden,
cell=opts.sender_cell, max_len=opts.max_len, num_layers=opts.sender_num_layers,
force_eos=force_eos)
if opts.receiver_cell == 'transformer':
receiver = Receiver(n_features=opts.n_features, n_hidden=opts.receiver_embedding)
receiver = core.TransformerReceiverDeterministic(receiver, opts.vocab_size, opts.max_len,
opts.receiver_embedding, opts.receiver_num_heads, opts.receiver_hidden,
opts.receiver_num_layers, causal=opts.causal_receiver)
else:
receiver = Receiver(n_features=opts.n_features, n_hidden=opts.receiver_hidden)
receiver = core.RnnReceiverDeterministic(receiver, opts.vocab_size, opts.receiver_embedding,
opts.receiver_hidden, cell=opts.receiver_cell,
num_layers=opts.receiver_num_layers)
empty_logger = LoggingStrategy.minimal()
game = core.SenderReceiverRnnReinforce(sender, receiver, loss, sender_entropy_coeff=opts.sender_entropy_coeff,
receiver_entropy_coeff=opts.receiver_entropy_coeff,
train_logging_strategy=empty_logger,
length_cost=opts.length_cost)
optimizer = core.build_optimizer(game.parameters())
callbacks = [EarlyStopperAccuracy(opts.early_stopping_thr),
core.ConsoleLogger(as_json=True, print_train_loss=True)]
if opts.checkpoint_dir:
checkpoint_name = f'{opts.name}_vocab{opts.vocab_size}_rs{opts.random_seed}_lr{opts.lr}_shid{opts.sender_hidden}_rhid{opts.receiver_hidden}_sentr{opts.sender_entropy_coeff}_reg{opts.length_cost}_max_len{opts.max_len}'
callbacks.append(core.CheckpointSaver(checkpoint_path=opts.checkpoint_dir, prefix=checkpoint_name))
trainer = core.Trainer(game=game, optimizer=optimizer, train_data=train_loader, validation_data=test_loader, callbacks=callbacks)
trainer.train(n_epochs=opts.n_epochs)
game.logging_strategy = LoggingStrategy.maximal() # now log everything
dump(trainer.game, opts.n_features, device, False)
core.close()
def main(params):
opts = get_params(params)
print(opts, flush=True)
# For compatibility, after https://github.com/facebookresearch/EGG/pull/130
# the meaning of `length` changed a bit. Before it included the EOS symbol; now
# it doesn't. To ensure that hyperparameters/CL arguments do not change,
# we subtract it here.
opts.max_len -= 1
device = opts.device
if opts.probs == "uniform":
probs = np.ones(opts.n_features)
elif opts.probs == "powerlaw":
probs = 1 / np.arange(1, opts.n_features + 1, dtype=np.float32)
else:
probs = np.array([float(x) for x in opts.probs.split(",")], dtype=np.float32)
probs /= probs.sum()
print("the probs are: ", probs, flush=True)
train_loader = OneHotLoader(
n_features=opts.n_features,
batch_size=opts.batch_size,
batches_per_epoch=opts.batches_per_epoch,
probs=probs,
)
# single batches with 1s on the diag
test_loader = UniformLoader(opts.n_features)
if opts.sender_cell == "transformer":
sender = Sender(n_features=opts.n_features, n_hidden=opts.sender_embedding)
sender = core.TransformerSenderReinforce(
agent=sender,
vocab_size=opts.vocab_size,
embed_dim=opts.sender_embedding,
max_len=opts.max_len,
num_layers=opts.sender_num_layers,
num_heads=opts.sender_num_heads,
hidden_size=opts.sender_hidden,
generate_style=opts.sender_generate_style,
causal=opts.causal_sender,
)
else:
sender = Sender(n_features=opts.n_features, n_hidden=opts.sender_hidden)
sender = core.RnnSenderReinforce(
sender,
opts.vocab_size,
opts.sender_embedding,
opts.sender_hidden,
cell=opts.sender_cell,
max_len=opts.max_len,
num_layers=opts.sender_num_layers,
)
if opts.receiver_cell == "transformer":
receiver = Receiver(
n_features=opts.n_features, n_hidden=opts.receiver_embedding
)
receiver = core.TransformerReceiverDeterministic(
receiver,
opts.vocab_size,
opts.max_len,
opts.receiver_embedding,
opts.receiver_num_heads,
opts.receiver_hidden,
opts.receiver_num_layers,
causal=opts.causal_receiver,
)
else:
receiver = Receiver(n_features=opts.n_features, n_hidden=opts.receiver_hidden)
receiver = core.RnnReceiverDeterministic(
receiver,
opts.vocab_size,
opts.receiver_embedding,
opts.receiver_hidden,
cell=opts.receiver_cell,
num_layers=opts.receiver_num_layers,
)
empty_logger = LoggingStrategy.minimal()
game = core.SenderReceiverRnnReinforce(
sender,
receiver,
loss,
sender_entropy_coeff=opts.sender_entropy_coeff,
receiver_entropy_coeff=opts.receiver_entropy_coeff,
train_logging_strategy=empty_logger,
length_cost=opts.length_cost,
)
optimizer = core.build_optimizer(game.parameters())
callbacks = [
EarlyStopperAccuracy(opts.early_stopping_thr),
core.ConsoleLogger(as_json=True, print_train_loss=True),
]
if opts.checkpoint_dir:
checkpoint_name = f"{opts.name}_vocab{opts.vocab_size}_rs{opts.random_seed}_lr{opts.lr}_shid{opts.sender_hidden}_rhid{opts.receiver_hidden}_sentr{opts.sender_entropy_coeff}_reg{opts.length_cost}_max_len{opts.max_len}"
callbacks.append(
core.CheckpointSaver(
checkpoint_path=opts.checkpoint_dir, prefix=checkpoint_name
)
)
trainer = core.Trainer(
game=game,
optimizer=optimizer,
train_data=train_loader,
validation_data=test_loader,
callbacks=callbacks,
)
trainer.train(n_epochs=opts.n_epochs)
game.logging_strategy = LoggingStrategy.maximal() # now log everything
dump(trainer.game, opts.n_features, device, False)
core.close()
