def runTest(n_layers, hidden_size, reverse, modelFile, beam_size, inp, corpus):
torch.set_grad_enabled(False)
voc, pairs = loadPrepareData(corpus)
embedding = nn.Embedding(voc.num_words, hidden_size)
encoder = EncoderRNN(hidden_size, embedding, n_layers)
attn_model = 'dot'
decoder = LuongAttnDecoderRNN(attn_model, embedding, hidden_size,
voc.num_words, n_layers)
checkpoint = torch.load(modelFile,
map_location=lambda storage, loc: storage)
encoder.load_state_dict(checkpoint['en'])
decoder.load_state_dict(checkpoint['de'])
# train mode set to false, effect only on dropout, batchNorm
encoder.train(False)
decoder.train(False)
encoder = encoder.to(device)
decoder = decoder.to(device)
if inp:
evaluateInput(encoder, decoder, voc, beam_size)
else:
evaluateRandomly(encoder, decoder, voc, pairs, reverse, beam_size, 20)
def main():
data_path = './data/chatbot.txt'
voc, pairs = loadPrepareData(data_path)
# 把含有低频词的句子扔掉
MIN_COUNT = Config.MIN_COUNT
pairs = trimRareWords(voc, pairs, MIN_COUNT)
training_batches = [batch2TrainData(voc, [random.choice(pairs) for _ in range(Config.batch_size)])
for _ in range(Config.total_step)]
# 词嵌入部分
embedding = nn.Embedding(voc.num_words, Config.hidden_size)
# 定义编码解码器
encoder = EncoderRNN(Config.hidden_size, embedding, Config.encoder_n_layers, Config.dropout)
decoder = LuongAttnDecoderRNN(Config.attn_model, embedding, Config.hidden_size, voc.num_words, Config.decoder_n_layers, Config.dropout)
# 定义优化器
encoder_optimizer = optim.Adam(encoder.parameters(), lr=Config.learning_rate)
decoder_optimizer = optim.Adam(decoder.parameters(), lr=Config.learning_rate * Config.decoder_learning_ratio)
start_iteration = 1
save_every = 4000 # 多少步保存一次模型
for iteration in range(start_iteration, Config.total_step + 1):
training_batch = training_batches[iteration - 1]
input_variable, lengths, target_variable, mask, max_target_len = training_batch
start_time = time.time()
# Run a training iteration with batch
loss = train(input_variable, lengths, target_variable, mask, max_target_len, encoder,
decoder, embedding, encoder_optimizer, decoder_optimizer, Config.batch_size, Config.clip)
time_str = datetime.datetime.now().isoformat()
log_str = "time: {}, Iteration: {}; Percent complete: {:.1f}%; loss: {:.4f}, spend_time: {:6f}".format(time_str, iteration, iteration / Config.total_step * 100, loss, time.time() - start_time)
rainbow(log_str)
# Save checkpoint
if iteration % save_every == 0:
save_path = './save_model/'
if not os.path.exists(save_path):
os.makedirs(save_path)
torch.save({
'iteration': iteration,
'encoder': encoder.state_dict(),
'decoder': decoder.state_dict(),
'en_opt': encoder_optimizer.state_dict(),
'de_opt': decoder_optimizer.state_dict(),
'loss': loss,
'voc_dict': voc.__dict__,
'embedding': embedding.state_dict()
}, os.path.join(save_path, '{}_{}_model.tar'.format(iteration, 'checkpoint')))
def trainIters(attn_model, hidden_size,encoder_n_layers, decoder_n_layers, save_dir, n_iteration, batch_size, \
learning_rate, decoder_learning_ratio, print_every, save_every, clip, dropout, \
corpus_name, datafile, modelFile=None, need_trim=True):
# load train data
voc, pairs = loadPrepareData(datafile)
if need_trim:
# Trim voc and pairs
pairs = trimRareWords(voc, pairs, MIN_COUNT)
# Load batches for each iteration
training_batches = [
batch2TrainData(voc, [random.choice(pairs) for _ in range(batch_size)])
for _ in range(n_iteration)
]
if modelFile:
# If loading on same machine the model was trained on
checkpoint = torch.load(modelFile)
# 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']
embedding = nn.Embedding(voc.num_words, hidden_size)
if modelFile:
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)
# get model params
if modelFile:
encoder.load_state_dict(encoder_sd)
decoder.load_state_dict(decoder_sd)
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 modelFile:
encoder_optimizer.load_state_dict(encoder_optimizer_sd)
decoder_optimizer.load_state_dict(decoder_optimizer_sd)
# Initializations
print('Initializing ...')
start_iteration = 1
print_loss = 0
if modelFile:
start_iteration = checkpoint['iteration'] + 1
# Training loop
print("Training...")
encoder.train()
decoder.train()
for iteration in range(start_iteration, n_iteration + 1):
training_batch = training_batches[iteration - 1]
# Extract fields from batch
input_variable, lengths, target_variable, mask, max_target_len = training_batch
# Run a training iteration with batch
loss = train(input_variable, lengths, target_variable, mask,
max_target_len, encoder, decoder, encoder_optimizer,
decoder_optimizer, batch_size, clip)
print_loss += loss
# Print progress
if iteration % print_every == 0:
print_loss_avg = print_loss / print_every
print("Iteration: {}; Percent complete: {:.1f}%; Average loss: {:.4f}".format(iteration, \
iteration / n_iteration * 100, print_loss_avg))
print_loss = 0
# Save checkpoint
if (iteration % save_every == 0):
directory = os.path.join(save_dir, "model", '{}-{}_{}'.format(encoder_n_layers, \
decoder_n_layers, hidden_size))
if not os.path.exists(directory):
os.makedirs(directory)
torch.save(
{
'iteration': iteration,
'en': encoder.state_dict(),
'de': decoder.state_dict(),
'en_opt': encoder_optimizer.state_dict(),
'de_opt': decoder_optimizer.state_dict(),
'loss': loss,
'embedding': embedding.state_dict()
},
os.path.join(directory,
'{}_{}.tar'.format(iteration, 'checkpoint')))
# Evaluate sentence
output_words = evaluate(encoder, decoder, searcher, voc,
input_sentence)
# Format and print response sentence
output_words[:] = [
x for x in output_words if not (x == 'EOS' or x == 'PAD')
]
print('Bot:', ' '.join(output_words))
except KeyError:
print("Error: Encountered unknown word.")
if __name__ == '__main__':
data_path = './data/chatbot.txt'
voc, pairs = loadPrepareData(data_path)
# 把含有低频词的句子扔掉
MIN_COUNT = Config.MIN_COUNT
pairs = trimRareWords(voc, pairs, MIN_COUNT)
loadFilename = './save_model/' + '4000_checkpoint_model.tar'
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['encoder']
decoder_sd = checkpoint['decoder']
encoder_optimizer_sd = checkpoint['en_opt']
decoder_optimizer_sd = checkpoint['de_opt']
embedding_sd = checkpoint['embedding']