PG-net模型训练部分

注意:在这里最后几个epoch才优化coverage_loss

import os
import args
import torch
import random
from tqdm import tqdm
import evaluate
from optimizer import BertAdam
from dataset.dataloader import Dureader
from dataset.file_utils import PYTORCH_PRETRAINED_BERT_CACHE
from model_dir.modeling import BertForQuestionAnswering,BertConfig

random.seed(args.seed)
torch.manual_seed(args.seed)
device = args.device
#print(device)
device_ids = [0]

if len(device_ids) > 0:
    torch.cuda.manual_seed_all(args.seed)
model = BertForQuestionAnswering.from_pretrained('./chinese_roberta_wwm_ext')

if len(device_ids) > 1:
    model = torch.nn.DataParallel(model,device_ids = device_ids)
    model = model.cuda(device = device_ids[0])
elif len(device_ids) == 1:
    model.to(device)
    
param_optimizer = list(model.named_parameters())
param_optimizer = [n for n in param_optimizer if 'pooler' not in n[0]]
no_decay = ['bias','LayerNorm.bias','LayerNorm.weight']
optimizer_grouped_parameters = [
    {'params':[p for n,p in param_optimizer if not any(nd in n for nd in no_decay)],'weight_decay':0.01},
    {'params':[p for n,p in param_optimizer if any(nd in n for nd in no_decay)]}
]
optimizer = BertAdam(optimizer_grouped_parameters,lr = args.learning_rate,warmup = 0.1,t_total = args.num_train_optimization_steps)


data = Dureader()
train_dataloader,dev_dataloader = data.train_iter,data.dev_iter
best_loss = 100000.0
model.train()
for epoch in range(args.num_train_epochs):
    main_losses,ide_losses = 0,0
    train_loss,train_loss_total = 0.0,0.0
    n_words,n_words_total = 0,0
    n_sent,n_sents_total = 0,0
    
    for step,batch in enumerate(tqdm(train_dataloader,desc = 'Epoch')):
            input_ids, input_mask, input_ids_q, input_mask_q, answer_ids, answer_mask, \
            segment_ids, can_answer = \
                batch.input_ids, batch.input_mask, batch.input_ids_q, batch.input_mask_q,\
                batch.answer_ids, batch.answer_mask, batch.segment_ids,batch.can_answer
            answer_inputs = answer_ids[:,:-1]    # 去掉EOS
            answer_targets = answer_ids[:,1:]    # 去掉BOS
            answer_len = answer_mask.sum(1) - 1
            if len(device_ids) > 1:
                input_ids, input_mask, input_ids_q, input_mask_q, answer_inputs, answer_len, answer_targets,\
                segment_ids, can_answer = \
                    input_ids.cuda(device=device_ids[0]), input_mask.cuda(device=device_ids[0]), \
                    input_ids_q.cuda(device=device_ids[0]), input_mask_q.cuda(device=device_ids[0]), \
                    answer_inputs.cuda(device=device_ids[0]), answer_len.cuda(device=device_ids[0]), \
                    answer_targets.cuda(device=device_ids[0]), \
                    segment_ids.cuda(device=device_ids[0]), can_answer.cuda(device=device_ids[0])
            elif len(device_ids) == 1:
                input_ids, input_mask, input_ids_q, input_mask_q, answer_inputs, answer_len, answer_targets, \
                segment_ids, can_answer = \
                    input_ids.to(device), input_mask.to(device), input_ids_q.to(device), input_mask_q.to(device), \
                    answer_inputs.to(device), answer_len.to(device), answer_targets.to(device), \
                    segment_ids.to(device), can_answer.to(device)
            
            loss,coverage_loss = model(input_ids,input_ids_q,token_type_ids = segment_ids,attention_mask = input_mask,
                                       attention_mask_q = input_mask_q,dec_inputs = answer_inputs,dec_inputs_len = answer_len,
                                       dec_targets = answer_targets,can_answer = can_answer)
            if epoch >= args.num_train_epochs - 2:
                loss = loss + coverage_loss
            train_loss_total += float(loss.item())
            n_words_total += torch.sum(answer_len)
            n_sents_total += answer_len.size(0)
            if step % args.display_freq == 0 and step:
                loss_int = (train_loss_total - train_loss)
                n_words_int = (n_words_total - n_words)
                loss_per_words = loss_int / n_words_int.item()
                avg_loss = loss_per_words
                print('Epoch {0:<3}'.format(epoch),'Step {0:<10}'.format(step),'Avg_loss {0:<10.2f}'.format(avg_loss))
                train_loss,n_words,n_sents = (train_loss_total,n_words_total.item(),n_sents_total)
            elif step == 0:
                loss_int = (train_loss_total - train_loss)
                n_words_int = (n_words_total - n_words)
                loss_per_words = loss_int / n_words_int.item()
                avg_loss = loss_per_words

                print('Epoch {0:<3}'.format(epoch),
                      'Step {0:<10}'.format(step),
                      'Avg_loss {0:<10.2f}'.format(avg_loss))
                train_loss, n_words, n_sents = (train_loss_total, n_words_total.item(), n_sents_total)
            
            if len(device_ids) > 1:
                loss = loss.mean()
            if args.gradient_accumulation_steps > 1:
                loss = loss / args.gradient_accumulation_steps
            
            loss.backward()
            
            
            # 更新梯度
            if (step+1) % args.gradient_accumulation_steps == 0:
                optimizer.step()
                optimizer.zero_grad()
            
            
            # 验证
            if step % args.log_step == 4:
                eval_loss = evaluate.evaluate(model,dev_dataloader,device_ids)
                if eval_loss < best_loss:
                    best_loss = eval_loss
                    if len(device_ids) > 1:
                        torch.save(model.module.state_dict(),'./model_dir/debug')
                    if len(device_ids) == 1:
                        torch.save(model.state_dict(), './model_dir/' + "debug")
                model.train()