def main(params): print(torch.cuda.is_available()) 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) #elif opts.probs == "creneau": # ones = np.ones(int(opts.n_features/2)) # tens = 10*np.ones(opts.n_features-int(opts.n_features/2)) # probs = np.concatenate((tens,ones),axis=0) #elif opts.probs == "toy": # fives = 5*np.ones(int(opts.n_features/10)) # ones = np.ones(opts.n_features-int(opts.n_features/10)) # probs = np.concatenate((fives,ones),axis=0) #elif opts.probs == "escalier": # ones = np.ones(int(opts.n_features/4)) # tens = 10*np.ones(int(opts.n_features/4)) # huns = 100*np.ones(int(opts.n_features/4)) # thous = 1000*np.ones(opts.n_features-3*int(opts.n_features/4)) # probs = np.concatenate((thous,huns,tens,ones),axis=0) 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) if not opts.impatient: 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) else: receiver = Receiver(n_features=opts.receiver_hidden, n_hidden=opts.vocab_size) # If impatient 1 receiver = RnnReceiverImpatient(receiver, opts.vocab_size, opts.receiver_embedding, opts.receiver_hidden, cell=opts.receiver_cell, num_layers=opts.receiver_num_layers, max_len=opts.max_len, n_features=opts.n_features) # If impatient 2 #receiver = RnnReceiverImpatient2(receiver, opts.vocab_size, opts.receiver_embedding, # opts.receiver_hidden, cell=opts.receiver_cell, # num_layers=opts.receiver_num_layers, max_len=opts.max_len, n_features=opts.n_features) if not opts.impatient: game = core.SenderReceiverRnnReinforce(sender, receiver, loss, sender_entropy_coeff=opts.sender_entropy_coeff, receiver_entropy_coeff=opts.receiver_entropy_coeff, length_cost=opts.length_cost,unigram_penalty=opts.unigram_pen,reg=opts.reg) else: game = SenderImpatientReceiverRnnReinforce(sender, receiver, loss_impatient, sender_entropy_coeff=opts.sender_entropy_coeff, receiver_entropy_coeff=opts.receiver_entropy_coeff, length_cost=opts.length_cost,unigram_penalty=opts.unigram_pen,reg=opts.reg) optimizer = core.build_optimizer(game.parameters()) trainer = core.Trainer(game=game, optimizer=optimizer, train_data=train_loader, validation_data=test_loader, callbacks=[EarlyStopperAccuracy(opts.early_stopping_thr)]) for epoch in range(int(opts.n_epochs)): print("Epoch: "+str(epoch)) if epoch%100==0: trainer.optimizer.defaults["lr"]/=2 trainer.train(n_epochs=1) if opts.checkpoint_dir: trainer.save_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}') if not opts.impatient: acc_vec,messages=dump(trainer.game, opts.n_features, device, False,epoch) else: acc_vec,messages=dump_impatient(trainer.game, opts.n_features, device, False,epoch) # ADDITION TO SAVE MESSAGES all_messages=[] for x in messages: x = x.cpu().numpy() all_messages.append(x) all_messages = np.asarray(all_messages) if epoch%50==0: torch.save(sender.state_dict(), opts.dir_save+"/sender/sender_weights"+str(epoch)+".pth") torch.save(receiver.state_dict(), opts.dir_save+"/receiver/receiver_weights"+str(epoch)+".pth") #print(acc_vec) np.save(opts.dir_save+'/messages/messages_'+str((epoch))+'.npy', all_messages) np.save(opts.dir_save+'/accuracy/accuracy_'+str((epoch))+'.npy', acc_vec) 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) 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() 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) if not opts.impatient: 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) else: receiver = Receiver(n_features=opts.receiver_hidden, n_hidden=opts.vocab_size) # If impatient 1 receiver = RnnReceiverImpatient(receiver, opts.vocab_size, opts.receiver_embedding, opts.receiver_hidden, cell=opts.receiver_cell, num_layers=opts.receiver_num_layers, max_len=opts.max_len, n_features=opts.n_features) # If impatient 2 #receiver = RnnReceiverImpatient2(receiver, opts.vocab_size, opts.receiver_embedding, # opts.receiver_hidden, cell=opts.receiver_cell, # num_layers=opts.receiver_num_layers, max_len=opts.max_len, n_features=opts.n_features) sender.load_state_dict(torch.load(opts.sender_weights,map_location=torch.device('cpu'))) receiver.load_state_dict(torch.load(opts.receiver_weights,map_location=torch.device('cpu'))) if not opts.impatient: game = core.SenderReceiverRnnReinforce(sender, receiver, loss, sender_entropy_coeff=opts.sender_entropy_coeff, receiver_entropy_coeff=opts.receiver_entropy_coeff, length_cost=opts.length_cost,unigram_penalty=opts.unigram_pen) else: game = SenderImpatientReceiverRnnReinforce(sender, receiver, loss, sender_entropy_coeff=opts.sender_entropy_coeff, receiver_entropy_coeff=opts.receiver_entropy_coeff, length_cost=opts.length_cost,unigram_penalty=opts.unigram_pen) optimizer = core.build_optimizer(game.parameters()) trainer = core.Trainer(game=game, optimizer=optimizer, train_data=train_loader, validation_data=test_loader, callbacks=[EarlyStopperAccuracy(opts.early_stopping_thr)]) # Test impose message if not opts.impatient: acc_vec,messages=dump(trainer.game, opts.n_features, device, False) else: acc_vec,messages=dump_impatient(trainer.game, opts.n_features, device, False,save_dir=opts.save_dir) all_messages=[] for x in messages: x = x.cpu().numpy() all_messages.append(x) all_messages = np.asarray(all_messages) messages=-1*np.ones((opts.n_features,opts.max_len)) for i in range(len(all_messages)): for j in range(all_messages[i].shape[0]): messages[i,j]=all_messages[i][j] np.save(opts.save_dir+"messages_analysis.npy",messages) 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) elif opts.probs == 'perso': probs = opts.n_features+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) game = core.SenderReceiverRnnReinforce(sender, receiver, loss, sender_entropy_coeff=opts.sender_entropy_coeff, receiver_entropy_coeff=opts.receiver_entropy_coeff, length_cost=opts.length_cost) optimizer = core.build_optimizer(game.parameters()) trainer = core.Trainer(game=game, optimizer=optimizer, train_data=train_loader, validation_data=test_loader, callbacks=[EarlyStopperAccuracy(opts.early_stopping_thr), core.ConsoleLogger(as_json=True, print_train_loss=True)]) """ mode accuracy chope a chaque epoch accs=[] all_messages,acc=dump(trainer.game, opts.n_features, device, False) np.save('messages_0.npy', all_messages) accs.append(acc) for i in range(int(opts.n_epochs)): print(i) trainer.train(n_epochs=1) all_messages,acc=dump(trainer.game, opts.n_features, device, False) np.save('messages_'+str((i+1))+'.npy', all_messages) accs.append(acc) np.save('accuracy.npy',accs) """ trainer.train(n_epochs=opts.n_epochs) #if opts.checkpoint_dir: #trainer.save_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}') #for i in range(30): # for k in range(30): # if i<k: # all_messages=dump(trainer.game, opts.n_features, device, False,pos_m=i,pos_M=k) all_messages=dump(trainer.game, opts.n_features, device, False) print(all_messages) #freq=np.zeros(30) #for message in all_messages[0]: # if i in range(message.shape[0]): # freq[int(message[i])]+=1 #print(freq) core.close()
def main(params): opts = get_params(params) print(opts) device = opts.device train_loader = OneHotLoader(n_bits=opts.n_bits, bits_s=opts.bits_s, bits_r=opts.bits_r, batch_size=opts.batch_size, batches_per_epoch=opts.n_examples_per_epoch / opts.batch_size) test_loader = UniformLoader(n_bits=opts.n_bits, bits_s=opts.bits_s, bits_r=opts.bits_r) test_loader.batch = [x.to(device) for x in test_loader.batch] if not opts.variable_length: sender = Sender(n_bits=opts.n_bits, n_hidden=opts.sender_hidden, vocab_size=opts.vocab_size) if opts.mode == 'gs': sender = core.GumbelSoftmaxWrapper(agent=sender, temperature=opts.temperature) receiver = Receiver(n_bits=opts.n_bits, n_hidden=opts.receiver_hidden) receiver = core.SymbolReceiverWrapper( receiver, vocab_size=opts.vocab_size, agent_input_size=opts.receiver_hidden) game = core.SymbolGameGS(sender, receiver, diff_loss) elif opts.mode == 'rf': sender = core.ReinforceWrapper(agent=sender) receiver = Receiver(n_bits=opts.n_bits, n_hidden=opts.receiver_hidden) receiver = core.SymbolReceiverWrapper( receiver, vocab_size=opts.vocab_size, agent_input_size=opts.receiver_hidden) receiver = core.ReinforceDeterministicWrapper(agent=receiver) game = core.SymbolGameReinforce( sender, receiver, diff_loss, sender_entropy_coeff=opts.sender_entropy_coeff) elif opts.mode == 'non_diff': sender = core.ReinforceWrapper(agent=sender) receiver = ReinforcedReceiver(n_bits=opts.n_bits, n_hidden=opts.receiver_hidden) receiver = core.SymbolReceiverWrapper( receiver, vocab_size=opts.vocab_size, agent_input_size=opts.receiver_hidden) game = core.SymbolGameReinforce( sender, receiver, non_diff_loss, sender_entropy_coeff=opts.sender_entropy_coeff, receiver_entropy_coeff=opts.receiver_entropy_coeff) else: if opts.mode != 'rf': print('Only mode=rf is supported atm') opts.mode = 'rf' if opts.sender_cell == 'transformer': receiver = Receiver(n_bits=opts.n_bits, n_hidden=opts.receiver_hidden) sender = Sender( n_bits=opts.n_bits, n_hidden=opts.sender_hidden, vocab_size=opts.sender_hidden) # TODO: not really vocab sender = core.TransformerSenderReinforce( agent=sender, vocab_size=opts.vocab_size, embed_dim=opts.sender_emb, max_len=opts.max_len, num_layers=1, num_heads=1, hidden_size=opts.sender_hidden) else: receiver = Receiver(n_bits=opts.n_bits, n_hidden=opts.receiver_hidden) sender = Sender( n_bits=opts.n_bits, n_hidden=opts.sender_hidden, vocab_size=opts.sender_hidden) # TODO: not really vocab sender = core.RnnSenderReinforce(agent=sender, vocab_size=opts.vocab_size, embed_dim=opts.sender_emb, hidden_size=opts.sender_hidden, max_len=opts.max_len, cell=opts.sender_cell) if opts.receiver_cell == 'transformer': receiver = Receiver(n_bits=opts.n_bits, n_hidden=opts.receiver_emb) receiver = core.TransformerReceiverDeterministic( receiver, opts.vocab_size, opts.max_len, opts.receiver_emb, num_heads=1, hidden_size=opts.receiver_hidden, num_layers=1) else: receiver = Receiver(n_bits=opts.n_bits, n_hidden=opts.receiver_hidden) receiver = core.RnnReceiverDeterministic(receiver, opts.vocab_size, opts.receiver_emb, opts.receiver_hidden, cell=opts.receiver_cell) game = core.SenderReceiverRnnGS(sender, receiver, diff_loss) game = core.SenderReceiverRnnReinforce( sender, receiver, diff_loss, sender_entropy_coeff=opts.sender_entropy_coeff, receiver_entropy_coeff=opts.receiver_entropy_coeff) optimizer = torch.optim.Adam([ dict(params=sender.parameters(), lr=opts.sender_lr), dict(params=receiver.parameters(), lr=opts.receiver_lr) ]) loss = game.loss intervention = CallbackEvaluator(test_loader, device=device, is_gs=opts.mode == 'gs', loss=loss, var_length=opts.variable_length, input_intervention=True) trainer = core.Trainer(game=game, optimizer=optimizer, train_data=train_loader, validation_data=test_loader, callbacks=[ core.ConsoleLogger(as_json=True), EarlyStopperAccuracy(opts.early_stopping_thr), intervention ]) trainer.train(n_epochs=opts.n_epochs) 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()
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() 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) if not opts.impatient: 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) else: receiver = Receiver(n_features=opts.receiver_hidden, n_hidden=opts.vocab_size) # If impatient 1 receiver = RnnReceiverImpatient( receiver, opts.vocab_size, opts.receiver_embedding, opts.receiver_hidden, cell=opts.receiver_cell, num_layers=opts.receiver_num_layers, max_len=opts.max_len, n_features=opts.n_features) sender.load_state_dict( torch.load(opts.sender_weights, map_location=torch.device('cpu'))) receiver.load_state_dict( torch.load(opts.receiver_weights, map_location=torch.device('cpu'))) if not opts.impatient: game = core.SenderReceiverRnnReinforce( sender, receiver, loss, sender_entropy_coeff=opts.sender_entropy_coeff, receiver_entropy_coeff=opts.receiver_entropy_coeff, length_cost=opts.length_cost, unigram_penalty=opts.unigram_pen) else: game = SenderImpatientReceiverRnnReinforce( sender, receiver, loss, sender_entropy_coeff=opts.sender_entropy_coeff, receiver_entropy_coeff=opts.receiver_entropy_coeff, length_cost=opts.length_cost, unigram_penalty=opts.unigram_pen) optimizer = core.build_optimizer(game.parameters()) trainer = core.Trainer( game=game, optimizer=optimizer, train_data=train_loader, validation_data=test_loader, callbacks=[EarlyStopperAccuracy(opts.early_stopping_thr)]) # Debut test position position_sieve = np.zeros((opts.n_features, opts.max_len)) for position in range(opts.max_len): dataset = [[torch.eye(opts.n_features).to(device), None]] if opts.impatient: sender_inputs, messages, receiver_inputs, receiver_outputs, _ = \ dump_test_position_impatient(trainer.game, dataset, position=position, voc_size=opts.vocab_size, gs=False, device=device, variable_length=True) else: sender_inputs, messages, receiver_inputs, receiver_outputs, _ = \ dump_test_position(trainer.game, dataset, position=position, voc_size=opts.vocab_size, gs=False, device=device, variable_length=True) acc_pos = [] for sender_input, message, receiver_output in zip( sender_inputs, messages, receiver_outputs): input_symbol = sender_input.argmax() output_symbol = receiver_output.argmax() acc = (input_symbol == output_symbol).float().item() acc_pos.append(acc) acc_pos = np.array(acc_pos) position_sieve[:, position] = acc_pos # Put -1 for position after message_length _, messages = dump(trainer.game, opts.n_features, device, False) # Convert messages to numpy array messages_np = [] for x in messages: x = x.cpu().numpy() messages_np.append(x) for i in range(len(messages_np)): # Message i message_i = messages_np[i] id_0 = np.where(message_i == 0)[0] if id_0.shape[0] > 0: for j in range(id_0[0] + 1, opts.max_len): position_sieve[i, j] = -1 np.save("analysis/position_sieve.npy", position_sieve) core.close()