def train(mode, checkpoint_path): # Data data_train = IWSLT15EnViDataSet(en_path="../data/train-en-vi/train.en", vi_path="../data/train-en-vi/train.vi") data_loader = DataLoader(data_train, batch_size=BATCH_SIZE, shuffle=False, drop_last=False) if mode == EN2VI: src_vocab_size, tgt_vocab_size = data_train.en_vocab_size, data_train.vi_vocab_size else: src_vocab_size, tgt_vocab_size = data_train.vi_vocab_size, data_train.en_vocab_size print("Loading data done!") # Model & Optimizer model = NMT(mode=mode, src_vocab_size=src_vocab_size, tgt_vocab_size=tgt_vocab_size) model.to(device) criterion = MaskedPaddingCrossEntropyLoss().to(device) optimizer = Adam(model.parameters()) prev_epoch = 0 if checkpoint_path.exists(): # Resume training model, optimizer, prev_epoch = load_checkpoint(model, optimizer, checkpoint_path) print(f"Resume training from {prev_epoch} epochs!") else: model.apply(xavier_init_weights) print("Training from start!") model.train() for epoch in range(N_EPOCHS - prev_epoch): print(f"\nEpoch: {epoch+prev_epoch+1}") for b, (en_tokens, en_valid_len, vi_tokens, vi_valid_len) in enumerate(data_loader): en_tokens, vi_tokens = en_tokens.to(device), vi_tokens.to(device) en_valid_len, vi_valid_len = en_valid_len.to( device), vi_valid_len.to(device) en_padding_masks = mask_padding(en_tokens, en_valid_len, device) vi_padding_masks = mask_padding(vi_tokens, vi_valid_len, device) if mode == EN2VI: src, tgt = en_tokens, vi_tokens tgt_valid_len = vi_valid_len src_masks, tgt_masks = en_padding_masks, vi_padding_masks else: src, tgt = vi_tokens, en_tokens tgt_valid_len = en_valid_len src_masks, tgt_masks = vi_padding_masks, en_padding_masks optimizer.zero_grad() # Encoder's forward pass: encoder_state = model.encoder(src, src_masks) # Decoder's forward pass decoder_X = torch.tensor([[DEFAULT_SOS_INDEX] * tgt.shape[0]], device=device).reshape(-1, 1) decoder_state = encoder_state loss = torch.tensor(0, device=device, dtype=torch.float) for i in range(1, tgt.shape[1]): decoder_state, logit_pred = model.decoder( decoder_X, decoder_state) loss += criterion(pred=logit_pred[:, 0, :], label=tgt[:, i], device=device).sum() # Teacher forcing decoder_X = tgt[:, i].reshape(-1, 1) loss.backward() clip_grad_norm_(model.parameters(), max_norm=1.0) optimizer.step() if b % 50 == 0: seq_loss = loss / (MAX_LENGTH - 1) print(f"\tBatch {b}; Loss: {seq_loss:.2f}; " f"Mean Token Loss: {seq_loss/tgt_valid_len.sum():.4f}") ## Free up GPU memory del src, tgt, en_valid_len, vi_valid_len, decoder_state, logit_pred, loss torch.cuda.empty_cache() save_checkpoint(mode, src_vocab_size, tgt_vocab_size, model, optimizer, data_train.tokenizer_en, data_train.tokenizer_vi, prev_epoch + epoch + 1, checkpoint_path) for en in ens: vi = translate_en2vi(en_sentence=en, length=MAX_LENGTH, model=model, tokenizer_en=data_train.tokenizer_en, tokenizer_vi=data_train.tokenizer_vi, device=device) print("en:", en, "=> vi:", vi)
class Trainer: """ 训练类,使用训练集训练模型 Args: _hparams (NameSpace): 人为设定的超参数,默认值见config.py,也可以在命令行指定。 """ def __init__(self, _hparams): self.hparams = _hparams set_seed(_hparams.fixed_seed) self.train_loader = get_dataloader(_hparams.train_src_path, _hparams.train_dst_path, _hparams.batch_size, _hparams.num_workers) self.src_vocab, self.dst_vocab = load_vocab(_hparams.train_src_pkl, _hparams.train_dst_pkl) self.device = torch.device(_hparams.device) self.model = NMT(_hparams.embed_size, _hparams.hidden_size, self.src_vocab, self.dst_vocab, self.device, _hparams.dropout_rate).to(self.device) self.optimizer = torch.optim.Adam(self.model.parameters(), lr=_hparams.lr) def train(self): print('*' * 20, 'train', '*' * 20) hist_valid_scores = [] patience = 0 num_trial = 0 for epoch in range(int(self.hparams.max_epochs)): self.model.train() epoch_loss_val = 0 epoch_steps = len(self.train_loader) for step, data_pairs in tqdm(enumerate(self.train_loader)): sents = [(dp.src, dp.dst) for dp in data_pairs] src_sents, tgt_sents = zip(*sents) self.optimizer.zero_grad() batch_size = len(src_sents) example_losses = -self.model(src_sents, tgt_sents) batch_loss = example_losses.sum() train_loss = batch_loss / batch_size epoch_loss_val += train_loss.item() train_loss.backward() torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.hparams.clip_gradient) self.optimizer.step() epoch_loss_val /= epoch_steps print('epoch: {}, epoch_loss_val: {}'.format(epoch, epoch_loss_val)) # perform validation if epoch % self.hparams.valid_niter == 0: print('*' * 20, 'validate', '*' * 20) dev_ppl = evaluate_ppl(self.model, self.hparams.val_src_path, self.hparams.val_dst_path, self.hparams.batch_val_size, self.hparams.num_workers) valid_metric = -dev_ppl is_better = len(hist_valid_scores) == 0 or valid_metric > max(hist_valid_scores) hist_valid_scores.append(valid_metric) if is_better: patience = 0 print('save currently the best model to {}'.format(self.hparams.model_save_path)) self.model.save(self.hparams.model_save_path) torch.save(self.optimizer.state_dict(), self.hparams.optimizer_save_path) elif patience < self.hparams.patience: patience += 1 print('hit patience %d' % patience) if patience == self.hparams.patience: num_trial += 1 print('hit #{} trial'.format(num_trial)) if num_trial == self.hparams.max_num_trial: print('early stop!') exit(0) # 兼容设计,考虑Adam不需要人工调整lr,而其他优化器需要 if hasattr(self.optimizer, 'param_group'): # decay lr, and restore from previously best checkpoint lr = self.optimizer.param_groups[0]['lr'] * self.hparams.lr_decay print('load previously best model and decay learning rate to %f' % lr) params = torch.load(self.hparams.model_save_path, map_location=lambda storage, loc: storage) self.model.load_state_dict(params['state_dict']) self.model = self.model.to(self.device) print('restore parameters of the optimizers') self.optimizer.load_state_dict(torch.load(self.hparams.optimizer_save_path)) # set new lr for param_group in self.optimizer.param_groups: param_group['lr'] = lr # reset patience patience = 0 print('*' * 20, 'end validate', '*' * 20) print('*' * 20, 'end train', '*' * 20)
if use_w2v: w2v_train = pickle.load(open('./utils/w2vtrain.pkl', 'rb')) w2v_test = pickle.load(open('./utils/w2vtest.pkl', 'rb')) model = NMT(dim_in, dim_out, w2v_train=w2v_train, w2v_test=w2v_test, ues_attention=attention, ratio=teach_force).to(device) else: model = NMT(dim_in, dim_out, ues_attention=attention, ratio=teach_force).to(device) optimizer = Adam(model.parameters(), lr=LR, weight_decay=1e-4) weight = Calc_P(train_y, dim_out) Loss_fn = MyLoss(weight).to(device) timer = Timer(epoch_size) init_output_log(save_dir) print('model load done.') # get model train_len = len(train_x) dev_len = len(dev_x) train_x, train_y, train_mask = data_padding(train_x, train_y) dev_x, dev_y, dev_mask = data_padding(dev_x, dev_y) print('training start.')
def main(options): use_cuda = (len(options.gpuid) >= 1) if options.gpuid: cuda.set_device(options.gpuid[0]) src_train, src_dev, src_test, src_vocab = torch.load(open(options.data_file + "." + options.src_lang, 'rb')) trg_train, trg_dev, trg_test, trg_vocab = torch.load(open(options.data_file + "." + options.trg_lang, 'rb')) batched_train_src, batched_train_src_mask, sort_index = utils.tensor.advanced_batchize(src_train, options.batch_size, src_vocab.stoi["<blank>"]) batched_train_trg, batched_train_trg_mask = utils.tensor.advanced_batchize_no_sort(trg_train, options.batch_size, trg_vocab.stoi["<blank>"], sort_index) batched_dev_src, batched_dev_src_mask, sort_index = utils.tensor.advanced_batchize(src_dev, options.batch_size, src_vocab.stoi["<blank>"]) batched_dev_trg, batched_dev_trg_mask = utils.tensor.advanced_batchize_no_sort(trg_dev, options.batch_size, trg_vocab.stoi["<blank>"], sort_index) trg_vocab_size = len(trg_vocab) src_vocab_size = len(src_vocab) word_emb_size = 300 hidden_size = 1024 nmt = NMT(src_vocab_size, trg_vocab_size, word_emb_size, hidden_size, src_vocab, trg_vocab, attn_model = "general", use_cuda = True) if use_cuda > 0: nmt.cuda() if options.distributed: nmt = torch.nn.DataParallel(nmt) else: nmt.cpu() criterion = torch.nn.NLLLoss() # Configure optimization lr = options.learning_rate optimizer = eval("torch.optim." + options.optimizer)(nmt.parameters(), lr) # main training loop last_dev_avg_loss = float("inf") for epoch_i in range(options.epochs): logging.info("At {0}-th epoch.".format(epoch_i)) # Set training mode nmt.train() # srange generates a lazy sequence of shuffled range for i, batch_i in enumerate(utils.rand.srange(len(batched_train_src))): train_src_batch = Variable(batched_train_src[batch_i]) # of size (src_seq_len, batch_size) train_trg_batch = Variable(batched_train_trg[batch_i]) # of size (src_seq_len, batch_size) train_src_mask = Variable(batched_train_src_mask[batch_i]) train_trg_mask = Variable(batched_train_trg_mask[batch_i]) if use_cuda: train_src_batch = train_src_batch.cuda() train_trg_batch = train_trg_batch.cuda() train_src_mask = train_src_mask.cuda() train_trg_mask = train_trg_mask.cuda() sys_out_batch = nmt(train_src_batch, train_trg_batch, True) del train_src_batch train_trg_mask = train_trg_mask.view(-1) train_trg_batch = train_trg_batch.view(-1) train_trg_batch = train_trg_batch.masked_select(train_trg_mask) train_trg_mask = train_trg_mask.unsqueeze(1).expand(len(train_trg_mask), trg_vocab_size) sys_out_batch = sys_out_batch.view(-1, trg_vocab_size) sys_out_batch = sys_out_batch.masked_select(train_trg_mask).view(-1, trg_vocab_size) loss = criterion(sys_out_batch, train_trg_batch) logging.debug("loss at batch {0}: {1}".format(i, loss.data[0])) optimizer.zero_grad() loss.backward() # # gradient clipping torch.nn.utils.clip_grad_norm(nmt.parameters(), 5.0) optimizer.step() # validation -- this is a crude esitmation because there might be some paddings at the end dev_loss = 0.0 # Set validation mode nmt.eval() for batch_i in range(len(batched_dev_src)): dev_src_batch = Variable(batched_dev_src[batch_i], volatile=True) dev_trg_batch = Variable(batched_dev_trg[batch_i], volatile=True) dev_src_mask = Variable(batched_dev_src_mask[batch_i], volatile=True) dev_trg_mask = Variable(batched_dev_trg_mask[batch_i], volatile=True) if use_cuda: dev_src_batch = dev_src_batch.cuda() dev_trg_batch = dev_trg_batch.cuda() dev_src_mask = dev_src_mask.cuda() dev_trg_mask = dev_trg_mask.cuda() sys_out_batch = nmt(dev_src_batch, dev_trg_batch, False) dev_trg_mask = dev_trg_mask.view(-1) dev_trg_batch = dev_trg_batch.view(-1) dev_trg_batch = dev_trg_batch.masked_select(dev_trg_mask) dev_trg_mask = dev_trg_mask.unsqueeze(1).expand(len(dev_trg_mask), trg_vocab_size) sys_out_batch = sys_out_batch.view(-1, trg_vocab_size) sys_out_batch = sys_out_batch.masked_select(dev_trg_mask).view(-1, trg_vocab_size) loss = criterion(sys_out_batch, dev_trg_batch) logging.debug("dev loss at batch {0}: {1}".format(batch_i, loss.data[0])) dev_loss += loss dev_avg_loss = dev_loss / len(batched_dev_src) logging.info("Average loss value per instance is {0} at the end of epoch {1}".format(dev_avg_loss.data[0], epoch_i)) # if (last_dev_avg_loss - dev_avg_loss).data[0] < options.estop: # logging.info("Early stopping triggered with threshold {0} (previous dev loss: {1}, current: {2})".format(epoch_i, last_dev_avg_loss.data[0], dev_avg_loss.data[0])) # break torch.save(nmt, open(options.model_file + ".nll_{0:.2f}.epoch_{1}".format(dev_avg_loss.data[0], epoch_i), 'wb'), pickle_module=dill) last_dev_avg_loss = dev_avg_loss
def train(args: Dict): train_data_src = read_corpus(args['--train-src'], source='src') train_data_tgt = read_corpus(args['--train-tgt'], source='tgt') dev_data_src = read_corpus(args['--dev-src'], source='src') dev_data_tgt = read_corpus(args['--dev-tgt'], source='tgt') # [(src_0, tgt_0), (src_1, tgt_1), ..., ] train_data = list(zip(train_data_src, train_data_tgt)) dev_data = list(zip(dev_data_src, dev_data_tgt)) train_batch_size = int(args['--batch-size']) clip_grad = float(args['--clip-grad']) valid_niter = int(args['--valid-niter']) log_every = int(args['--log-every']) model_save_path = args['--save-to'] # vocab = Vocab.load(args['--vocab']) vocab = Vocab.build(train_data_src, train_data_tgt, int(args['--vocab-size']), 1) model = NMT(embed_size=int(args['--embed-size']), hidden_size=int(args['--hidden-size']), dropout_rate=float(args['--dropout']), vocab=vocab) model.train() print(model) uniform_init = float(args['--uniform-init']) if np.abs(uniform_init) > 0.: print('uniformly initialize parameters [-%f, +%f]' % (uniform_init, uniform_init), file=sys.stderr) for p in model.parameters(): p.data.uniform_(-uniform_init, uniform_init) # vocab_mask = torch.ones(len(vocab.tgt)) # vocab_mask[vocab.tgt['<pad>']] = 0 device = torch.device("cuda:0" if args['--cuda'] else "cpu") print('use device: %s' % device, file=sys.stderr) model = model.to(device) model.save(model_save_path) optimizer = torch.optim.Adam(model.parameters(), lr=float(args['--lr'])) num_trial = 0 train_iter = patience = cum_loss = report_loss = cum_tgt_words = report_tgt_words = 0 cum_examples = report_examples = epoch = valid_num = 0 hist_valid_scores = [] train_time = begin_time = time.time() print('begin Maximum Likelihood training') while True: epoch += 1 for src_sents, tgt_sents in batch_iter(train_data, batch_size=train_batch_size, shuffle=True): train_iter += 1 optimizer.zero_grad() batch_size = len(src_sents) #################### forward pass and compute loss ######################### # example_losses = -model(src_sents, tgt_sents) # (batch_size,) example_losses = model(src_sents, tgt_sents) # [batch_size,] batch_loss = example_losses.sum() loss = batch_loss / batch_size #################### backward pass to compute gradients #################### loss.backward() # clip gradient grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), clip_grad) #################### update model parameters ############################### optimizer.step() #################### do some statistics #################################### batch_losses_val = batch_loss.item() report_loss += batch_losses_val cum_loss += batch_losses_val tgt_words_num_to_predict = sum( len(s[1:]) for s in tgt_sents) # omitting leading `<s>` report_tgt_words += tgt_words_num_to_predict cum_tgt_words += tgt_words_num_to_predict report_examples += batch_size cum_examples += batch_size #################### print log ############################################# if train_iter % log_every == 0: print( 'epoch %d, iter %d, avg. loss %.2f, avg. ppl %.2f ' 'cum. examples %d, speed %.2f words/sec, time elapsed %.2f sec' % ( epoch, train_iter, report_loss / report_examples, # math.exp(report_loss / report_tgt_words), (report_loss / report_tgt_words), cum_examples, report_tgt_words / (time.time() - train_time), time.time() - begin_time), file=sys.stderr) train_time = time.time() report_loss = report_tgt_words = report_examples = 0. ##################### perform validation ################################## if train_iter % valid_niter == 0: print( 'epoch %d, iter %d, cum. loss %.2f, cum. ppl %.2f cum. examples %d' % (epoch, train_iter, cum_loss / cum_examples, np.exp(cum_loss / cum_tgt_words), cum_examples), file=sys.stderr) cum_loss = cum_examples = cum_tgt_words = 0. valid_num += 1 print('begin validation ...', file=sys.stderr) # compute dev. ppl and bleu dev_ppl = evaluate_ppl( model, dev_data, batch_size=128) # dev batch size can be a bit larger valid_metric = -dev_ppl print('validation: iter %d, dev. ppl %f' % (train_iter, dev_ppl), file=sys.stderr) is_better = len(hist_valid_scores ) == 0 or valid_metric > max(hist_valid_scores) hist_valid_scores.append(valid_metric) # hypotheses = beam_search(model, dev_data_src, # beam_size=4, # max_decoding_time_step=10) if is_better: patience = 0 print('save currently the best model to [%s]' % model_save_path, file=sys.stderr) model.save(model_save_path) # also save the optimizers' state torch.save(optimizer.state_dict(), model_save_path + '.optim') elif patience < int(args['--patience']): patience += 1 print('hit patience %d' % patience, file=sys.stderr) if patience == int(args['--patience']): num_trial += 1 print('hit #%d trial' % num_trial, file=sys.stderr) if num_trial == int(args['--max-num-trial']): print('early stop!', file=sys.stderr) exit(0) # decay lr, and restore from previously best checkpoint lr = optimizer.param_groups[0]['lr'] * float( args['--lr-decay']) print( 'load previously best model and decay learning rate to %f' % lr, file=sys.stderr) # load model params = torch.load( model_save_path, map_location=lambda storage, loc: storage) model.load_state_dict(params['state_dict']) model = model.to(device) print('restore parameters of the optimizers', file=sys.stderr) optimizer.load_state_dict( torch.load(model_save_path + '.optim')) # set new lr for param_group in optimizer.param_groups: param_group['lr'] = lr # reset patience patience = 0 if epoch == int(args['--max-epoch']): print('reached maximum number of epochs!', file=sys.stderr) exit(0)
def main(options): original_model = torch.load(open(options.original_model_file, 'rb')) nmt = NMT(original_model) use_cuda = (len(options.gpuid) >= 1) if options.gpuid: cuda.set_device(options.gpuid[0]) src_train, src_dev, src_test, src_vocab = torch.load( open(options.data_file + "." + options.src_lang, 'rb')) trg_train, trg_dev, trg_test, trg_vocab = torch.load( open(options.data_file + "." + options.trg_lang, 'rb')) batched_test_src, batched_test_src_mask, _ = utils.tensor.advanced_batchize( src_test, 1, src_vocab.stoi["<pad>"]) batched_test_trg, batched_test_trg_mask, _ = utils.tensor.advanced_batchize( trg_test, 1, trg_vocab.stoi["<pad>"]) batched_train_src, batched_train_src_mask, _ = utils.tensor.advanced_batchize( src_train, 1, src_vocab.stoi["<pad>"]) batched_train_trg, batched_train_trg_mask, _ = utils.tensor.advanced_batchize( trg_train, 1, trg_vocab.stoi["<pad>"]) batched_dev_src, batched_dev_src_mask, _ = utils.tensor.advanced_batchize( src_dev, options.batch_size, src_vocab.stoi["<pad>"]) batched_dev_trg, batched_dev_trg_mask, _ = utils.tensor.advanced_batchize( trg_dev, options.batch_size, trg_vocab.stoi["<pad>"]) trg_vocab_size = len(trg_vocab) src_vocab_size = len(src_vocab) # for i, batch_i in enumerate(utils.rand.srange(len(batched_train_src))): # train_src_batch = Variable(batched_train_src[batch_i]) # of size (src_seq_len, batch_size) # train_trg_batch = Variable(batched_train_trg[batch_i]) # of size (src_seq_len, batch_size) # train_src_mask = Variable(batched_train_src_mask[batch_i]) # train_trg_mask = Variable(batched_train_trg_mask[batch_i]) # if use_cuda: # train_src_batch = train_src_batch.cuda() # train_trg_batch = train_trg_batch.cuda() # train_src_mask = train_src_mask.cuda() # train_trg_mask = train_trg_mask.cuda() # sys_out_batch = nmt(train_src_batch, train_trg_batch.size()[0]) # # print(sys_out_batch.size()) # _, sys_out_batch = torch.max(sys_out_batch, dim=2) # sys_out_batch = sys_out_batch.view(-1) # sent = [] # # print(sys_out_batch) # for w in sys_out_batch: # # print(w) # sent.append(trg_vocab.itos[w.data[0]]) # print(sent) # # Initialize encoder with weights parameters from original model # encoder = nn.LSTM(300, 512, bidirectional=True) # encoder.weight_ih_l0 = nn.Parameter(original_model['encoder.rnn.weight_ih_l0']) # encoder.weight_hh_l0 = nn.Parameter(original_model['encoder.rnn.weight_hh_l0']) # encoder.bias_ih_l0 = nn.Parameter(original_model['encoder.rnn.bias_ih_l0']) # encoder.bias_hh_l0 = nn.Parameter(original_model['encoder.rnn.bias_hh_l0']) # encoder.weight_ih_l0_reverse = nn.Parameter(original_model['encoder.rnn.weight_ih_l0_reverse']) # encoder.weight_hh_l0_reverse = nn.Parameter(original_model['encoder.rnn.weight_hh_l0_reverse']) # encoder.bias_ih_l0_reverse = nn.Parameter(original_model['encoder.rnn.bias_ih_l0_reverse']) # encoder.bias_hh_l0_reverse = nn.Parameter(original_model['encoder.rnn.bias_hh_l0_reverse']) # # Initialize decoder with weights parameters from original model # decoder = nn.LSTM(1324, 1024) # decoder.weight_ih_l0 = nn.Parameter(original_model['decoder.rnn.layers.0.weight_ih']) # decoder.weight_hh_l0 = nn.Parameter(original_model['decoder.rnn.layers.0.weight_hh']) # decoder.bias_ih_l0 = nn.Parameter(original_model['decoder.rnn.layers.0.bias_ih']) # decoder.bias_hh_l0 = nn.Parameter(original_model['decoder.rnn.layers.0.bias_hh']) # if use_cuda > 0: # encoder.cuda() # decoder.cuda() # else: # encoder.cpu() # decoder.cpu() # # Initialize embeddings # encoder_embedding = nn.Embedding(36616, 300) # decoder_embedding = nn.Embedding(23262, 300) # encoder_embedding.weight = nn.Parameter(original_model['encoder.embeddings.emb_luts.0.weight']) # decoder_embedding.weight = nn.Parameter(original_model['decoder.embeddings.emb_luts.0.weight']) # # Initialize Ws # wi = nn.Linear(1024,1024, bias=False) # wi.weight = nn.Parameter(original_model['decoder.attn.linear_in.weight']) # wo = nn.Linear(2048, 1024, bias=False) # wo.weight = nn.Parameter(original_model['decoder.attn.linear_out.weight']) # generator = nn.Linear(1024, 23262) # generator.weight = nn.Parameter(original_model['0.weight']) # generator.bias = nn.Parameter(original_model['0.bias']) criterion = torch.nn.NLLLoss() # encoder_optimizer = eval("torch.optim." + options.optimizer)(encoder.parameters(), options.learning_rate) # decoder_optimizer = eval("torch.optim." + options.optimizer)(decoder.parameters(), options.learning_rate) # soft_max = nn.Softmax() optimizer = eval("torch.optim." + options.optimizer)(nmt.parameters(), options.learning_rate) # main training loop last_dev_avg_loss = float("inf") for epoch_i in range(options.epochs): logging.info("At {0}-th epoch.".format(epoch_i)) h_t_1 = Variable(torch.ones(1024)) # srange generates a lazy sequence of shuffled range for i, batch_i in enumerate(utils.rand.srange(len(batched_train_src))): train_src_batch = Variable(batched_train_src[batch_i] ) # of size (src_seq_len, batch_size) train_trg_batch = Variable(batched_train_trg[batch_i] ) # of size (src_seq_len, batch_size) train_src_mask = Variable(batched_train_src_mask[batch_i]) train_trg_mask = Variable(batched_train_trg_mask[batch_i]) if use_cuda: train_src_batch = train_src_batch.cuda() train_trg_batch = train_trg_batch.cuda() train_src_mask = train_src_mask.cuda() train_trg_mask = train_trg_mask.cuda() # encoder_input = encoder_embedding(train_trg_batch) # sys_out_batch, (encoder_hidden_states, _) = encoder(encoder_input) # (trg_seq_len, batch_size, trg_vocab_size) # TODO: add more arguments as necessary # h = Variable(torch.FloatTensor(sys_out_batch.size()[1], 1024).fill_(1./1024)) # c = Variable(torch.FloatTensor(sys_out_batch.size()[1], 1024).fill_(0)) # softmax = torch.nn.Softmax() # tanh = torch.nn.Tanh() # # _,w = torch.max(softmax(generator(h)), dim=1) # w = softmax(generator(h)) # result = Variable(torch.FloatTensor(sys_out_batch.size()[0], sys_out_batch.size()[1], 23262)) # for i in range(sys_out_batch.size()[0]): # wht1 = wi(h).view(1, -1, 1024).expand_as(sys_out_batch) # score = softmax(torch.sum(sys_out_batch * wht1, dim=2)).view(sys_out_batch.size()[0],sys_out_batch.size()[1],1) # st = torch.sum(score * sys_out_batch, dim=0) # ct = tanh(wo(torch.cat([st, h], dim=1))) # _, w = torch.max(w, dim=1) # input = torch.cat([decoder_embedding(w), ct], dim=1) # input = input.view(1, input.size()[0], input.size()[1]) # _,(b,c) = decoder(input, (h,c)) # h = b[0] # c = c[0] # w = softmax(generator(h)) # result[i] = w # # result.append(w) # sys_out_batch = result sys_out_batch = nmt(train_src_batch, train_trg_batch.size()[0]) # s_vector = [] # for hs in sys_out_batch: # score = hs.matmul(wi).matmul(h_t_1) # score = score.unsqueeze(0) # a_h_s = soft_max(score) # # print a_h_s, hs.squeeze(0) # s_vector.append(a_h_s.squeeze(0).dot(hs.squeeze(0))) # s_tilda = sum(s_vector) # c_t = nn.Tanh(wo.matmul(torch.cat(s_tilda, h_t_1))) # sys.exit() # train_trg_mask = train_trg_mask.view(-1) # train_trg_batch = train_trg_batch.view(-1) # train_trg_batch = train_trg_batch.masked_select(train_trg_mask) # train_trg_mask = train_trg_mask.unsqueeze(1).expand(len(train_trg_mask), trg_vocab_size) # sys_out_batch = sys_out_batch.view(-1, trg_vocab_size) # sys_out_batch = sys_out_batch.masked_select(train_trg_mask).view(-1, trg_vocab_size) print(train_trg_mask.size()) train_trg_mask = train_trg_mask.view(-1) train_trg_batch = train_trg_batch.view(-1) train_trg_batch = train_trg_batch.masked_select(train_trg_mask) train_trg_mask = train_trg_mask.unsqueeze(1).expand( len(train_trg_mask), trg_vocab_size - 1) # print(trainin.size()) # print(train_trg_batch[:,:-1].size()) sys_out_batch = sys_out_batch.view(-1, trg_vocab_size - 1) print(trg_vocab_size) print(train_trg_mask.size()) sys_out_batch = sys_out_batch.masked_select(train_trg_mask).view( -1, trg_vocab_size - 1) print(sys_out_batch.size()) loss = criterion(sys_out_batch, train_trg_batch) logging.debug("loss at batch {0}: {1}".format(i, loss.data[0])) optimizer.zero_grad() loss.backward() optimizer.step() # validation -- this is a crude esitmation because there might be some paddings at the end dev_loss = 0.0 for batch_i in range(len(batched_dev_src)): dev_src_batch = Variable(batched_dev_src[batch_i], volatile=True) dev_trg_batch = Variable(batched_dev_trg[batch_i], volatile=True) dev_src_mask = Variable(batched_dev_src_mask[batch_i], volatile=True) dev_trg_mask = Variable(batched_dev_trg_mask[batch_i], volatile=True) if use_cuda: dev_src_batch = dev_src_batch.cuda() dev_trg_batch = dev_trg_batch.cuda() dev_src_mask = dev_src_mask.cuda() dev_trg_mask = dev_trg_mask.cuda() # encoder_input = encoder_embedding(dev_trg_batch) # sys_out_batch = encoder(encoder_input) # (trg_seq_len, batch_size, trg_vocab_size) # TODO: add more arguments as necessary sys_out_batch = nmt(dev_src_batch, dev_trg_batch.size()[0]) dev_trg_mask = dev_trg_mask.view(-1) dev_trg_batch = dev_trg_batch.view(-1) dev_trg_batch = dev_trg_batch.masked_select(dev_trg_mask) dev_trg_mask = dev_trg_mask.unsqueeze(1).expand( len(dev_trg_mask), trg_vocab_size - 1) sys_out_batch = sys_out_batch.view(-1, trg_vocab_size - 1) sys_out_batch = sys_out_batch.masked_select(dev_trg_mask).view( -1, trg_vocab_size - 1) loss = criterion(sys_out_batch, dev_trg_batch) logging.debug("dev loss at batch {0}: {1}".format( batch_i, loss.data[0])) dev_loss += loss dev_avg_loss = dev_loss / len(batched_dev_in) logging.info( "Average loss value per instance is {0} at the end of epoch {1}". format(dev_avg_loss.data[0], epoch_i)) if (last_dev_avg_loss - dev_avg_loss).data[0] < options.estop: logging.info( "Early stopping triggered with threshold {0} (previous dev loss: {1}, current: {2})" .format(epoch_i, last_dev_avg_loss.data[0], dev_avg_loss.data[0])) break torch.save( nmt.state_dict(), open( options.model_file + ".nll_{0:.2f}.epoch_{1}".format(dev_avg_loss.data[0], epoch_i), 'wb'), pickle_module=dill) last_dev_avg_loss = dev_avg_loss
def main(options): use_cuda = (len(options.gpuid) >= 1) # if options.gpuid: # cuda.set_device(options.gpuid[0]) src_train, src_dev, src_test, src_vocab = torch.load( open(options.data_file + "." + options.src_lang, 'rb')) trg_train, trg_dev, trg_test, trg_vocab = torch.load( open(options.data_file + "." + options.trg_lang, 'rb')) batched_train_src, batched_train_src_mask, sort_index = utils.tensor.advanced_batchize( src_train, options.batch_size, src_vocab.stoi["<blank>"]) batched_train_trg, batched_train_trg_mask = utils.tensor.advanced_batchize_no_sort( trg_train, options.batch_size, trg_vocab.stoi["<blank>"], sort_index) batched_dev_src, batched_dev_src_mask, sort_index = utils.tensor.advanced_batchize( src_dev, options.batch_size, src_vocab.stoi["<blank>"]) batched_dev_trg, batched_dev_trg_mask = utils.tensor.advanced_batchize_no_sort( trg_dev, options.batch_size, trg_vocab.stoi["<blank>"], sort_index) print "preprocessing batched data..." processed_src = list() processed_trg = list() processed_src_mask = list() processed_trg_mask = list() for batch_i in range(len(batched_train_src)): if batched_train_src[batch_i].size( 0) <= 35 and batched_train_trg[batch_i].size(0) <= 35: processed_src.append(batched_train_src[batch_i]) processed_trg.append(batched_train_trg[batch_i]) processed_src_mask.append(batched_train_src_mask[batch_i]) processed_trg_mask.append(batched_train_trg_mask[batch_i]) batched_train_src = processed_src batched_train_trg = processed_trg batched_train_src_mask = processed_src_mask batched_train_trg_mask = processed_trg_mask processed_src = list() processed_trg = list() processed_src_mask = list() processed_trg_mask = list() for batch_i in range(len(batched_dev_src)): if batched_dev_src[batch_i].size( 0) <= 35 and batched_dev_trg[batch_i].size(0) <= 35: processed_src.append(batched_dev_src[batch_i]) processed_trg.append(batched_dev_trg[batch_i]) processed_src_mask.append(batched_dev_src_mask[batch_i]) processed_trg_mask.append(batched_dev_trg_mask[batch_i]) batched_dev_src = processed_src batched_dev_trg = processed_trg batched_dev_src_mask = processed_src_mask batched_dev_trg_mask = processed_trg_mask del processed_src, processed_trg, processed_trg_mask, processed_src_mask trg_vocab_size = len(trg_vocab) src_vocab_size = len(src_vocab) word_emb_size = 50 hidden_size = 1024 nmt = NMT(src_vocab_size, trg_vocab_size, word_emb_size, hidden_size, src_vocab, trg_vocab, attn_model="general", use_cuda=True) discriminator = Discriminator(src_vocab_size, trg_vocab_size, word_emb_size, src_vocab, trg_vocab, use_cuda=True) if use_cuda > 0: #nmt = torch.nn.DataParallel(nmt,device_ids=options.gpuid).cuda() nmt.cuda() #discriminator = torch.nn.DataParallel(discriminator,device_ids=options.gpuid).cuda() discriminator.cuda() else: nmt.cpu() discriminator.cpu() criterion_g = torch.nn.NLLLoss().cuda() criterion = torch.nn.CrossEntropyLoss().cuda() # Configure optimization optimizer_g = eval("torch.optim." + options.optimizer)( nmt.parameters(), options.learning_rate) optimizer_d = eval("torch.optim." + options.optimizer)( discriminator.parameters(), options.learning_rate) # main training loop f1 = open("train_loss", "a") f2 = open("dev_loss", "a") last_dev_avg_loss = float("inf") for epoch_i in range(options.epochs): logging.info("At {0}-th epoch.".format(epoch_i)) # srange generates a lazy sequence of shuffled range train_loss_g = 0.0 train_loss_d = 0.0 train_loss_g_nll = 0.0 train_loss_g_ce = 0.0 train_loss_nll_batch_num = 0 train_loss_ce_batch_num = 0 for i, batch_i in enumerate(utils.rand.srange(len(batched_train_src))): if i == 1500: break # if i==5: # break train_src_batch = Variable(batched_train_src[batch_i] ) # of size (src_seq_len, batch_size) train_trg_batch = Variable(batched_train_trg[batch_i] ) # of size (src_seq_len, batch_size) train_src_mask = Variable(batched_train_src_mask[batch_i]) train_trg_mask = Variable(batched_train_trg_mask[batch_i]) if use_cuda: train_src_batch = train_src_batch.cuda() train_trg_batch = train_trg_batch.cuda() train_src_mask = train_src_mask.cuda() train_trg_mask = train_trg_mask.cuda() # train discriminator sys_out_batch = nmt(train_src_batch, train_trg_batch, True).detach() _, predict_batch = sys_out_batch.topk(1) del _ predict_batch = predict_batch.squeeze(2) real_dis_label_out = discriminator(train_src_batch, train_trg_batch, True) fake_dis_label_out = discriminator(train_src_batch, predict_batch, True) optimizer_d.zero_grad() loss_d_real = criterion( real_dis_label_out, Variable( torch.ones(options.batch_size * len(options.gpuid)).long()).cuda()) loss_d_real.backward() loss_d_fake = criterion( fake_dis_label_out, Variable( torch.zeros(options.batch_size * len(options.gpuid)).long()).cuda()) #loss_d_fake.backward(retain_graph=True) loss_d_fake.backward() loss_d = loss_d_fake.data[0] + loss_d_real.data[0] del loss_d_fake, loss_d_real logging.debug("D loss at batch {0}: {1}".format(i, loss_d)) f1.write("D train loss at batch {0}: {1}\n".format(i, loss_d)) optimizer_d.step() if use_cuda > 0: sys_out_batch = sys_out_batch.cuda() train_trg_batch = train_trg_batch.cuda() else: sys_out_batch = sys_out_batch.cpu() train_trg_batch = train_trg_batch.cpu() # train nmt sys_out_batch = nmt(train_src_batch, train_trg_batch, True) _, predict_batch = sys_out_batch.topk(1) predict_batch = predict_batch.squeeze(2) fake_dis_label_out = discriminator(train_src_batch, predict_batch, True) if random.random() > 0.5: train_trg_mask = train_trg_mask.view(-1) train_trg_batch = train_trg_batch.view(-1) train_trg_batch = train_trg_batch.masked_select(train_trg_mask) train_trg_mask = train_trg_mask.unsqueeze(1).expand( len(train_trg_mask), trg_vocab_size) sys_out_batch = sys_out_batch.view(-1, trg_vocab_size) sys_out_batch = sys_out_batch.masked_select( train_trg_mask).view(-1, trg_vocab_size) loss_g = criterion_g(sys_out_batch, train_trg_batch) train_loss_g_nll += loss_g train_loss_nll_batch_num += 1 f1.write("G train NLL loss at batch {0}: {1}\n".format( i, loss_g.data[0])) else: loss_g = criterion( fake_dis_label_out, Variable( torch.ones(options.batch_size * len(options.gpuid)).long()).cuda()) train_loss_g_ce += loss_g train_loss_ce_batch_num += 1 f1.write("G train CE loss at batch {0}: {1}\n".format( i, loss_g.data[0])) logging.debug("G loss at batch {0}: {1}".format(i, loss_g.data[0])) optimizer_g.zero_grad() loss_g.backward() # # gradient clipping torch.nn.utils.clip_grad_norm(nmt.parameters(), 5.0) optimizer_g.step() train_loss_d += loss_d train_avg_loss_g_nll = train_loss_g_nll / train_loss_nll_batch_num train_avg_loss_g_ce = train_loss_g_ce / train_loss_ce_batch_num train_avg_loss_d = train_loss_d / len(train_src_batch) logging.info( "G TRAIN Average NLL loss value per instance is {0} at the end of epoch {1}" .format(train_avg_loss_g_nll, epoch_i)) logging.info( "G TRAIN Average CE loss value per instance is {0} at the end of epoch {1}" .format(train_avg_loss_g_ce, epoch_i)) logging.info( "D TRAIN Average loss value per instance is {0} at the end of epoch {1}" .format(train_avg_loss_d, epoch_i)) # validation -- this is a crude esitmation because there might be some paddings at the end # dev_loss_g_nll = 0.0 # dev_loss_g_ce = 0.0 # dev_loss_d = 0.0 # for batch_i in range(len(batched_dev_src)): # dev_src_batch = Variable(batched_dev_src[batch_i], volatile=True) # dev_trg_batch = Variable(batched_dev_trg[batch_i], volatile=True) # dev_src_mask = Variable(batched_dev_src_mask[batch_i], volatile=True) # dev_trg_mask = Variable(batched_dev_trg_mask[batch_i], volatile=True) # if use_cuda: # dev_src_batch = dev_src_batch.cuda() # dev_trg_batch = dev_trg_batch.cuda() # dev_src_mask = dev_src_mask.cuda() # dev_trg_mask = dev_trg_mask.cuda() # sys_out_batch = nmt(dev_src_batch, dev_trg_batch, False).detach() # _,predict_batch = sys_out_batch.topk(1) # predict_batch = predict_batch.squeeze(2) # real_dis_label_out = discriminator(dev_src_batch, dev_trg_batch, True).detach() # fake_dis_label_out = discriminator(dev_src_batch, predict_batch, True).detach() # if use_cuda > 0: # sys_out_batch = sys_out_batch.cuda() # dev_trg_batch = dev_trg_batch.cuda() # else: # sys_out_batch = sys_out_batch.cpu() # dev_trg_batch = dev_trg_batch.cpu() # dev_trg_mask = dev_trg_mask.view(-1) # dev_trg_batch = dev_trg_batch.view(-1) # dev_trg_batch = dev_trg_batch.masked_select(dev_trg_mask) # dev_trg_mask = dev_trg_mask.unsqueeze(1).expand(len(dev_trg_mask), trg_vocab_size) # sys_out_batch = sys_out_batch.view(-1, trg_vocab_size) # sys_out_batch = sys_out_batch.masked_select(dev_trg_mask).view(-1, trg_vocab_size) # loss_g_nll = criterion_g(sys_out_batch, dev_trg_batch) # loss_g_ce = criterion(fake_dis_label_out, Variable(torch.ones(options.batch_size*len(options.gpuid)).long(),volatile=True).cuda()) # loss_d = criterion(real_dis_label_out, Variable(torch.ones(options.batch_size*len(options.gpuid)).long(),volatile=True).cuda()) + criterion(fake_dis_label_out, Variable(torch.zeros(options.batch_size*len(options.gpuid)).long(),volatile=True).cuda()) # logging.debug("G dev NLL loss at batch {0}: {1}".format(batch_i, loss_g_nll.data[0])) # logging.debug("G dev CE loss at batch {0}: {1}".format(batch_i, loss_g_ce.data[0])) # f2.write("G dev NLL loss at batch {0}: {1}\n".format(batch_i, loss_g_nll.data[0])) # f2.write("G dev CE loss at batch {0}: {1}\n".format(batch_i, loss_g_ce.data[0])) # logging.debug("D dev loss at batch {0}: {1}".format(batch_i, loss_d.data[0])) # f2.write("D dev loss at batch {0}: {1}\n".format(batch_i, loss_d.data[0])) # dev_loss_g_nll += loss_g_nll # dev_loss_g_ce += loss_g_ce # dev_loss_d += loss_d # dev_avg_loss_g_nll = dev_loss_g_nll / len(batched_dev_src) # dev_avg_loss_g_ce = dev_loss_g_ce / len(batched_dev_src) # dev_avg_loss_d = dev_loss_d / len(batched_dev_src) # logging.info("G DEV Average NLL loss value per instance is {0} at the end of epoch {1}".format(dev_avg_loss_g_nll.cpu().data[0], epoch_i)) # logging.info("G DEV Average CE loss value per instance is {0} at the end of epoch {1}".format(dev_avg_loss_g_ce.cpu().data[0], epoch_i)) # logging.info("D DEV Average loss value per instance is {0} at the end of epoch {1}".format(dev_avg_loss_d.data[0], epoch_i)) # # if (last_dev_avg_loss - dev_avg_loss).data[0] < options.estop: # # logging.info("Early stopping triggered with threshold {0} (previous dev loss: {1}, current: {2})".format(epoch_i, last_dev_avg_loss.data[0], dev_avg_loss.data[0])) # # break torch.save(nmt, open( "nmt.nll_{0:.2f}.epoch_{1}".format( train_avg_loss_g_nll.cpu().data[0], epoch_i), 'wb'), pickle_module=dill) torch.save(discriminator, open( "discriminator.nll_{0:.2f}.epoch_{1}".format( train_avg_loss_d.data[0], epoch_i), 'wb'), pickle_module=dill) f1.close() f2.close()
def main(options): use_cuda = (len(options.gpuid) >= 1) if options.gpuid: cuda.set_device(options.gpuid[0]) src_train, src_dev, src_test, src_vocab = torch.load( open(options.data_file + "." + options.src_lang, 'rb')) trg_train, trg_dev, trg_test, trg_vocab = torch.load( open(options.data_file + "." + options.trg_lang, 'rb')) """src_train = get_lm_input(src_train) src_dev = get_lm_input(src_dev) src_test = get_lm_input(src_test) trg_train = get_lm_output(trg_train) trg_dev = get_lm_output(trg_dev) trg_test = get_lm_output(trg_test)""" batched_train_src, batched_train_src_mask, sort_index = tensor.advanced_batchize( src_train, options.batch_size, src_vocab.stoi["<blank>"]) batched_train_trg, batched_train_trg_mask = tensor.advanced_batchize_no_sort( trg_train, options.batch_size, trg_vocab.stoi["<blank>"], sort_index) batched_dev_src, batched_dev_src_mask, sort_index = tensor.advanced_batchize( src_dev, options.batch_size, src_vocab.stoi["<blank>"]) batched_dev_trg, batched_dev_trg_mask = tensor.advanced_batchize_no_sort( trg_dev, options.batch_size, trg_vocab.stoi["<blank>"], sort_index) (src_vocab_size, trg_vocab_size) = len(src_vocab), len(trg_vocab) nmt = NMT((src_vocab_size, trg_vocab_size), use_cuda) # TODO: add more arguments as necessary #nmt = init_model((src_vocab_size, trg_vocab_size), use_cuda) if use_cuda > 0: nmt.cuda() else: nmt.cpu() criterion = torch.nn.NLLLoss() optimizer = eval("torch.optim." + options.optimizer)(nmt.parameters(), options.learning_rate) # main training loop last_dev_avg_loss = float("inf") for epoch_i in range(options.epochs): logging.info("At {0}-th epoch.".format(epoch_i)) # srange generates a lazy sequence of shuffled range for i, batch_i in enumerate(rand.srange(len(batched_train_src))): #if random.random() > 0.5: if False: #i > 1500: # TODO REMOVE THIS !!!!!!!!!!!!!!!!!!!! #model.save("intermediate_ds_new", nmt); break if i % 200 == 0: model.save("intermediate_ds", nmt) pass train_src_batch = Variable(batched_train_src[batch_i] ) # of size (src_seq_len, batch_size) train_trg_batch = Variable(batched_train_trg[batch_i] ) # of size (src_seq_len, batch_size) train_src_mask = Variable(batched_train_src_mask[batch_i]) train_trg_mask = Variable(batched_train_trg_mask[batch_i]) if use_cuda: train_src_batch = train_src_batch.cuda() train_trg_batch = train_trg_batch.cuda() train_src_mask = train_src_mask.cuda() train_trg_mask = train_trg_mask.cuda() sys_out_batch, translated_sentence_wd_index = nmt( train_src_batch, train_trg_batch ) # (trg_seq_len, batch_size, trg_vocab_size) # TODO: add more arguments as necessary train_trg_mask = train_trg_mask.view(-1) train_trg_batch = train_trg_batch.view(-1) train_trg_batch = train_trg_batch.masked_select(train_trg_mask) train_trg_mask = train_trg_mask.unsqueeze(1).expand( len(train_trg_mask), trg_vocab_size) sys_out_batch = sys_out_batch.view(-1, trg_vocab_size) sys_out_batch = sys_out_batch.masked_select(train_trg_mask).view( -1, trg_vocab_size) loss = criterion(sys_out_batch, train_trg_batch) logging.debug("loss at batch {0}: {1}".format(i, loss.data[0])) optimizer.zero_grad() loss.backward() optimizer.step() # validation -- this is a crude esitmation because there might be some paddings at the end dev_loss = 0.0 for batch_i in range(len(batched_dev_src)): #if random.random() > 0.5: if False: #i > 1500: # TODO REMOVE THIS !!!!!!!!!!!!!!!!!!!! #model.save("intermediate_ds_new", nmt); break dev_src_batch = Variable(batched_dev_src[batch_i], volatile=True) dev_trg_batch = Variable(batched_dev_trg[batch_i], volatile=True) dev_src_mask = Variable(batched_dev_src_mask[batch_i], volatile=True) dev_trg_mask = Variable(batched_dev_trg_mask[batch_i], volatile=True) if use_cuda: dev_src_batch = dev_src_batch.cuda() dev_trg_batch = dev_trg_batch.cuda() dev_src_mask = dev_src_mask.cuda() dev_trg_mask = dev_trg_mask.cuda() sys_out_batch, translated_sentence_wd_index = nmt( dev_src_batch ) # (trg_seq_len, batch_size, trg_vocab_size) # TODO: add more arguments as necessary actual_trg_max_len = dev_trg_mask.data.shape[0] dev_trg_mask = dev_trg_mask.view(-1) dev_trg_batch = dev_trg_batch.view(-1) dev_trg_batch = dev_trg_batch.masked_select(dev_trg_mask) dev_trg_mask = dev_trg_mask.unsqueeze(1).expand( len(dev_trg_mask), trg_vocab_size) # TODO Remove this!!!!!! predicted_trg_max_len = sys_out_batch.data.shape[0] if actual_trg_max_len > predicted_trg_max_len: sys_out_batch = torch.cat( (sys_out_batch, torch.ones((actual_trg_max_len - predicted_trg_max_len, options.batch_size, trg_vocab_size)))) else: sys_out_batch = sys_out_batch[0:actual_trg_max_len] # TODO Remove this ^^^ !!!!!! sys_out_batch = sys_out_batch.view(-1, trg_vocab_size) sys_out_batch = sys_out_batch.masked_select(dev_trg_mask).view( -1, trg_vocab_size) loss = criterion(sys_out_batch, dev_trg_batch) logging.debug("dev loss at batch {0}: {1}".format( batch_i, loss.data[0])) dev_loss += loss #if True: break dev_avg_loss = dev_loss / len(batched_dev_src) logging.info( "Average loss value per instance is {0} at the end of epoch {1}". format(dev_avg_loss.data[0], epoch_i)) if (last_dev_avg_loss - dev_avg_loss).data[0] < options.estop: logging.info( "Early stopping triggered with threshold {0} (previous dev loss: {1}, current: {2})" .format(epoch_i, last_dev_avg_loss.data[0], dev_avg_loss.data[0])) break torch.save( nmt, open( options.model_file + ".nll_{0:.2f}.epoch_{1}".format(dev_avg_loss.data[0], epoch_i), 'wb'), pickle_module=dill) last_dev_avg_loss = dev_avg_loss import datetime current_dt_tm = datetime.datetime.now().strftime("%Y_%m_%d_%H_%M_%S") model.save("%s_%s" % (options.model_file, current_dt_tm), nmt)
def train(args: Dict): """ Train the NMT Model. @param args (Dict): args from cmd line """ do_bleu = '--ignore-test-bleu' not in args or not args['--ignore-test-bleu'] train_data_src = read_corpus(args['--train-src'], source='src', dev_mode=dev_mode) train_data_tgt = read_corpus(args['--train-tgt'], source='tgt', dev_mode=dev_mode) dev_data_src = read_corpus(args['--dev-src'], source='src', dev_mode=dev_mode) dev_data_tgt = read_corpus(args['--dev-tgt'], source='tgt', dev_mode=dev_mode) if do_bleu: test_data_src = read_corpus(args['--test-src'], source='src', dev_mode=dev_mode) test_data_tgt = read_corpus(args['--test-tgt'], source='tgt', dev_mode=dev_mode) train_data = list(zip(train_data_src, train_data_tgt)) dev_data = list(zip(dev_data_src, dev_data_tgt)) max_tokens_in_sentence = int(args['--max-decoding-time-step']) train_data = clean_data(train_data, max_tokens_in_sentence) dev_data = clean_data(dev_data, max_tokens_in_sentence) train_batch_size = int(args['--batch-size']) dev_batch_size = 128 clip_grad = float(args['--clip-grad']) valid_niter = int(args['--valid-niter']) bleu_niter = int(args['--bleu-niter']) log_every = int(args['--log-every']) model_save_path = args['--save-to'] vocab = Vocab.load(args['--vocab'], args['--word_freq']) model = NMT(embed_size=int(args['--embed-size']), hidden_size=int(args['--hidden-size']), dropout_rate=float(args['--dropout']), vocab=vocab) writer = SummaryWriter() # model = TransformerNMT(vocab, num_hidden_layers=3) model.train() uniform_init = float(args['--uniform-init']) if np.abs(uniform_init) > 0.: print('uniformly initialize parameters [-%f, +%f]' % (uniform_init, uniform_init), file=sys.stderr) for p in model.parameters(): if p.dim() > 1: torch.nn.init.xavier_uniform_(p) else: p.data.uniform_(-uniform_init, uniform_init) vocab_mask = torch.ones(len(vocab.tgt)) vocab_mask[vocab.tgt['<pad>']] = 0 device = torch.device("cuda:0" if args['--cuda'] else "cpu") print('use device: %s' % device, file=sys.stderr) model = model.to(device) optimizer = torch.optim.Adam(model.parameters(), lr=float(args['--lr'])) num_trial = 0 train_iter = patience = cum_loss = report_loss = cum_tgt_words = report_tgt_words = 0 cum_examples = report_examples = epoch = valid_num = 0 hist_valid_scores = [] train_time = begin_time = time.time() print("Sorting dataset based on difficulty...") dataset = (train_data, dev_data) ordered_dataset = load_order(args['--order-name'], dataset, vocab) # TODO: order = balance_order(order, dataset) (train_data, dev_data) = ordered_dataset visualize_scoring_examples = False if visualize_scoring_examples: visualize_scoring(ordered_dataset, vocab) n_iters = math.ceil(len(train_data) / train_batch_size) print("n_iters per epoch is {}: ({} / {})".format(n_iters, len(train_data), train_batch_size)) max_epoch = int(args['--max-epoch']) max_iters = max_epoch * n_iters print('begin Maximum Likelihood training') print('Using order function: {}'.format(args['--order-name'])) print('Using pacing function: {}'.format(args['--pacing-name'])) while True: epoch += 1 for _ in range(n_iters): # Get pacing data according to train_iter current_train_data, current_dev_data = pacing_data( train_data, dev_data, time=train_iter, warmup_iters=int(args["--warmup-iters"]), method=args['--pacing-name'], tb=writer) # Uniformly sample batches from the paced dataset src_sents, tgt_sents = get_pacing_batch( current_train_data, batch_size=train_batch_size, shuffle=True) train_iter += 1 # ERROR START optimizer.zero_grad() batch_size = len(src_sents) example_losses = -model(src_sents, tgt_sents) # (batch_size,) batch_loss = example_losses.sum() loss = batch_loss / batch_size loss.backward() # clip gradient grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), clip_grad) optimizer.step() batch_losses_val: int = batch_loss.item() report_loss += batch_losses_val cum_loss += batch_losses_val tgt_words_num_to_predict = sum( len(s[1:]) for s in tgt_sents) # omitting leading `<s>` report_tgt_words += tgt_words_num_to_predict cum_tgt_words += tgt_words_num_to_predict report_examples += batch_size cum_examples += batch_size if train_iter % log_every == 0: print( 'epoch %d, iter %d, avg. loss %.2f, avg. ppl %.2f ' 'cum. examples %d, speed %.2f words/sec, time elapsed %.2f sec' % (epoch, train_iter, report_loss / report_examples, math.exp(report_loss / report_tgt_words), cum_examples, report_tgt_words / (time.time() - train_time), time.time() - begin_time), file=sys.stderr) writer.add_scalar('Loss/train', report_loss / report_examples, train_iter) writer.add_scalar('ppl/train', math.exp(report_loss / report_tgt_words), train_iter) train_time = time.time() report_loss = report_tgt_words = report_examples = 0. # evaluate BLEU if train_iter % bleu_niter == 0 and do_bleu: bleu = decode_with_params( model, test_data_src, test_data_tgt, int(args['--beam-size']), int(args['--max-decoding-time-step'])) writer.add_scalar('bleu/test', bleu, train_iter) # perform validation if train_iter % valid_niter == 0: print( 'epoch %d, iter %d, cum. loss %.2f, cum. ppl %.2f cum. examples %d' % (epoch, train_iter, cum_loss / cum_examples, np.exp(cum_loss / cum_tgt_words), cum_examples), file=sys.stderr) cum_loss = cum_examples = cum_tgt_words = 0. valid_num += 1 print('begin validation ...', file=sys.stderr) # compute dev. ppl and bleu # dev batch size can be a bit larger dev_ppl = evaluate_ppl(model, current_dev_data, batch_size=dev_batch_size) valid_metric = -dev_ppl writer.add_scalar('ppl/valid', dev_ppl, train_iter) cum_loss = cum_examples = cum_tgt_words = 0. valid_num += 1 print('validation: iter %d, dev. ppl %f' % (train_iter, dev_ppl), file=sys.stderr) is_better = len(hist_valid_scores ) == 0 or valid_metric > max(hist_valid_scores) hist_valid_scores.append(valid_metric) if is_better: patience = 0 print('save currently the best model to [%s]' % model_save_path, file=sys.stderr) model.save(model_save_path) # also save the optimizers' state torch.save(optimizer.state_dict(), model_save_path + '.optim') elif patience < int(args['--patience']): patience += 1 print('hit patience %d' % patience, file=sys.stderr) if patience == int(args['--patience']): num_trial += 1 print('hit #%d trial' % num_trial, file=sys.stderr) if num_trial == int(args['--max-num-trial']): print('early stop!', file=sys.stderr) exit(0) # decay lr, and restore from previously best checkpoint lr = optimizer.param_groups[0]['lr'] * \ float(args['--lr-decay']) print( 'load previously best model and decay learning rate to %f' % lr, file=sys.stderr) # load model params = torch.load( model_save_path, map_location=lambda storage, loc: storage) model.load_state_dict(params['state_dict']) model = model.to(device) print('restore parameters of the optimizers', file=sys.stderr) optimizer.load_state_dict( torch.load(model_save_path + '.optim')) # set new lr for param_group in optimizer.param_groups: param_group['lr'] = lr # reset patience patience = 0 if epoch >= int(args['--max-epoch']): print('reached maximum number of epochs!', file=sys.stderr) exit(0)
def main(options): src_train, src_dev, src_test, src_vocab = torch.load( open(options.data_file + "." + options.src_lang, 'rb')) trg_train, trg_dev, trg_test, trg_vocab = torch.load( open(options.data_file + "." + options.trg_lang, 'rb')) batched_train_src, batched_train_src_mask, sort_index = utils.tensor.advanced_batchize( src_train, options.batch_size, src_vocab.stoi["<blank>"]) batched_train_trg, batched_train_trg_mask = utils.tensor.advanced_batchize_no_sort( trg_train, options.batch_size, trg_vocab.stoi["<blank>"], sort_index) batched_dev_src, batched_dev_src_mask, sort_index = utils.tensor.advanced_batchize( src_dev, options.batch_size, src_vocab.stoi["<blank>"]) batched_dev_trg, batched_dev_trg_mask = utils.tensor.advanced_batchize_no_sort( trg_dev, options.batch_size, trg_vocab.stoi["<blank>"], sort_index) src_vocab_size = len(src_vocab) trg_vocab_size = len(trg_vocab) nmt = NMT(src_vocab_size, trg_vocab_size) # TODO: add more arguments as necessary if torch.cuda.is_available(): nmt.cuda() else: nmt.cpu() criterion = torch.nn.NLLLoss() optimizer = eval("torch.optim." + options.optimizer)(nmt.parameters(), options.learning_rate) # main training loop last_dev_avg_loss = float("inf") for epoch_i in range(options.epochs): logging.info("At {0}-th epoch.".format(epoch_i)) # srange generates a lazy sequence of shuffled range for i, batch_i in enumerate(utils.rand.srange(len(batched_train_src))): train_src_batch = to_var(batched_train_src[batch_i] ) # of size (src_seq_len, batch_size) train_trg_batch = to_var(batched_train_trg[batch_i] ) # of size (src_seq_len, batch_size) train_src_mask = to_var(batched_train_src_mask[batch_i]) train_trg_mask = to_var(batched_train_trg_mask[batch_i]) sys_out_batch = nmt( train_src_batch, train_trg_batch, training=True ) # (trg_seq_len, batch_size, trg_vocab_size) # TODO: add more arguments as necessary train_trg_mask = train_trg_mask.view(-1) train_trg_batch = train_trg_batch.view(-1) train_trg_batch = train_trg_batch.masked_select(train_trg_mask) train_trg_mask = train_trg_mask.unsqueeze(1).expand( len(train_trg_mask), trg_vocab_size) sys_out_batch = sys_out_batch.view(-1, trg_vocab_size) sys_out_batch = sys_out_batch.masked_select(train_trg_mask).view( -1, trg_vocab_size) loss = criterion(sys_out_batch, train_trg_batch) # logging.debug("loss at batch {0}: {1}".format(i, loss.data[0])) optimizer.zero_grad() loss.backward() optimizer.step() # validation -- this is a crude esitmation because there might be some paddings at the end dev_loss = 0.0 for batch_i in range(len(batched_dev_src)): dev_src_batch = to_var(batched_dev_src[batch_i], volatile=True) dev_trg_batch = to_var(batched_dev_trg[batch_i], volatile=True) dev_src_mask = to_var(batched_dev_src_mask[batch_i], volatile=True) dev_trg_mask = to_var(batched_dev_trg_mask[batch_i], volatile=True) sys_out_batch = nmt( dev_src_batch, dev_trg_batch ) # (trg_seq_len, batch_size, trg_vocab_size) # TODO: add more arguments as necessary dev_trg_mask = dev_trg_mask.view(-1) dev_trg_batch = dev_trg_batch.view(-1) dev_trg_batch = dev_trg_batch.masked_select(dev_trg_mask) dev_trg_mask = dev_trg_mask.unsqueeze(1).expand( len(dev_trg_mask), trg_vocab_size) sys_out_batch = sys_out_batch.view(-1, trg_vocab_size) sys_out_batch = sys_out_batch.masked_select(dev_trg_mask).view( -1, trg_vocab_size) #sys_out_batch = sys_out_batch.masked_select(dev_trg_mask).view(-1, trg_vocab_size) loss = criterion(sys_out_batch, dev_trg_batch) # logging.debug("dev loss at batch {0}: {1}".format(batch_i, loss.data[0])) dev_loss += loss dev_avg_loss = dev_loss / len(batched_dev_src) logging.info( "Average loss value per instance is {0} at the end of epoch {1}". format(dev_avg_loss.data[0], epoch_i)) # if (last_dev_avg_loss - dev_avg_loss).data[0] < options.estop: # logging.info("Early stopping triggered with threshold {0} (previous dev loss: {1}, current: {2})".format(epoch_i, last_dev_avg_loss.data[0], dev_avg_loss.data[0])) # break torch.save( nmt, open( options.model_file + ".nll_{0:.2f}.epoch_{1}".format(dev_avg_loss.data[0], epoch_i), 'wb'), pickle_module=dill) last_dev_avg_loss = dev_avg_loss