def load_encoder_decoder(voc, checkpoint, configs):
"""
Initialize encoder and decoder, and load from file if prev states exists
:param voc: Vocabulary
:param checkpoint: dict
:param configs: dict
:return: Encoder, LuongAttentionDecoderRNN
"""
logging.info('Building encoder and decoder ...')
# Initialize word embeddings
embedding = nn.Embedding(voc.num_words, configs["hidden_size"])
# Initialize encoder & decoder models
encoder = EncoderRNN(configs["hidden_size"], embedding,
configs["encoder_n_layers"], configs["dropout"])
decoder = LuongAttentionDecoderRNN(embedding, voc.num_words, configs)
if checkpoint:
voc.__dict__ = checkpoint['voc_dict']
embedding.load_state_dict(checkpoint['embedding'])
encoder.load_state_dict(checkpoint['en'])
decoder.load_state_dict(checkpoint['de'])
logging.info('Models built and ready to go!')
return encoder.to(get_device()), decoder.to(get_device())
def main(args):
config_path = os.path.join(args.config_path, 'config.json')
with open(config_path) as f:
config = json.load(f)
print('[-] Loading pickles')
dataset_path = Path(config["dataset_path"])
input_lang = CustomUnpickler(open(dataset_path / 'input_lang.pkl',
'rb')).load()
output_lang = CustomUnpickler(open(dataset_path / 'output_lang.pkl',
'rb')).load()
pairs = CustomUnpickler(open(dataset_path / 'pairs.pkl', 'rb')).load()
# input_lang = load_pkl(dataset_path / 'input_lang.pkl')
# output_lang = load_pkl(dataset_path / 'output_lang.pkl')
# pairs = load_pkl(dataset_path / 'pairs.pkl')
hidden_size = config["model_cfg"]["hidden_size"]
max_len = config["max_len"]
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
encoder = EncoderRNN(input_lang.n_words, hidden_size).to(device)
decoder = AttnDecoderRNN(hidden_size,
output_lang.n_words,
max_len,
dropout_p=0.1).to(device)
print('[-] Loading models')
ckpt = torch.load(config["models_path"] + 'models.ckpt')
encoder.load_state_dict(ckpt['encoder'])
encoder.to(device)
decoder.load_state_dict(ckpt['decoder'])
decoder.to(device)
evaluator = Evaluater(device, encoder, decoder, input_lang, output_lang,
max_len)
# Evaluate random samples
evaluator.evaluateRandomly(pairs)
evaluator.evaluateAndShowAttention("elle a cinq ans de moins que moi .")
# evaluator.evaluateAndShowAttention("elle est trop petit .")
# evaluator.evaluateAndShowAttention("je ne crains pas de mourir .")
# evaluator.evaluateAndShowAttention("c est un jeune directeur plein de talent .")
plt.savefig('attention.png')
def main():
# 加载词库，加载数据集
voc = Lang('data/WORDMAP.json')
print("词库数量 " + str(voc.n_words))
train_data = SaDataset('train', voc)
val_data = SaDataset('valid', voc)
# 初始化模型
encoder = EncoderRNN(voc.n_words, hidden_size, encoder_n_layers, dropout)
# 将模型使用device进行计算，如果是gpu,则会使用显存，如果是cpu，则会使用内存
encoder = encoder.to(device)
# 初始化优化器 优化器的目的是让梯度下降，手段是调整模型的参数，optim是一个pytorch的一个包，adam是一个优化算法，梯度下降
print('Building optimizers ...')
'''
需要优化的参数
学习率
'''
optimizer = optim.Adam(encoder.parameters(), lr=learning_rate)
# 基础准确率
best_acc = 0
epochs_since_improvement = 0
# epochs 训练的次数
for epoch in range(0, epochs):
# Decay learning rate if there is no improvement for 8 consecutive epochs, and terminate training after 20
if epochs_since_improvement == 20:
break
if epochs_since_improvement > 0 and epochs_since_improvement % 8 == 0:
adjust_learning_rate(optimizer, 0.8)
# 训练一次
train(epoch, train_data, encoder, optimizer)
# 使用验证集对训练结果进行验证，防止过拟合
val_acc, val_loss = valid(val_data, encoder)
print('\n * ACCURACY - {acc:.3f}, LOSS - {loss:.3f}\n'.format(acc=val_acc, loss=val_loss))
# 检查是否有提升
is_best = val_acc > best_acc
best_acc = max(best_acc, val_acc)
if not is_best:
epochs_since_improvement += 1
print("\nEpochs since last improvement: %d\n" % (epochs_since_improvement,))
else:
epochs_since_improvement = 0
# Save checkpoint
save_checkpoint(epoch, encoder, optimizer, val_acc, is_best)
# Reshuffle samples 将验证集合测试集打乱
np.random.shuffle(train_data.samples)
np.random.shuffle(val_data.samples)
def init():
print("\tInitialising sentences")
print("\t\tLoading and cleaning json files")
json_of_convs = load_all_json_conv('./Dataset/messages')
print("\t\tLoading two person convs")
duo_conversations = get_chat_friend_and_me(json_of_convs)
print("\t\tMaking two person convs discussions")
discussions = get_discussions(duo_conversations)
print("\t\tCreating pairs for training")
pairs_of_sentences = make_pairs(discussions)
print(f"\t\t{len(pairs_of_sentences)} different pairs")
print("\t\tCreating Vocabulary")
voc = Voc()
print("\t\tPopulating Vocabulary")
voc.createVocFromPairs(pairs_of_sentences)
print(f"\t\tVocabulary of : {voc.num_words} differents words")
print('\tBuilding encoder and decoder ...')
embedding = nn.Embedding(voc.num_words, HIDDEN_SIZE)
encoder = EncoderRNN(HIDDEN_SIZE, embedding, ENCODER_N_LAYERS, DROPOUT)
decoder = LuongAttnDecoderRNN(ATTN_MODEL, embedding, HIDDEN_SIZE,
voc.num_words, DECODER_N_LAYERS, DROPOUT)
encoder_optimizer = optim.Adam(encoder.parameters(), lr=LEARNING_RATE)
decoder_optimizer = optim.Adam(decoder.parameters(),
lr=LEARNING_RATE * DECODER_LEARNING_RATIO)
checkpoint = None
if LOADFILENAME:
print("\t\tLoading last training")
checkpoint = torch.load(LOADFILENAME)
# If loading a model trained on GPU to CPU
# checkpoint=torch.load(loadFilename,map_location=torch.device('cpu'))
encoder_sd = checkpoint['en']
decoder_sd = checkpoint['de']
encoder_optimizer_sd = checkpoint['en_opt']
decoder_optimizer_sd = checkpoint['de_opt']
embedding_sd = checkpoint['embedding']
voc.__dict__ = checkpoint['voc_dict']
print("\t\tPopulating from last training")
embedding.load_state_dict(embedding_sd)
encoder.load_state_dict(encoder_sd)
decoder.load_state_dict(decoder_sd)
encoder_optimizer.load_state_dict(encoder_optimizer_sd)
decoder_optimizer.load_state_dict(decoder_optimizer_sd)
encoder = encoder.to(DEVICE)
decoder = decoder.to(DEVICE)
return (encoder, decoder, encoder_optimizer, decoder_optimizer, embedding,
voc, pairs_of_sentences, checkpoint)
def main():
input_lang = Lang('data/WORDMAP_en.json')
output_lang = Lang('data/WORDMAP_zh.json')
print("input_lang.n_words: " + str(input_lang.n_words))
print("output_lang.n_words: " + str(output_lang.n_words))
train_data = TranslationDataset('train')
val_data = TranslationDataset('valid')
# Initialize encoder & decoder models
encoder = EncoderRNN(input_lang.n_words, hidden_size, encoder_n_layers,
dropout)
decoder = LuongAttnDecoderRNN(attn_model, hidden_size, output_lang.n_words,
decoder_n_layers, dropout)
# Use appropriate device
encoder = encoder.to(device)
decoder = decoder.to(device)
# Initialize optimizers
print('Building optimizers ...')
encoder_optimizer = optim.Adam(encoder.parameters(), lr=learning_rate)
decoder_optimizer = optim.Adam(decoder.parameters(), lr=learning_rate)
# Initializations
print('Initializing ...')
train_batch_time = ExpoAverageMeter() # forward prop. + back prop. time
train_losses = ExpoAverageMeter() # loss (per word decoded)
val_batch_time = ExpoAverageMeter()
val_losses = ExpoAverageMeter()
best_loss = 100000
epochs_since_improvement = 0
# Epochs
for epoch in range(start_epoch, epochs):
# Decay learning rate if there is no improvement for 8 consecutive epochs, and terminate training after 20
if epochs_since_improvement == 20:
break
if epochs_since_improvement > 0 and epochs_since_improvement % 8 == 0:
adjust_learning_rate(decoder_optimizer, 0.8)
adjust_learning_rate(encoder_optimizer, 0.8)
# One epoch's training
# Ensure dropout layers are in train mode
encoder.train()
decoder.train()
start = time.time()
# Batches
for i_batch in range(len(train_data)):
input_variable, lengths, target_variable, mask, max_target_len = train_data[
i_batch]
train_loss = train(input_variable, lengths, target_variable, mask,
max_target_len, encoder, decoder,
encoder_optimizer, decoder_optimizer)
# Keep track of metrics
train_losses.update(train_loss)
train_batch_time.update(time.time() - start)
start = time.time()
# Print status
if i_batch % print_every == 0:
print(
'[{0}] Epoch: [{1}][{2}/{3}]\t'
'Batch Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'.format(
timestamp(),
epoch,
i_batch,
len(train_data),
batch_time=train_batch_time,
loss=train_losses))
# One epoch's validation
start = time.time()
# Batches
for i_batch in range(len(val_data)):
input_variable, lengths, target_variable, mask, max_target_len = val_data[
i_batch]
val_loss = valid(input_variable, lengths, target_variable, mask,
max_target_len, encoder, decoder)
# Keep track of metrics
val_losses.update(val_loss)
val_batch_time.update(time.time() - start)
start = time.time()
# Print status
if i_batch % print_every == 0:
print(
'Validation: [{0}/{1}]\t'
'Batch Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'.format(
i_batch,
len(val_data),
batch_time=val_batch_time,
loss=val_losses))
val_loss = val_losses.avg
print('\n * LOSS - {loss:.3f}\n'.format(loss=val_loss))
# Check if there was an improvement
is_best = val_loss < best_loss
best_loss = min(best_loss, val_loss)
if not is_best:
epochs_since_improvement += 1
print("\nEpochs since last improvement: %d\n" %
(epochs_since_improvement, ))
else:
epochs_since_improvement = 0
save_checkpoint(epoch, encoder, decoder, encoder_optimizer,
decoder_optimizer, input_lang, output_lang, val_loss,
is_best)
# Initialize search module
searcher = GreedySearchDecoder(encoder, decoder)
for input_sentence, target_sentence in pick_n_valid_sentences(
input_lang, output_lang, 10):
decoded_words = evaluate(searcher, input_sentence, input_lang,
output_lang)
print('> {}'.format(input_sentence))
print('= {}'.format(target_sentence))
print('< {}'.format(''.join(decoded_words)))
# Reshuffle train and valid samples
np.random.shuffle(train_data.samples)
np.random.shuffle(val_data.samples)
def main(opts):
# set manual_seed and build vocab
print(opts, flush=True)
setup(opts, opts.seed)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Usando {device} :)")
# create a batch training environment that will also preprocess text
vocab = read_vocab(opts.train_vocab)
tok = Tokenizer(opts.remove_punctuation == 1, opts.reversed == 1, vocab=vocab, encoding_length=opts.max_cap_length)
# create language instruction encoder
encoder_kwargs = {
'opts': opts,
'vocab_size': len(vocab),
'embedding_size': opts.word_embedding_size,
'hidden_size': opts.rnn_hidden_size,
'padding_idx': padding_idx,
'dropout_ratio': opts.rnn_dropout,
'bidirectional': opts.bidirectional == 1,
'num_layers': opts.rnn_num_layers
}
print('Using {} as encoder ...'.format(opts.lang_embed))
if 'lstm' in opts.lang_embed:
encoder = EncoderRNN(**encoder_kwargs)
else:
raise ValueError('Unknown {} language embedding'.format(opts.lang_embed))
print(encoder)
# create policy model
policy_model_kwargs = {
'opts':opts,
'img_fc_dim': opts.img_fc_dim,
'img_fc_use_batchnorm': opts.img_fc_use_batchnorm == 1,
'img_dropout': opts.img_dropout,
'img_feat_input_dim': opts.img_feat_input_dim,
'rnn_hidden_size': opts.rnn_hidden_size,
'rnn_dropout': opts.rnn_dropout,
'max_len': opts.max_cap_length,
'max_navigable': opts.max_navigable
}
if opts.arch == 'regretful':
model = Regretful(**policy_model_kwargs)
elif opts.arch == 'self-monitoring':
model = SelfMonitoring(**policy_model_kwargs)
elif opts.arch == 'speaker-baseline':
model = SpeakerFollowerBaseline(**policy_model_kwargs)
else:
raise ValueError('Unknown {} model for seq2seq agent'.format(opts.arch))
print(model)
encoder = encoder.to(device)
model = model.to(device)
params = list(encoder.parameters()) + list(model.parameters())
optimizer = torch.optim.Adam(params, lr=opts.learning_rate)
# optionally resume from a checkpoint
if opts.resume:
model, encoder, optimizer, best_success_rate = resume_training(opts, model, encoder, optimizer)
# if a secondary exp name is specified, this is useful when resuming from a previous saved
# experiment and save to another experiment, e.g., pre-trained on synthetic data and fine-tune on real data
if opts.exp_name_secondary:
opts.exp_name += opts.exp_name_secondary
feature, img_spec = load_features(opts.img_feat_dir, opts.blind)
if opts.test_submission:
assert opts.resume, 'The model was not resumed before running for submission.'
test_env = ('test', (R2RPanoBatch(opts, feature, img_spec, batch_size=opts.batch_size,
splits=['test'], tokenizer=tok), Evaluation(['test'], opts)))
agent_kwargs = {
'opts': opts,
'env': test_env[1][0],
'results_path': "",
'encoder': encoder,
'model': model,
'feedback': opts.feedback
}
agent = PanoSeq2SeqAgent(**agent_kwargs)
# setup trainer
trainer = PanoSeq2SeqTrainer(opts, agent, optimizer)
epoch = opts.start_epoch - 1
trainer.eval(epoch, test_env)
return
# set up R2R environments
if not opts.train_data_augmentation:
train_env = R2RPanoBatch(opts, feature, img_spec, batch_size=opts.batch_size, seed=opts.seed,
splits=['train'], tokenizer=tok)
else:
train_env = R2RPanoBatch(opts, feature, img_spec, batch_size=opts.batch_size, seed=opts.seed,
splits=['synthetic'], tokenizer=tok)
val_craft_splits = ['craft_seen', 'craft_unseen']
val_splits = ['val_seen', 'val_unseen']
if opts.craft_eval:
val_splits += val_craft_splits
val_envs = {split: (R2RPanoBatch(opts, feature, img_spec, batch_size=opts.batch_size,
splits=[split], tokenizer=tok), Evaluation([split], opts))
for split in val_splits}
# create agent
agent_kwargs = {
'opts': opts,
'env': train_env,
'results_path': "",
'encoder': encoder,
'model': model,
'feedback': opts.feedback
}
agent = PanoSeq2SeqAgent(**agent_kwargs)
# setup trainer
trainer = PanoSeq2SeqTrainer(opts, agent, optimizer, opts.train_iters_epoch)
if opts.eval_only:
success_rate = []
for val_env in val_envs.items():
success_rate.append(trainer.eval(opts.start_epoch - 1, val_env, tb_logger=None))
return
# set up tensorboard logger
tb_logger = set_tb_logger(opts.log_dir, opts.exp_name, opts.resume)
sys.stdout.flush()
best_success_rate = best_success_rate if opts.resume else 0.0
for epoch in range(opts.start_epoch, opts.max_num_epochs + 1):
trainer.train(epoch, train_env, tb_logger)
if epoch % opts.eval_every_epochs == 0:
success_rate = []
for val_env in val_envs.items():
success_rate.append(trainer.eval(epoch, val_env, tb_logger))
success_rate_compare = success_rate[1]
if is_experiment():
# remember best val_seen success rate and save checkpoint
is_best = success_rate_compare >= best_success_rate
best_success_rate = max(success_rate_compare, best_success_rate)
print("--> Highest val_unseen success rate: {}".format(best_success_rate))
sys.stdout.flush()
# save the model if it is the best so far
save_checkpoint({
'opts': opts,
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'encoder_state_dict': encoder.state_dict(),
'best_success_rate': best_success_rate,
'optimizer': optimizer.state_dict(),
'max_episode_len': opts.max_episode_len,
}, is_best, checkpoint_dir=opts.checkpoint_dir, name=opts.exp_name)
if opts.train_data_augmentation and epoch == opts.epochs_data_augmentation:
train_env = R2RPanoBatch(opts, feature, img_spec, batch_size=opts.batch_size, seed=opts.seed,
splits=['train'], tokenizer=tok)
print("--> Finished training")
def main(opts):
# set manual_seed and build vocab
setup(opts, opts.seed)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# create a batch training environment that will also preprocess text
vocab = read_vocab(opts.train_vocab)
tok = Tokenizer(
opts.remove_punctuation == 1,
opts.reversed == 1,
vocab=vocab,
encoding_length=opts.max_cap_length,
)
# create language instruction encoder
encoder_kwargs = {
"opts": opts,
"vocab_size": len(vocab),
"embedding_size": opts.word_embedding_size,
"hidden_size": opts.rnn_hidden_size,
"padding_idx": padding_idx,
"dropout_ratio": opts.rnn_dropout,
"bidirectional": opts.bidirectional == 1,
"num_layers": opts.rnn_num_layers,
}
print("Using {} as encoder ...".format(opts.lang_embed))
if "lstm" in opts.lang_embed:
encoder = EncoderRNN(**encoder_kwargs)
else:
raise ValueError("Unknown {} language embedding".format(
opts.lang_embed))
print(encoder)
# create policy model
policy_model_kwargs = {
"opts": opts,
"img_fc_dim": opts.img_fc_dim,
"img_fc_use_batchnorm": opts.img_fc_use_batchnorm == 1,
"img_dropout": opts.img_dropout,
"img_feat_input_dim": opts.img_feat_input_dim,
"rnn_hidden_size": opts.rnn_hidden_size,
"rnn_dropout": opts.rnn_dropout,
"max_len": opts.max_cap_length,
"max_navigable": opts.max_navigable,
}
if opts.arch == "self-monitoring":
model = SelfMonitoring(**policy_model_kwargs)
elif opts.arch == "speaker-baseline":
model = SpeakerFollowerBaseline(**policy_model_kwargs)
else:
raise ValueError("Unknown {} model for seq2seq agent".format(
opts.arch))
print(model)
encoder = encoder.to(device)
model = model.to(device)
params = list(encoder.parameters()) + list(model.parameters())
optimizer = torch.optim.Adam(params, lr=opts.learning_rate)
# optionally resume from a checkpoint
if opts.resume:
model, encoder, optimizer, best_success_rate = resume_training(
opts, model, encoder, optimizer)
# if a secondary exp name is specified, this is useful when resuming from a previous saved
# experiment and save to another experiment, e.g., pre-trained on synthetic data and fine-tune on real data
if opts.exp_name_secondary:
opts.exp_name += opts.exp_name_secondary
feature, img_spec = load_features(opts.img_feat_dir)
if opts.test_submission:
assert (opts.resume
), "The model was not resumed before running for submission."
test_env = (
"test",
(
R2RPanoBatch(
opts,
feature,
img_spec,
batch_size=opts.batch_size,
splits=["test"],
tokenizer=tok,
),
Evaluation(["test"]),
),
)
agent_kwargs = {
"opts": opts,
"env": test_env[1][0],
"results_path": "",
"encoder": encoder,
"model": model,
"feedback": opts.feedback,
}
agent = PanoSeq2SeqAgent(**agent_kwargs)
# setup trainer
trainer = PanoSeq2SeqTrainer(opts, agent, optimizer)
epoch = opts.start_epoch - 1
trainer.eval(epoch, test_env)
return
# set up R2R environments
if not opts.train_data_augmentation:
train_env = R2RPanoBatch(
opts,
feature,
img_spec,
batch_size=opts.batch_size,
seed=opts.seed,
splits=["train"],
tokenizer=tok,
)
else:
train_env = R2RPanoBatch(
opts,
feature,
img_spec,
batch_size=opts.batch_size,
seed=opts.seed,
splits=["synthetic"],
tokenizer=tok,
)
val_envs = {
split: (
R2RPanoBatch(
opts,
feature,
img_spec,
batch_size=opts.batch_size,
splits=[split],
tokenizer=tok,
),
Evaluation([split]),
)
for split in ["val_seen", "val_unseen"]
}
# create agent
agent_kwargs = {
"opts": opts,
"env": train_env,
"results_path": "",
"encoder": encoder,
"model": model,
"feedback": opts.feedback,
}
agent = PanoSeq2SeqAgent(**agent_kwargs)
# setup trainer
trainer = PanoSeq2SeqTrainer(opts, agent, optimizer,
opts.train_iters_epoch)
if opts.eval_beam or opts.eval_only:
success_rate = []
for val_env in val_envs.items():
success_rate.append(
trainer.eval(opts.start_epoch - 1, val_env, tb_logger=None))
return
# set up tensorboard logger
tb_logger = set_tb_logger(opts.log_dir, opts.exp_name, opts.resume)
best_success_rate = best_success_rate if opts.resume else 0.0
for epoch in range(opts.start_epoch, opts.max_num_epochs + 1):
trainer.train(epoch, train_env, tb_logger)
if epoch % opts.eval_every_epochs == 0:
success_rate = []
for val_env in val_envs.items():
success_rate.append(trainer.eval(epoch, val_env, tb_logger))
success_rate_compare = success_rate[1]
if is_experiment():
# remember best val_seen success rate and save checkpoint
is_best = success_rate_compare >= best_success_rate
best_success_rate = max(success_rate_compare,
best_success_rate)
print("--> Highest val_unseen success rate: {}".format(
best_success_rate))
# save the model if it is the best so far
save_checkpoint(
{
"opts": opts,
"epoch": epoch + 1,
"state_dict": model.state_dict(),
"encoder_state_dict": encoder.state_dict(),
"best_success_rate": best_success_rate,
"optimizer": optimizer.state_dict(),
"max_episode_len": opts.max_episode_len,
},
is_best,
checkpoint_dir=opts.checkpoint_dir,
name=opts.exp_name,
)
if (opts.train_data_augmentation
and epoch == opts.epochs_data_augmentation):
train_env = R2RPanoBatch(
opts,
feature,
img_spec,
batch_size=opts.batch_size,
seed=opts.seed,
splits=["train"],
tokenizer=tok,
)
print("--> Finished training")
def main():
train_loader = ChatbotDataset('train')
val_loader = ChatbotDataset('valid')
# Initialize word embeddings
embedding = nn.Embedding(voc.num_words, hidden_size)
# Initialize encoder & decoder models
encoder = EncoderRNN(hidden_size, embedding, encoder_n_layers, dropout)
decoder = LuongAttnDecoderRNN(attn_model, embedding, hidden_size, voc.num_words, decoder_n_layers, dropout)
# Use appropriate device
encoder = encoder.to(device)
decoder = decoder.to(device)
# Initialize optimizers
print('Building optimizers ...')
encoder_optimizer = optim.Adam(encoder.parameters(), lr=learning_rate)
decoder_optimizer = optim.Adam(decoder.parameters(), lr=learning_rate)
# Initializations
print('Initializing ...')
batch_time = AverageMeter() # forward prop. + back prop. time
losses = AverageMeter() # loss (per word decoded)
# Epochs
for epoch in range(start_epoch, epochs):
# One epoch's training
# Ensure dropout layers are in train mode
encoder.train()
decoder.train()
start = time.time()
# Batches
for i in range(train_loader.__len__()):
input_variable, lengths, target_variable, mask, max_target_len = train_loader.__getitem__(i)
loss = train(input_variable, lengths, target_variable, mask, max_target_len, encoder, decoder,
encoder_optimizer, decoder_optimizer)
# Keep track of metrics
losses.update(loss, max_target_len)
batch_time.update(time.time() - start)
start = time.time()
if i % print_every == 0:
print('[{0}] Epoch: [{1}][{2}/{3}]\t'
'Batch Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'.format(timestamp(), epoch, i, len(train_loader),
batch_time=batch_time,
loss=losses))
# One epoch's validation
val_loss = validate(val_loader, encoder, decoder)
print('\n * LOSS - {loss:.3f}\n'.format(loss=val_loss))
# Initialize search module
searcher = GreedySearchDecoder(encoder, decoder)
for sentence in pick_n_valid_sentences(10):
decoded_words = evaluate(searcher, sentence)
print('Human: {}'.format(sentence))
print('Bot: {}'.format(''.join(decoded_words)))
# Save checkpoint
if epoch % save_every == 0:
directory = save_dir
if not os.path.exists(directory):
os.makedirs(directory)
torch.save({
'epoch': epoch,
'en': encoder.state_dict(),
'de': decoder.state_dict(),
'en_opt': encoder_optimizer.state_dict(),
'de_opt': decoder_optimizer.state_dict(),
'loss': loss,
'voc': voc.__dict__
}, os.path.join(directory, '{}_{}_{}.tar'.format('checkpoint', epoch, val_loss)))
'voc_dict': voc.__dict__,
'embedding': embedding.state_dict()
},
os.path.join(directory,
'{}_{}.tar'.format(iteration, 'checkpoint')))
print('Building encoder and decoder ...')
# word embedding
embedding = nn.Embedding(VOC.num_words, hp.hidden_size)
encoder = EncoderRNN(hp.hidden_size, embedding, hp.n_layers, hp.dropout)
decoder = LuongAttnDecoderRNN(hp.attn_model, embedding, hp.hidden_size,
VOC.num_words, hp.n_layers, hp.dropout)
encoder = encoder.to(device)
decoder = decoder.to(device)
print('Models built and ready to go!')
encoder.train()
decoder.train()
print('Building optimizers ...')
encoder_optimizer = optim.Adam(encoder.parameters(), lr=hp.lr)
decoder_optimizer = optim.Adam(decoder.parameters(),
lr=hp.lr * hp.decoder_learning_ratio)
encoder_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
encoder_optimizer, 5)
decoder_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
decoder_optimizer, 5)
if loadFilename:
def main():
corpus_name = "cornell movie-dialogs corpus"
corpus = os.path.join("data", corpus_name)
printLines(os.path.join(corpus, "movie_lines.txt"))
# Define path to new file
datafile = os.path.join(corpus, "formatted_movie_lines.txt")
linefile = os.path.join(corpus, "movie_lines.txt")
conversationfile = os.path.join(corpus, "movie_conversations.txt")
# Initialize lines dict, conversations list, and field ids
MOVIE_LINES_FIELDS = ["lineID", "characterID", "movieID", "character", "text"]
MOVIE_CONVERSATIONS_FIELDS = ["character1ID", "character2ID", "movieID", "utteranceIDs"]
# Load lines and process conversations
preprocess = Preprocess(datafile, linefile, conversationfile, MOVIE_LINES_FIELDS, MOVIE_CONVERSATIONS_FIELDS)
preprocess.loadLines()
preprocess.loadConversations()
preprocess.writeCSV()
# Load/Assemble voc and pairs
save_dir = os.path.join("data", "save")
dataset = Dataset(corpus, corpus_name, datafile)
voc, pairs = dataset.loadPrepareData()
# # Print some pairs to validate
# print("\npairs:")
# for pair in pairs[:10]:
# print(pair)
# Trim voc and pairs
pairs = dataset.trimRareWords(voc, pairs, MIN_COUNT)
# Example for validation
small_batch_size = 5
batches = dataset.batch2TrainData(voc, [random.choice(pairs) for _ in range(small_batch_size)])
input_variable, lengths, target_variable, mask, max_target_len = batches
print("input_variable:", input_variable)
print("lengths:", lengths)
print("target_variable:", target_variable)
print("mask:", mask)
print("max_target_len:", max_target_len)
# Configure models
model_name = 'cb_model'
attn_model = 'dot'
#attn_model = 'general'
#attn_model = 'concat'
hidden_size = 500
encoder_n_layers = 2
decoder_n_layers = 2
dropout = 0.1
batch_size = 64
# Set checkpoint to load from; set to None if starting from scratch
loadFilename = None
checkpoint_iter = 4000
#loadFilename = os.path.join(save_dir, model_name, corpus_name,
# '{}-{}_{}'.format(encoder_n_layers, decoder_n_layers, hidden_size),
# '{}_checkpoint.tar'.format(checkpoint_iter))
if loadFilename:
# If loading on same machine the model was trained on
checkpoint = torch.load(loadFilename)
# If loading a model trained on GPU to CPU
#checkpoint = torch.load(loadFilename, map_location=torch.device('cpu'))
encoder_sd = checkpoint['en']
decoder_sd = checkpoint['de']
encoder_optimizer_sd = checkpoint['en_opt']
decoder_optimizer_sd = checkpoint['de_opt']
embedding_sd = checkpoint['embedding']
voc.__dict__ = checkpoint['voc_dict']
print('Building encoder and decoder ...')
# Initialize word embeddings
embedding = nn.Embedding(voc.num_words, hidden_size)
if loadFilename:
embedding.load_state_dict(embedding_sd)
# Initialize encoder & decoder models
encoder = EncoderRNN(hidden_size, embedding, encoder_n_layers, dropout)
decoder = LuongAttnDecoderRNN(attn_model, embedding, hidden_size, voc.num_words, decoder_n_layers, dropout)
if loadFilename:
encoder.load_state_dict(encoder_sd)
decoder.load_state_dict(decoder_sd)
# Use appropriate device
encoder = encoder.to(device)
decoder = decoder.to(device)
print('Models built and ready to go!')
# Configure training/optimization
clip = 50.0
teacher_forcing_ratio = 1.0
learning_rate = 0.0001
decoder_learning_ratio = 5.0
n_iteration = 4000
print_every = 1
save_every = 500
# Ensure dropout layers are in train mode
encoder.train()
decoder.train()
# Initialize optimizers
print('Building optimizers ...')
encoder_optimizer = optim.Adam(encoder.parameters(), lr=learning_rate)
decoder_optimizer = optim.Adam(decoder.parameters(), lr=learning_rate * decoder_learning_ratio)
if loadFilename:
encoder_optimizer.load_state_dict(encoder_optimizer_sd)
decoder_optimizer.load_state_dict(decoder_optimizer_sd)
# Run training iterations
print("Starting Training!")
model = Model(dataset.batch2TrainData, teacher_forcing_ratio)
model.trainIters(model_name, voc, pairs, encoder, decoder, encoder_optimizer, decoder_optimizer,
embedding, encoder_n_layers, decoder_n_layers, save_dir, n_iteration, batch_size,
print_every, save_every, clip, corpus_name, loadFilename)
# Set dropout layers to eval mode
encoder.eval()
decoder.eval()
# Initialize search module
searcher = GreedySearchDecoder(encoder, decoder)
