def main():
if csv:
uttrs, resps = get_uttr_resp_csv(files)
else:
if data_type == "single":
uttrs, resps = get_uttr_resps_oneline_single(files)
else:
uttrs, resps = get_uttr_resp_oneline_pair(files)
logger.info("处理uttrance")
for i, uttr in enumerate(uttrs):
result = ""
for ch in uttr:
if is_chinese_punc(ch) or is_chinese_char(ch):
result += f" {ch} "
else:
result += ch
result = result.strip().replace(" ", " ")
uttrs[i] = result
dialogues = []
for i, (uttr, resp) in tqdm(enumerate(zip(uttrs, resps))):
d = copy.deepcopy(template)
d["id"] = i
d["utterance"] = uttr
d["responses"][0]["reply"] = resp
dialogues.append(d)
with open(Path(basename) / to_file, 'w', encoding="utf8") as f:
json.dump(dialogues, f, ensure_ascii=False, indent=4)
def evaluate(self, dataloader: DataLoader):
result = dict()
y_pred = self.predict(dataloader)
y_true = dataloader.label
if isinstance(self._task, tasks.Classification):
y_pred_label = np.argmax(y_pred[:10], axis=-1)
# 记录前10个真实标签和预测标签
logger.info(
f"The former 10 true label is {y_true[:10]} | pred label is {y_pred_label}"
)
elif isinstance(self._task, tasks.Ranking):
y_true = y_true.reshape(len(y_true), 1)
loss = sum([
c(torch.tensor(y_pred), torch.tensor(y_true))
for c in self._criterions
])
self._writer.add_scalar("Loss/eval", loss.item(), self._iteration)
try:
id_left = dataloader.id_left
except:
pass
if isinstance(self._task, tasks.Ranking):
for metric in self._task.metrics:
result[metric] = self._eval_metric_on_data_frame(
metric, id_left, y_true, y_pred.squeeze(axis=-1))
else:
for metric in self._task.metrics:
result[metric] = metric(y_true, y_pred)
return result, loss
def save(self):
"""
Save the trainer.
`Trainer` parameters like epoch, best_so_far, model, optimizer
and early_stopping will be savad to specific file path.
:param path: Path to save trainer.
"""
checkpoint = self._save_dir.joinpath('trainer.pt')
logger.info(f" ** save trainer model to {checkpoint} ** ")
if self._data_parallel:
model = self._model.module.state_dict()
else:
model = self._model.state_dict()
state = {
'epoch': self._epoch,
'model': model,
'optimizer': self._optimizer.state_dict(),
'early_stopping': self._early_stopping.state_dict(),
}
if self._epoch_scheduler:
state['epoch_scheduler'] = self._epoch_scheduler.state_dict()
if self._step_scheduler:
state['step_scheduler'] = self._step_scheduler.state_dict()
torch.save(state, checkpoint)
def save_json(basename, json_files, tokenizer, model, device):
for file in json_files:
logger.info(f"正在处理文件 {file} ...")
with open(Path(basename) / file, 'r', encoding="utf8") as f:
records = json.load(f)
for i, record in tqdm(enumerate(records)):
uttr = record['utterance']
uttr_list = re.split(r"\[U=\d+\]", uttr)
uttr = uttr_list[-1]
uttr = replace_kv(uttr).replace(" ", "")
responses = record['responses']
for j, reply in enumerate(responses):
ans = reply['reply']
ans = replace_kv(ans)
pred = albert_predict(text_a=[uttr],
text_b=[ans],
tokenizer=tokenizer,
model=model,
device=device)
lan_style = labels[pred[0]]
## 保存语言风格
reply['language_style'] = lan_style
responses[j] = reply
records[i]['responses'] = responses
save_file = to_file_prefix + "_" + file
logger.info(f"处理完成并保存在 {save_file}")
with open(save_file, 'w', encoding='utf8') as to_f:
json.dump(records, to_f, ensure_ascii=False, indent=4)
def _load_model(self,
model: BaseModel,
device: typing.Union[torch.device, int, list,
None] = None):
if not isinstance(model, BaseModel):
raise ValueError(f"model should be a `BaseModel` instance. "
f"But got {type(model)}")
logger.info(" ** load model and device ** ")
self._task = model.params['task']
self._data_parallel = False
self._model = model
# 如果指定了多个GPU，则默认是数据并行
if isinstance(device, list) and len(device):
logger.info(" ** data parallel ** ")
self._data_parallel = True
self._model = torch.nn.DataParallel(self._model, device_ids=device)
self._device = device[0]
else:
if not (isinstance(device, torch.device)
or isinstance(device, int)):
device = torch.device(
"cuda" if torch.cuda.is_available() else "cpu")
self._device = device
self._model.to(self._device)
def get_dataloader(train_data, valid_data):
train_dataset = PairDataset(train_data, num_neg=0)
valid_dataset = PairDataset(valid_data, num_neg=0)
padding = MultiQAPadding(fixed_length_uttr=fixed_length_uttr,
fixed_length_resp=fixed_length_resp,
fixed_length_turn=fixed_length_turn,
data_type=data_type)
train_dataloader = DictDataLoader(train_dataset,
batch_size=batch_size,
turns=fixed_length_turn,
stage='train',
shuffle=True,
sort=False,
callback=padding)
valid_dataloader = DictDataLoader(valid_dataset,
batch_size=batch_size,
turns=fixed_length_turn,
stage='dev',
shuffle=False,
sort=False,
callback=padding)
for i, (x, y) in enumerate(train_dataloader):
# 打印Utterance的形状
logger.info(f"The shape of utternace is {x[constants.UTTRS].shape}")
if i == 0:
break
return train_dataloader, valid_dataloader
def save_model(self):
"""Save the model."""
checkpoint = self._save_dir.joinpath('model.pt')
logger.info(f" ** save raw model to {checkpoint} ** ")
if self._data_parallel:
torch.save(self._model.module.state_dict(), checkpoint)
else:
torch.save(self._model.state_dict(), checkpoint)
def get_uttr_resp_oneline_pair(files):
uttrs, resps = [], []
for file in files:
logger.info(f"正在处理文件 {file}")
with open(Path(basename) / file, 'r', encoding="utf8") as f:
for line in tqdm(f):
line = line.strip()
if line:
line_list = line.split("\t")
uttrs.append(line_list[0].strip().replace(" ", ""))
resps.append(line_list[-1].strip().replace(" ", ""))
return uttrs, resps
def restore_model(self, checkpoint: typing.Union[str, Path]):
"""
Restore model.
:param checkpoint: A checkpoint from which to continue training.
"""
logger.info(" ** restore raw model ** ")
state = torch.load(checkpoint, map_location=self._device)
if self._data_parallel:
self._model.module.load_state_dict(state)
else:
self._model.load_state_dict(state)
def _load_path(self, checkpoint: typing.Union[str, Path],
save_dir: typing.Union[str, Path]):
logger.info(" ** restore exist model checkpoint ** ")
if not save_dir:
save_dir = Path('.').joinpath('save')
if not Path(save_dir).exists():
Path(save_dir).mkdir(parents=True)
self._save_dir = Path(save_dir)
if checkpoint:
if self._save_all:
self.restore(checkpoint)
else:
self.restore_model(checkpoint)
def __init__(self,
model: BaseModel,
optimizer: Optimizer,
trainloader: DataLoader,
validloader: DataLoader,
device: typing.Union[torch.device, int, list, None] = None,
writer: SummaryWriter = None,
start_epoch: int = 1,
epochs: int = 10,
validate_interval: typing.Optional[int] = None,
epoch_scheduler: typing.Any = None,
step_scheduler: typing.Any = None,
clip_norm: typing.Union[float, int] = None,
l1_reg: float = 0.0,
l2_reg: float = 0.0,
patience: typing.Optional[int] = None,
key: typing.Any = None,
checkpoint: typing.Union[str, Path] = None,
save_dir: typing.Union[str, Path] = None,
save_all: bool = False,
verbose: int = 1,
**kwargs):
log_dir = Path(save_dir).joinpath("runs")
self._writer = writer or SummaryWriter(log_dir=log_dir)
logger.info(f" ** save tensorboard to {self._writer.get_logdir()} ** ")
self._load_model(model, device)
self._load_dataloader(trainloader, validloader, validate_interval)
self._optimizer = optimizer
self._epoch_scheduler = epoch_scheduler # 由于这里的scheduler是在每个epoch之后调用一次，所以在定义的时候注意设置更新的step数和epoch一致
self._step_scheduler = step_scheduler # 这个是针对step的scheduler
self._clip_norm = clip_norm # 梯度裁剪
# 正则化系数
self._l1_reg = l1_reg
self._l2_reg = l2_reg
self._criterions = self._task.losses
if not key:
key = self._task.metrics[0]
self._early_stopping = EarlyStopping(patience=patience, key=key)
self._start_epoch = start_epoch
self._epochs = epochs
self._iteration = 0
self._verbose = verbose
self._save_all = save_all
self._load_path(checkpoint, save_dir)
def restore(self, checkpoint: typing.Union[str, Path] = None):
"""
Restore trainer.
:param checkpoint: A checkpoint from which to continue training.
"""
logger.info(" ** restore trainer model ** ")
state = torch.load(checkpoint, map_location=self._device)
if self._data_parallel:
self._model.module.load_state_dict(state['model'])
else:
self._model.load_state_dict(state['model'])
self._optimizer.load_state_dict(state['optimizer'])
self._start_epoch = state['epoch'] + 1
self._early_stopping.load_state_dict(state['early_stopping'])
if self._epoch_scheduler:
self._epoch_scheduler.load_state_dict(state['epoch_scheduler'])
if self._step_scheduler:
self._step_scheduler.load_state_dict(state['step_scheduler'])
def get_uttr_resps_oneline_single(files):
uttrs = []
resps = []
flag = 0
for file in files:
logger.info(f"正在处理文件 {file}")
with open(Path(basename) / file, 'r', encoding="utf8") as f:
for i, line in tqdm(enumerate(f)):
line = line.strip()
j = i
if not line or flag:
if not line:
flag += 1
j = j - flag
if line:
if j & 1 == 0:
uttrs.append(line)
else:
resps.append(line)
return uttrs, resps
def preprocess_train_and_val(train_file, valid_file):
logger.info("读取数据")
train_data = pd.read_csv(train_file)
valid_data = pd.read_csv(valid_file)
train_data.dropna(axis=0, subset=[constants.UTTRS, constants.RESP], inplace=True)
valid_data.dropna(axis=0, subset=[constants.UTTRS, constants.RESP], inplace=True)
preprocessor = CNAlbertPreprocessorForMultiQA(Path(albert_path) / vocab_file,
uttr_len=fixed_length_uttr,
resp_len=fixed_length_resp,
add_special_tokens=add_specical_tokens)
logger.info(f"训练集数据量为：{train_data.shape[0]} | 验证集数据量为：{valid_data.shape[0]}")
logger.info("使用Preprocessor处理数据")
train_data = preprocessor.transform(train_data,
uttr_col=constants.UTTRS,
resp_col=constants.RESP)
valid_data = preprocessor.transform(valid_data,
uttr_col=constants.UTTRS,
resp_col=constants.RESP)
use_cols = ['D_num', 'turns', 'utterances', 'response', 'utterances_len', 'response_len', 'label']
train_data = train_data[use_cols]
valid_data = valid_data[use_cols]
label_type = int
train_data[constants.LABEL] = train_data[constants.LABEL].astype(label_type)
valid_data[constants.LABEL] = valid_data[constants.LABEL].astype(label_type)
return train_data, valid_data
def load_and_cache_examples(args, task, tokenizer, data_type='train'):
if args.local_rank not in [-1, 0] and not evaluate:
torch.distributed.barrier(
) # Make sure only the first process in distributed training process the dataset, and the others will use the cache
processor = processors[task]()
# Load data features from cache or dataset file
cached_features_file = os.path.join(
args.data_dir, 'cached_crf-{}_{}_{}_{}'.format(
data_type,
list(filter(None, args.model_name_or_path.split('/'))).pop(),
str(args.train_max_seq_length if data_type ==
'train' else args.eval_max_seq_length), str(task)))
if os.path.exists(cached_features_file) and not args.overwrite_cache:
logger.info("Loading features from cached file %s",
cached_features_file)
features = torch.load(cached_features_file)
else:
logger.info("Creating features from dataset file at %s", args.data_dir)
label_list = processor.get_labels()
if data_type == 'train':
examples = processor.get_train_examples(args.data_dir)
elif data_type == 'dev':
examples = processor.get_dev_examples(args.data_dir)
else:
examples = processor.get_test_examples(args.data_dir)
features = convert_examples_to_features(examples=examples,
tokenizer=tokenizer,
label_list=label_list,
max_seq_length=args.train_max_seq_length if data_type=='train' \
else args.eval_max_seq_length,
cls_token_at_end=bool(args.model_type in ["xlnet"]),
pad_on_left=bool(args.model_type in ['xlnet']),
cls_token = tokenizer.cls_token,
cls_token_segment_id=2 if args.model_type in ["xlnet"] else 0,
sep_token=tokenizer.sep_token,
# pad on the left for xlnet
pad_token=tokenizer.convert_tokens_to_ids([tokenizer.pad_token])[0],
pad_token_segment_id=4 if args.model_type in ['xlnet'] else 0,
)
if args.local_rank in [-1, 0]:
logger.info("Saving features into cached file %s",
cached_features_file)
torch.save(features, cached_features_file)
if args.local_rank == 0 and not evaluate:
torch.distributed.barrier(
) # Make sure only the first process in distributed training process the dataset, and the others will use the cache
# Convert to Tensors and build dataset
all_input_ids = torch.tensor([f.input_ids for f in features],
dtype=torch.long)
all_input_mask = torch.tensor([f.input_mask for f in features],
dtype=torch.long)
all_segment_ids = torch.tensor([f.segment_ids for f in features],
dtype=torch.long)
all_label_ids = torch.tensor([f.label_ids for f in features],
dtype=torch.long)
all_lens = torch.tensor([f.input_len for f in features], dtype=torch.long)
dataset = TensorDataset(all_input_ids, all_input_mask, all_segment_ids,
all_lens, all_label_ids)
return dataset
def predict(args, model, tokenizer, prefix=""):
pred_output_dir = args.output_dir
if not os.path.exists(pred_output_dir) and args.local_rank in [-1, 0]:
os.makedirs(pred_output_dir)
test_dataset = load_and_cache_examples(args,
args.task_name,
tokenizer,
data_type='test')
print(len(test_dataset))
# Note that DistributedSampler samples randomly
test_sampler = SequentialSampler(
test_dataset) if args.local_rank == -1 else DistributedSampler(
test_dataset)
test_dataloader = DataLoader(test_dataset,
sampler=test_sampler,
batch_size=1,
collate_fn=collate_fn)
# Eval!
logger.info("***** Running prediction %s *****", prefix)
logger.info(" Num examples = %d", len(test_dataset))
logger.info(" Batch size = %d", 1)
results = []
output_submit_file = os.path.join(pred_output_dir, prefix,
"test_prediction.json")
pbar = ProgressBar(n_total=len(test_dataloader), desc="Predicting")
for step, batch in enumerate(test_dataloader):
model.eval()
batch = tuple(t.to(args.device) for t in batch)
with torch.no_grad():
inputs = {
"input_ids": batch[0],
"attention_mask": batch[1],
"start_positions": None,
"end_positions": None
}
if args.model_type != "distilbert":
# XLM and RoBERTa don"t use segment_ids
inputs["token_type_ids"] = (batch[2] if args.model_type
in ["bert", "xlnet"] else None)
outputs = model(**inputs)
start_logits, end_logits = outputs[:2]
R = bert_extract_item(start_logits, end_logits)
if R:
label_entities = [[args.id2label[x[0]], x[1], x[2]] for x in R]
else:
label_entities = []
json_d = {}
json_d['id'] = step
json_d['entities'] = label_entities
results.append(json_d)
pbar(step)
print(" ")
with open(output_submit_file, "w") as writer:
for record in results:
writer.write(json.dumps(record) + '\n')
def predict(args, model, tokenizer, prefix=""):
pred_output_dir = args.output_dir
if not os.path.exists(pred_output_dir) and args.local_rank in [-1, 0]:
os.makedirs(pred_output_dir)
test_dataset = load_and_cache_examples(args,
args.task_name,
tokenizer,
data_type='test')
# Note that DistributedSampler samples randomly
test_sampler = SequentialSampler(
test_dataset) if args.local_rank == -1 else DistributedSampler(
test_dataset)
test_dataloader = DataLoader(test_dataset,
sampler=test_sampler,
batch_size=1,
collate_fn=collate_fn)
# Eval!
logger.info("***** Running prediction %s *****", prefix)
logger.info(" Num examples = %d", len(test_dataset))
logger.info(" Batch size = %d", 1)
results = []
output_submit_file = os.path.join(pred_output_dir, prefix,
"test_prediction.json")
pbar = ProgressBar(n_total=len(test_dataloader), desc="Predicting")
if isinstance(model, nn.DataParallel):
model = model.module
for step, batch in enumerate(test_dataloader):
model.eval()
batch = tuple(t.to(args.device) for t in batch)
with torch.no_grad():
inputs = {
"input_ids": batch[0],
"attention_mask": batch[1],
"labels": None,
'input_lens': batch[4]
}
if args.model_type != "distilbert":
# XLM and RoBERTa don"t use segment_ids
inputs["token_type_ids"] = (batch[2] if args.model_type
in ["bert", "xlnet"] else None)
outputs = model(**inputs)
logits = outputs[0]
preds, _ = model.crf._obtain_labels(logits, args.id2label,
inputs['input_lens'])
preds = preds[0][1:-1] # [CLS]XXXX[SEP]
label_entities = get_entities(preds, args.id2label, args.markup)
json_d = {}
json_d['id'] = step
json_d['tag_seq'] = " ".join(preds)
json_d['entities'] = label_entities
results.append(json_d)
pbar(step)
print(" ")
with open(output_submit_file, "w") as writer:
for record in results:
writer.write(json.dumps(record) + '\n')
def process_and_save_txt(model, tokenizer, device):
logger.info("准备数据")
if files[0].endswith("json"):
text_lefts, text_rights = get_data_json(basename, files)
else:
text_lefts, text_rights = get_data_txt(basename, files)
logger.info("数据预处理完成")
logger.info("开始预测...")
preds = albert_predict(text_lefts,
text_b=text_rights,
tokenizer=tokenizer,
model=model,
device=device)
#print(Counter(preds))
logger.info("保存文件")
with open(f"{to_file_prefix}_result.txt", 'w', encoding="utf8") as f:
for q, a, l in zip(*(text_lefts, text_rights, preds)):
f.write(labels[l] + "\t")
f.write(q + "\t" + a)
f.write("\n")
def load_and_cache_examples(args, task, tokenizer, data_type='train'):
if args.local_rank not in [-1, 0] and not evaluate:
torch.distributed.barrier(
) # Make sure only the first process in distributed training process the dataset, and the others will use the cache
processor = processors[task]()
output_mode = "ranking"
# Load data features from cache or dataset file
cached_features_file = os.path.join(
args.data_dir, 'cached_{}_{}_{}_{}'.format(
data_type,
list(filter(None, args.model_name_or_path.split('/'))).pop(),
str(args.max_seq_length), str(task)))
if os.path.exists(cached_features_file):
logger.info("Loading features from cached file %s",
cached_features_file)
features = torch.load(cached_features_file)
else:
logger.info("Creating features from dataset file at %s", args.data_dir)
label_list = processor.get_labels()
if task in ['mnli', 'mnli-mm'] and 'roberta' in args.model_type:
# HACK(label indices are swapped in RoBERTa pretrained model)
label_list[1], label_list[2] = label_list[2], label_list[1]
if data_type == 'train':
logger.info("Get train examples ...")
examples = processor.get_train_examples(args.data_dir)
elif data_type == 'dev':
logger.info("Get dev examples ...")
examples = processor.get_dev_examples(args.data_dir)
else:
examples = processor.get_test_examples(args.data_dir)
features = convert_examples_to_features(
examples,
tokenizer,
label_list=label_list,
max_seq_length=args.max_seq_length,
output_mode=output_mode)
if args.local_rank in [-1, 0]:
logger.info("Saving features into cached file %s",
cached_features_file)
torch.save(features, cached_features_file)
if args.local_rank == 0 and not evaluate:
torch.distributed.barrier(
) # Make sure only the first process in distributed training process the dataset, and the others will use the cache
# Convert to Tensors and build dataset
all_input_ids = torch.tensor([f.input_ids for f in features],
dtype=torch.long)
all_attention_mask = torch.tensor([f.attention_mask for f in features],
dtype=torch.long)
all_token_type_ids = torch.tensor([f.token_type_ids for f in features],
dtype=torch.long)
all_lens = torch.tensor([f.input_len for f in features], dtype=torch.long)
if output_mode == "classification":
all_labels = torch.tensor([f.label for f in features],
dtype=torch.long)
else:
all_labels = torch.tensor([f.label for f in features],
dtype=torch.float)
dataset = TensorDataset(all_input_ids, all_attention_mask,
all_token_type_ids, all_lens, all_labels)
return dataset
def evaluate(args, model, tokenizer, prefix=""):
# Loop to handle MNLI double evaluation (matched, mis-matched)
eval_task_names = ("mnli", "mnli-mm") if args.task_name == "mnli" else (
args.task_name, )
eval_outputs_dirs = (args.output_dir, args.output_dir +
'-MM') if args.task_name == "mnli" else (
args.output_dir, )
results = {}
for eval_task, eval_output_dir in zip(eval_task_names, eval_outputs_dirs):
eval_dataset = load_and_cache_examples(args,
eval_task,
tokenizer,
data_type='dev')
if not os.path.exists(eval_output_dir) and args.local_rank in [-1, 0]:
os.makedirs(eval_output_dir)
args.eval_batch_size = args.per_gpu_eval_batch_size * max(
1, args.n_gpu)
# Note that DistributedSampler samples randomly
eval_sampler = SequentialSampler(
eval_dataset) if args.local_rank == -1 else DistributedSampler(
eval_dataset)
eval_dataloader = DataLoader(eval_dataset,
sampler=eval_sampler,
batch_size=args.eval_batch_size,
collate_fn=collate_fn)
# Eval!
logger.info("***** Running evaluation {} *****".format(prefix))
logger.info(" Num examples = %d", len(eval_dataset))
logger.info(" Batch size = %d", args.eval_batch_size)
eval_loss = 0.0
nb_eval_steps = 0
preds = None
out_label_ids = None
pbar = ProgressBar(n_total=len(eval_dataloader), desc="Evaluating")
for step, batch in enumerate(eval_dataloader):
model.eval()
batch = tuple(t.to(args.device) for t in batch)
with torch.no_grad():
inputs = {
'input_ids': batch[0],
'attention_mask': batch[1],
'labels': batch[3]
}
inputs['token_type_ids'] = batch[2]
outputs = model(**inputs)
tmp_eval_loss, logits = outputs[:2]
eval_loss += tmp_eval_loss.mean().item()
nb_eval_steps += 1
if preds is None:
preds = logits.detach().cpu().numpy()
out_label_ids = inputs['labels'].detach().cpu().numpy()
else:
preds = np.append(preds, logits.detach().cpu().numpy(), axis=0)
out_label_ids = np.append(
out_label_ids,
inputs['labels'].detach().cpu().numpy(),
axis=0)
pbar(step)
print(' ')
if 'cuda' in str(args.device):
torch.cuda.empty_cache()
eval_loss = eval_loss / nb_eval_steps
if args.output_mode == "classification":
preds = np.argmax(preds, axis=1)
elif args.output_mode == "regression":
preds = np.squeeze(preds)
elif args.output_mode == "ranking":
preds = np.greater(preds, 0.5).astype(np.float)
result = compute_metrics(eval_task, preds, out_label_ids)
results.update(result)
logger.info("***** Eval results {} *****".format(prefix))
for key in sorted(result.keys()):
logger.info(" %s = %s", key, str(result[key]))
return results
def train(args, train_dataset, model, tokenizer):
""" Train the model """
args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)
train_sampler = RandomSampler(
train_dataset) if args.local_rank == -1 else DistributedSampler(
train_dataset)
train_dataloader = DataLoader(train_dataset,
sampler=train_sampler,
batch_size=args.train_batch_size,
collate_fn=collate_fn)
if args.max_steps > 0:
t_total = args.max_steps
args.num_train_epochs = args.max_steps // (
len(train_dataloader) // args.gradient_accumulation_steps) + 1
else:
t_total = len(
train_dataloader
) // args.gradient_accumulation_steps * args.num_train_epochs
# Prepare optimizer and schedule (linear warmup and decay)
no_decay = ["bias", "LayerNorm.weight"]
optimizer_grouped_parameters = [
{
"params": [
p for n, p in model.named_parameters()
if not any(nd in n for nd in no_decay)
],
"weight_decay":
args.weight_decay,
},
{
"params": [
p for n, p in model.named_parameters()
if any(nd in n for nd in no_decay)
],
"weight_decay":
0.0
},
]
args.warmup_steps = int(t_total * args.warmup_proportion)
optimizer = AdamW(optimizer_grouped_parameters,
lr=args.learning_rate,
eps=args.adam_epsilon)
scheduler = get_linear_schedule_with_warmup(
optimizer,
num_warmup_steps=args.warmup_steps,
num_training_steps=t_total)
# Check if saved optimizer or scheduler states exist
if os.path.isfile(os.path.join(
args.model_name_or_path, "optimizer.pt")) and os.path.isfile(
os.path.join(args.model_name_or_path, "scheduler.pt")):
# Load in optimizer and scheduler states
optimizer.load_state_dict(
torch.load(os.path.join(args.model_name_or_path, "optimizer.pt")))
scheduler.load_state_dict(
torch.load(os.path.join(args.model_name_or_path, "scheduler.pt")))
if args.fp16:
try:
from apex import amp
except ImportError:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to use fp16 training."
)
model, optimizer = amp.initialize(model,
optimizer,
opt_level=args.fp16_opt_level)
# multi-gpu training (should be after apex fp16 initialization)
if args.n_gpu > 1:
model = torch.nn.DataParallel(model)
# Distributed training (should be after apex fp16 initialization)
if args.local_rank != -1:
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[args.local_rank],
output_device=args.local_rank,
find_unused_parameters=True)
# Train!
logger.info("***** Running training *****")
logger.info(" Num examples = %d", len(train_dataset))
logger.info(" Num Epochs = %d", args.num_train_epochs)
logger.info(" Instantaneous batch size per GPU = %d",
args.per_gpu_train_batch_size)
logger.info(
" Total train batch size (w. parallel, distributed & accumulation) = %d",
args.train_batch_size * args.gradient_accumulation_steps *
(torch.distributed.get_world_size() if args.local_rank != -1 else 1),
)
logger.info(" Gradient Accumulation steps = %d",
args.gradient_accumulation_steps)
logger.info(" Total optimization steps = %d", t_total)
global_step = 0
steps_trained_in_current_epoch = 0
# Check if continuing training from a checkpoint
if os.path.exists(args.model_name_or_path
) and "checkpoint" in args.model_name_or_path:
# set global_step to gobal_step of last saved checkpoint from model path
global_step = int(args.model_name_or_path.split("-")[-1].split("/")[0])
epochs_trained = global_step // (len(train_dataloader) //
args.gradient_accumulation_steps)
steps_trained_in_current_epoch = global_step % (
len(train_dataloader) // args.gradient_accumulation_steps)
logger.info(
" Continuing training from checkpoint, will skip to saved global_step"
)
logger.info(" Continuing training from epoch %d", epochs_trained)
logger.info(" Continuing training from global step %d", global_step)
logger.info(" Will skip the first %d steps in the first epoch",
steps_trained_in_current_epoch)
tr_loss, logging_loss = 0.0, 0.0
model.zero_grad()
seed_everything(
args.seed
) # Added here for reproductibility (even between python 2 and 3)
for _ in range(int(args.num_train_epochs)):
pbar = ProgressBar(n_total=len(train_dataloader), desc='Training')
for step, batch in enumerate(train_dataloader):
# Skip past any already trained steps if resuming training
if steps_trained_in_current_epoch > 0:
steps_trained_in_current_epoch -= 1
continue
model.train()
batch = tuple(t.to(args.device) for t in batch)
inputs = {
"input_ids": batch[0],
"attention_mask": batch[1],
"start_positions": batch[3],
"end_positions": batch[4]
}
if args.model_type != "distilbert":
# XLM and RoBERTa don"t use segment_ids
inputs["token_type_ids"] = (batch[2] if args.model_type
in ["bert", "xlnet"] else None)
outputs = model(**inputs)
loss = outputs[
0] # model outputs are always tuple in pytorch-transformers (see doc)
if args.n_gpu > 1:
loss = loss.mean(
) # mean() to average on multi-gpu parallel training
if args.gradient_accumulation_steps > 1:
loss = loss / args.gradient_accumulation_steps
if args.fp16:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
pbar(step, {'loss': loss.item()})
tr_loss += loss.item()
if (step + 1) % args.gradient_accumulation_steps == 0:
if args.fp16:
torch.nn.utils.clip_grad_norm_(
amp.master_params(optimizer), args.max_grad_norm)
else:
torch.nn.utils.clip_grad_norm_(model.parameters(),
args.max_grad_norm)
scheduler.step() # Update learning rate schedule
optimizer.step()
model.zero_grad()
global_step += 1
if args.local_rank in [
-1, 0
] and args.logging_steps > 0 and global_step % args.logging_steps == 0:
# Log metrics
print(" ")
if args.local_rank == -1:
# Only evaluate when single GPU otherwise metrics may not average well
evaluate(args, model, tokenizer)
if args.local_rank in [
-1, 0
] and args.save_steps > 0 and global_step % args.save_steps == 0:
# Save model checkpoint
output_dir = os.path.join(
args.output_dir, "checkpoint-{}".format(global_step))
if not os.path.exists(output_dir):
os.makedirs(output_dir)
model_to_save = (
model.module if hasattr(model, "module") else model
) # Take care of distributed/parallel training
model_to_save.save_pretrained(output_dir)
torch.save(args,
os.path.join(output_dir, "training_args.bin"))
tokenizer.save_vocabulary(output_dir)
logger.info("Saving model checkpoint to %s", output_dir)
torch.save(optimizer.state_dict(),
os.path.join(output_dir, "optimizer.pt"))
torch.save(scheduler.state_dict(),
os.path.join(output_dir, "scheduler.pt"))
logger.info("Saving optimizer and scheduler states to %s",
output_dir)
print(" ")
if 'cuda' in str(args.device):
torch.cuda.empty_cache()
return global_step, tr_loss / global_step
def main():
parser = argparse.ArgumentParser()
# Required parameters
parser.add_argument(
"--task_name",
default=None,
type=str,
required=True,
help="The name of the task to train selected in the list: " +
", ".join(processors.keys()))
parser.add_argument(
"--data_dir",
default=None,
type=str,
required=True,
help=
"The input data dir. Should contain the training files for the CoNLL-2003 NER task.",
)
parser.add_argument(
"--model_type",
default=None,
type=str,
required=True,
help="Model type selected in the list: " +
", ".join(MODEL_CLASSES.keys()),
)
parser.add_argument(
"--model_name_or_path",
default=None,
type=str,
required=True,
help="Path to pre-trained model or shortcut name selected in the list: "
+ ", ".join(ALL_MODELS),
)
parser.add_argument(
"--output_dir",
default=None,
type=str,
required=True,
help=
"The output directory where the model predictions and checkpoints will be written.",
)
# Other parameters
parser.add_argument('--markup',
default='bios',
type=str,
choices=['bios', 'bio'])
parser.add_argument('--loss_type',
default='ce',
type=str,
choices=['lsr', 'focal', 'ce'])
parser.add_argument(
"--labels",
default="",
type=str,
help=
"Path to a file containing all labels. If not specified, CoNLL-2003 labels are used.",
)
parser.add_argument(
"--config_name",
default="",
type=str,
help="Pretrained config name or path if not the same as model_name")
parser.add_argument(
"--tokenizer_name",
default="",
type=str,
help="Pretrained tokenizer name or path if not the same as model_name",
)
parser.add_argument(
"--cache_dir",
default="",
type=str,
help=
"Where do you want to store the pre-trained models downloaded from s3",
)
parser.add_argument(
"--train_max_seq_length",
default=128,
type=int,
help=
"The maximum total input sequence length after tokenization. Sequences longer "
"than this will be truncated, sequences shorter will be padded.",
)
parser.add_argument(
"--eval_max_seq_length",
default=512,
type=int,
help=
"The maximum total input sequence length after tokenization. Sequences longer "
"than this will be truncated, sequences shorter will be padded.",
)
parser.add_argument("--do_train",
action="store_true",
help="Whether to run training.")
parser.add_argument("--do_eval",
action="store_true",
help="Whether to run eval on the dev set.")
parser.add_argument("--do_predict",
action="store_true",
help="Whether to run predictions on the test set.")
parser.add_argument(
"--evaluate_during_training",
action="store_true",
help="Whether to run evaluation during training at each logging step.",
)
parser.add_argument(
"--do_lower_case",
action="store_true",
help="Set this flag if you are using an uncased model.")
parser.add_argument("--per_gpu_train_batch_size",
default=8,
type=int,
help="Batch size per GPU/CPU for training.")
parser.add_argument("--per_gpu_eval_batch_size",
default=8,
type=int,
help="Batch size per GPU/CPU for evaluation.")
parser.add_argument(
"--gradient_accumulation_steps",
type=int,
default=1,
help=
"Number of updates steps to accumulate before performing a backward/update pass.",
)
parser.add_argument("--learning_rate",
default=5e-5,
type=float,
help="The initial learning rate for Adam.")
parser.add_argument("--weight_decay",
default=0.01,
type=float,
help="Weight decay if we apply some.")
parser.add_argument("--adam_epsilon",
default=1e-8,
type=float,
help="Epsilon for Adam optimizer.")
parser.add_argument("--max_grad_norm",
default=1.0,
type=float,
help="Max gradient norm.")
parser.add_argument("--num_train_epochs",
default=3.0,
type=float,
help="Total number of training epochs to perform.")
parser.add_argument(
"--max_steps",
default=-1,
type=int,
help=
"If > 0: set total number of training steps to perform. Override num_train_epochs.",
)
parser.add_argument(
"--warmup_proportion",
default=0.1,
type=float,
help=
"Proportion of training to perform linear learning rate warmup for,E.g., 0.1 = 10% of training."
)
parser.add_argument("--logging_steps",
type=int,
default=50,
help="Log every X updates steps.")
parser.add_argument("--save_steps",
type=int,
default=50,
help="Save checkpoint every X updates steps.")
parser.add_argument(
"--eval_all_checkpoints",
action="store_true",
help=
"Evaluate all checkpoints starting with the same prefix as model_name ending and ending with step number",
)
parser.add_argument(
'--predict_all_checkpoints',
action="store_true",
help=
"Predict all checkpoints starting with the same prefix as model_name ending and ending with step number",
)
parser.add_argument("--no_cuda",
action="store_true",
help="Avoid using CUDA when available")
parser.add_argument("--overwrite_output_dir",
action="store_true",
help="Overwrite the content of the output directory")
parser.add_argument(
"--overwrite_cache",
action="store_true",
help="Overwrite the cached training and evaluation sets")
parser.add_argument("--seed",
type=int,
default=42,
help="random seed for initialization")
parser.add_argument(
"--fp16",
action="store_true",
help=
"Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit",
)
parser.add_argument(
"--fp16_opt_level",
type=str,
default="O1",
help=
"For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']."
"See details at https://nvidia.github.io/apex/amp.html",
)
parser.add_argument("--local_rank",
type=int,
default=-1,
help="For distributed training: local_rank")
parser.add_argument("--server_ip",
type=str,
default="",
help="For distant debugging.")
parser.add_argument("--server_port",
type=str,
default="",
help="For distant debugging.")
args = parser.parse_args()
args.output_dir = args.output_dir + '{}'.format(args.model_type)
if not os.path.exists(args.output_dir):
os.mkdir(args.output_dir)
init_logger(log_file=args.output_dir + '/{}-{}-{}.log'.format(
args.model_type, args.task_name,
time.strftime("%Y-%m-%d-%H:%M:%S", time.localtime())))
if os.path.exists(args.output_dir) and os.listdir(
args.output_dir
) and args.do_train and not args.overwrite_output_dir:
raise ValueError(
"Output directory ({}) already exists and is not empty. Use --overwrite_output_dir to overcome."
.format(args.output_dir))
# Setup distant debugging if needed
if args.server_ip and args.server_port:
# Distant debugging - see https://code.visualstudio.com/docs/python/debugging#_attach-to-a-local-script
import ptvsd
print("Waiting for debugger attach")
ptvsd.enable_attach(address=(args.server_ip, args.server_port),
redirect_output=True)
ptvsd.wait_for_attach()
# Setup CUDA, GPU & distributed training
if args.local_rank == -1 or args.no_cuda:
device = torch.device("cuda" if torch.cuda.is_available()
and not args.no_cuda else "cpu")
args.n_gpu = torch.cuda.device_count()
else: # Initializes the distributed backend which will take care of sychronizing nodes/GPUs
torch.cuda.set_device(args.local_rank)
device = torch.device("cuda", args.local_rank)
torch.distributed.init_process_group(backend="nccl")
args.n_gpu = 1
args.device = device
logger.warning(
"Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s",
args.local_rank,
device,
args.n_gpu,
bool(args.local_rank != -1),
args.fp16,
)
# Set seed
seed_everything(args.seed)
# Prepare NER task
args.task_name = args.task_name.lower()
if args.task_name not in processors:
raise ValueError("Task not found: %s" % (args.task_name))
processor = processors[args.task_name]()
label_list = processor.get_labels()
args.id2label = {i: label for i, label in enumerate(label_list)}
args.label2id = {label: i for i, label in enumerate(label_list)}
num_labels = len(label_list)
# Load pretrained model and tokenizer
if args.local_rank not in [-1, 0]:
torch.distributed.barrier(
) # Make sure only the first process in distributed training will download model & vocab
args.model_type = args.model_type.lower()
config_class, model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
config = config_class.from_pretrained(
args.config_name if args.config_name else args.model_name_or_path,
num_labels=num_labels,
loss_type=args.loss_type,
cache_dir=args.cache_dir if args.cache_dir else None,
soft_label=True)
tokenizer = tokenizer_class.from_pretrained(
args.tokenizer_name
if args.tokenizer_name else args.model_name_or_path,
do_lower_case=args.do_lower_case,
cache_dir=args.cache_dir if args.cache_dir else None,
)
model = model_class.from_pretrained(
args.model_name_or_path,
from_tf=bool(".ckpt" in args.model_name_or_path),
config=config)
if args.local_rank == 0:
torch.distributed.barrier(
) # Make sure only the first process in distributed training will download model & vocab
model.to(args.device)
logger.info("Training/evaluation parameters %s", args)
# Training
if args.do_train:
train_dataset = load_and_cache_examples(args,
args.task_name,
tokenizer,
data_type='train')
global_step, tr_loss = train(args, train_dataset, model, tokenizer)
logger.info(" global_step = %s, average loss = %s", global_step,
tr_loss)
# Saving best-practices: if you use defaults names for the model, you can reload it using from_pretrained()
if args.do_train and (args.local_rank == -1
or torch.distributed.get_rank() == 0):
# Create output directory if needed
if not os.path.exists(args.output_dir) and args.local_rank in [-1, 0]:
os.makedirs(args.output_dir)
logger.info("Saving model checkpoint to %s", args.output_dir)
# Save a trained model, configuration and tokenizer using `save_pretrained()`.
# They can then be reloaded using `from_pretrained()`
model_to_save = (model.module if hasattr(model, "module") else model
) # Take care of distributed/parallel training
model_to_save.save_pretrained(args.output_dir)
tokenizer.save_vocabulary(args.output_dir)
# Good practice: save your training arguments together with the trained model
torch.save(args, os.path.join(args.output_dir, "training_args.bin"))
#Evaluation
results = {}
if args.do_eval and args.local_rank in [-1, 0]:
tokenizer = tokenizer_class.from_pretrained(
args.output_dir, do_lower_case=args.do_lower_case)
checkpoints = [args.output_dir]
if args.eval_all_checkpoints:
checkpoints = list(
os.path.dirname(c) for c in sorted(
glob.glob(args.output_dir + "/**/" + WEIGHTS_NAME,
recursive=True)))
logging.getLogger("pytorch_transformers.modeling_utils").setLevel(
logging.WARN) # Reduce logging
logger.info("Evaluate the following checkpoints: %s", checkpoints)
for checkpoint in checkpoints:
global_step = checkpoint.split(
"-")[-1] if len(checkpoints) > 1 else ""
prefix = checkpoint.split(
'/')[-1] if checkpoint.find('checkpoint') != -1 else ""
model = model_class.from_pretrained(checkpoint)
model.to(args.device)
result = evaluate(args, model, tokenizer, prefix=prefix)
if global_step:
result = {
"{}_{}".format(global_step, k): v
for k, v in result.items()
}
results.update(result)
output_eval_file = os.path.join(args.output_dir, "eval_results.txt")
with open(output_eval_file, "w") as writer:
for key in sorted(results.keys()):
writer.write("{} = {}\n".format(key, str(results[key])))
if args.do_predict and args.local_rank in [-1, 0]:
tokenizer = tokenizer_class.from_pretrained(
args.output_dir, do_lower_case=args.do_lower_case)
checkpoints = [args.output_dir]
if args.predict_all_checkpoints > 0:
checkpoints = list(
os.path.dirname(c) for c in sorted(
glob.glob(args.output_dir + '/**/' + WEIGHTS_NAME,
recursive=True)))
logging.getLogger("transformers.modeling_utils").setLevel(
logging.WARN) # Reduce logging
checkpoints = [
x for x in checkpoints
if x.split('-')[-1] == str(args.predict_checkpoints)
]
logger.info("Predict the following checkpoints: %s", checkpoints)
for checkpoint in checkpoints:
prefix = checkpoint.split(
'/')[-1] if checkpoint.find('checkpoint') != -1 else ""
model = model_class.from_pretrained(checkpoint)
model.to(args.device)
predict(args, model, tokenizer, prefix=prefix)
def evaluate(args, model, tokenizer, prefix=""):
metric = SpanEntityScore(args.id2label)
eval_output_dir = args.output_dir
if not os.path.exists(eval_output_dir) and args.local_rank in [-1, 0]:
os.makedirs(eval_output_dir)
eval_features = load_and_cache_examples(args,
args.task_name,
tokenizer,
data_type='dev')
args.eval_batch_size = args.per_gpu_eval_batch_size * max(1, args.n_gpu)
# Eval!
logger.info("***** Running evaluation %s *****", prefix)
logger.info(" Num examples = %d", len(eval_features))
logger.info(" Batch size = %d", args.eval_batch_size)
eval_loss = 0.0
nb_eval_steps = 0
pbar = ProgressBar(n_total=len(eval_features), desc="Evaluating")
for step, f in enumerate(eval_features):
input_lens = f.input_len
input_ids = torch.tensor([f.input_ids[:input_lens]],
dtype=torch.long).to(args.device)
input_mask = torch.tensor([f.input_mask[:input_lens]],
dtype=torch.long).to(args.device)
segment_ids = torch.tensor([f.segment_ids[:input_lens]],
dtype=torch.long).to(args.device)
start_ids = torch.tensor([f.start_ids[:input_lens]],
dtype=torch.long).to(args.device)
end_ids = torch.tensor([f.end_ids[:input_lens]],
dtype=torch.long).to(args.device)
subjects = f.subjects
model.eval()
with torch.no_grad():
inputs = {
"input_ids": input_ids,
"attention_mask": input_mask,
"start_positions": start_ids,
"end_positions": end_ids
}
if args.model_type != "distilbert":
# XLM and RoBERTa don"t use segment_ids
inputs["token_type_ids"] = (segment_ids if args.model_type
in ["bert", "xlnet"] else None)
outputs = model(**inputs)
tmp_eval_loss, start_logits, end_logits = outputs[:3]
R = bert_extract_item(start_logits, end_logits)
T = subjects
metric.update(true_subject=T, pred_subject=R)
if args.n_gpu > 1:
tmp_eval_loss = tmp_eval_loss.mean(
) # mean() to average on multi-gpu parallel evaluating
eval_loss += tmp_eval_loss.item()
nb_eval_steps += 1
pbar(step)
print(' ')
eval_loss = eval_loss / nb_eval_steps
eval_info, entity_info = metric.result()
results = {f'{key}': value for key, value in eval_info.items()}
results['loss'] = eval_loss
logger.info("***** Eval results %s *****", prefix)
info = "-".join(
[f' {key}: {value:.4f} ' for key, value in results.items()])
logger.info(info)
logger.info("***** Entity results %s *****", prefix)
for key in sorted(entity_info.keys()):
logger.info("******* %s results ********" % key)
info = "-".join([
f' {key}: {value:.4f} ' for key, value in entity_info[key].items()
])
logger.info(info)
return results
def evaluate(args, model, tokenizer, prefix=""):
metric = SeqEntityScore(args.id2label, markup=args.markup)
eval_output_dir = args.output_dir
if not os.path.exists(eval_output_dir) and args.local_rank in [-1, 0]:
os.makedirs(eval_output_dir)
eval_dataset = load_and_cache_examples(args,
args.task_name,
tokenizer,
data_type='dev')
args.eval_batch_size = args.per_gpu_eval_batch_size * max(1, args.n_gpu)
# Note that DistributedSampler samples randomly
eval_sampler = SequentialSampler(
eval_dataset) if args.local_rank == -1 else DistributedSampler(
eval_dataset)
eval_dataloader = DataLoader(eval_dataset,
sampler=eval_sampler,
batch_size=args.eval_batch_size,
collate_fn=collate_fn)
# Eval!
logger.info("***** Running evaluation %s *****", prefix)
logger.info(" Num examples = %d", len(eval_dataset))
logger.info(" Batch size = %d", args.eval_batch_size)
eval_loss = 0.0
nb_eval_steps = 0
pbar = ProgressBar(n_total=len(eval_dataloader), desc="Evaluating")
if isinstance(model, nn.DataParallel):
model = model.module
for step, batch in enumerate(eval_dataloader):
model.eval()
batch = tuple(t.to(args.device) for t in batch)
with torch.no_grad():
inputs = {
"input_ids": batch[0],
"attention_mask": batch[1],
"labels": batch[3],
'input_lens': batch[4]
}
if args.model_type != "distilbert":
# XLM and RoBERTa don"t use segment_ids
inputs["token_type_ids"] = (batch[2] if args.model_type
in ["bert", "xlnet"] else None)
outputs = model(**inputs)
tmp_eval_loss, logits = outputs[:2]
tags, _ = model.crf._obtain_labels(logits, args.id2label,
inputs['input_lens'])
if args.n_gpu > 1:
tmp_eval_loss = tmp_eval_loss.mean(
) # mean() to average on multi-gpu parallel evaluating
eval_loss += tmp_eval_loss.item()
nb_eval_steps += 1
out_label_ids = inputs['labels'].cpu().numpy().tolist()
for i, label in enumerate(out_label_ids):
temp_1 = []
temp_2 = []
for j, m in enumerate(label):
if j == 0:
continue
elif out_label_ids[i][j] == args.label2id['[SEP]']:
metric.update(pred_paths=[temp_2], label_paths=[temp_1])
break
else:
temp_1.append(args.id2label[out_label_ids[i][j]])
temp_2.append(tags[i][j])
pbar(step)
print(' ')
eval_loss = eval_loss / nb_eval_steps
eval_info, entity_info = metric.result()
results = {f'{key}': value for key, value in eval_info.items()}
results['loss'] = eval_loss
logger.info("***** Eval results %s *****", prefix)
info = "-".join(
[f' {key}: {value:.4f} ' for key, value in results.items()])
logger.info(info)
logger.info("***** Entity results %s *****", prefix)
for key in sorted(entity_info.keys()):
logger.info("******* %s results ********" % key)
info = "-".join([
f' {key}: {value:.4f} ' for key, value in entity_info[key].items()
])
logger.info(info)
return results
from snlp.trainers import Trainer
from snlp.tools.log import logger
os.environ['CUDA_VISIBLE_DEVICES'] = "1"
fixed_length_uttr = 20
fixed_length_resp = 20
fixed_length_turn = 5
start = time.time()
cls_task = tasks.Classification(num_classes=2, losses = nn.CrossEntropyLoss())
cls_task.metrics = ['accuracy']
# 对数据进行预处理
file = "../sample_data/multi_qa.csv"
logger.info("读取数据 %s" % os.path.basename(file))
data = pd.read_csv(file)
data['label'] = 1
# 对数据进行预处理
## 合并训练预处理器是一种leaky
logger.info("使用Preprocessor处理数据")
preprocessor = CNPreprocessorForMultiQA(stopwords=['\t'])
preprocessor = preprocessor.fit(data, columns=['utterances', 'response'])
data = preprocessor.transform(data)
data = data[['D_num', 'turns', 'utterances', 'response', 'utterances_len', 'response_len', 'label']]
data['label'] = data['label'].astype(int)
# 划分训练集和测试集
train = data[:90]
valid = data[90:]
def main():
parser = ArgumentParser()
## Required parameters
parser.add_argument(
"--data_dir",
default=None,
type=str,
required=True,
help=
"The input data dir. Should contain the .tsv files (or other data files) for the task."
)
parser.add_argument("--config_path", default=None, type=str, required=True)
parser.add_argument("--vocab_path", default=None, type=str, required=True)
parser.add_argument(
"--output_dir",
default=None,
type=str,
required=True,
help=
"The output directory where the model predictions and checkpoints will be written."
)
parser.add_argument("--model_path", default='', type=str)
parser.add_argument('--data_name', default='albert', type=str)
parser.add_argument(
"--file_num",
type=int,
default=10,
help="Number of dynamic masking to pregenerate (with different masks)")
parser.add_argument(
"--reduce_memory",
action="store_true",
help=
"Store training data as on-disc memmaps to massively reduce memory usage"
)
parser.add_argument("--epochs",
type=int,
default=4,
help="Number of epochs to train for")
parser.add_argument(
"--do_lower_case",
action='store_true',
help="Set this flag if you are using an uncased model.")
parser.add_argument('--num_eval_steps', default=100)
parser.add_argument('--num_save_steps', default=200)
parser.add_argument("--local_rank",
type=int,
default=-1,
help="local_rank for distributed training on gpus")
parser.add_argument("--weight_decay",
default=0.01,
type=float,
help="Weight deay if we apply some.")
parser.add_argument("--no_cuda",
action='store_true',
help="Whether not to use CUDA when available")
parser.add_argument(
'--gradient_accumulation_steps',
type=int,
default=1,
help=
"Number of updates steps to accumulate before performing a backward/update pass."
)
parser.add_argument("--train_batch_size",
default=4,
type=int,
help="Total batch size for training.")
parser.add_argument(
'--loss_scale',
type=float,
default=0,
help=
"Loss scaling to improve fp16 numeric stability. Only used when fp16 set to True.\n"
"0 (default value): dynamic loss scaling.\n"
"Positive power of 2: static loss scaling value.\n")
parser.add_argument("--warmup_proportion",
default=0.1,
type=float,
help="Linear warmup over warmup_steps.")
parser.add_argument("--adam_epsilon",
default=1e-8,
type=float,
help="Epsilon for Adam optimizer.")
parser.add_argument('--max_grad_norm', default=1.0, type=float)
parser.add_argument("--learning_rate",
default=0.00176,
type=float,
help="The initial learning rate for Adam.")
parser.add_argument('--seed',
type=int,
default=42,
help="random seed for initialization")
parser.add_argument(
'--fp16_opt_level',
type=str,
default='O2',
help=
"For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']."
"See details at https://nvidia.github.io/apex/amp.html")
parser.add_argument(
'--fp16',
action='store_true',
help="Whether to use 16-bit float precision instead of 32-bit")
args = parser.parse_args()
args.data_dir = Path(args.data_dir)
args.output_dir = Path(args.output_dir)
pregenerated_data = args.data_dir / "corpus/train"
init_logger(log_file=str(args.output_dir / "train_albert_model.log"))
assert pregenerated_data.is_dir(), \
"--pregenerated_data should point to the folder of files made by prepare_lm_data_mask.py!"
samples_per_epoch = 0
for i in range(args.file_num):
data_file = pregenerated_data / f"{args.data_name}_file_{i}.json"
metrics_file = pregenerated_data / f"{args.data_name}_file_{i}_metrics.json"
if data_file.is_file() and metrics_file.is_file():
metrics = json.loads(metrics_file.read_text())
samples_per_epoch += metrics['num_training_examples']
else:
if i == 0:
exit("No training data was found!")
print(
f"Warning! There are fewer epochs of pregenerated data ({i}) than training epochs ({args.epochs})."
)
print(
"This script will loop over the available data, but training diversity may be negatively impacted."
)
break
logger.info(f"samples_per_epoch: {samples_per_epoch}")
if args.local_rank == -1 or args.no_cuda:
device = torch.device(f"cuda" if torch.cuda.is_available()
and not args.no_cuda else "cpu")
args.n_gpu = torch.cuda.device_count()
else:
torch.cuda.set_device(args.local_rank)
device = torch.device("cuda", args.local_rank)
args.n_gpu = 1
# Initializes the distributed backend which will take care of sychronizing nodes/GPUs
torch.distributed.init_process_group(backend='nccl')
logger.info(
f"device: {device} , distributed training: {bool(args.local_rank != -1)}, 16-bits training: {args.fp16}, "
f"share_type: {args.share_type}")
if args.gradient_accumulation_steps < 1:
raise ValueError(
f"Invalid gradient_accumulation_steps parameter: {args.gradient_accumulation_steps}, should be >= 1"
)
args.train_batch_size = args.train_batch_size // args.gradient_accumulation_steps
seed_everything(args.seed)
tokenizer = BertTokenizer.from_pretrained(args.vocab_path,
do_lower_case=args.do_lower_case)
total_train_examples = samples_per_epoch * args.epochs
num_train_optimization_steps = int(total_train_examples /
args.train_batch_size /
args.gradient_accumulation_steps)
if args.local_rank != -1:
num_train_optimization_steps = num_train_optimization_steps // torch.distributed.get_world_size(
)
args.warmup_steps = int(num_train_optimization_steps *
args.warmup_proportion)
bert_config = AlbertConfig.from_pretrained(args.config_path,
share_type=args.share_type)
model = AlbertForPreTraining(config=bert_config)
if args.model_path:
model = AlbertForPreTraining.from_pretrained(args.model_path)
model.to(device)
# Prepare optimizer
param_optimizer = list(model.named_parameters())
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':
args.weight_decay
}, {
'params':
[p for n, p in param_optimizer if any(nd in n for nd in no_decay)],
'weight_decay':
0.0
}]
optimizer = AdamW(params=optimizer_grouped_parameters,
lr=args.learning_rate,
eps=args.adam_epsilon)
scheduler = get_linear_schedule_with_warmup(
optimizer,
num_warmup_steps=args.warmup_steps,
num_training_steps=num_train_optimization_steps)
# optimizer = Lamb(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
if args.fp16:
try:
from apex import amp
except ImportError:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to use fp16 training."
)
model, optimizer = amp.initialize(model,
optimizer,
opt_level=args.fp16_opt_level)
if args.n_gpu > 1:
model = torch.nn.DataParallel(model)
if args.local_rank != -1:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.local_rank], output_device=args.local_rank)
global_step = 0
mask_metric = LMAccuracy()
sop_metric = LMAccuracy()
tr_mask_acc = AverageMeter()
tr_sop_acc = AverageMeter()
tr_loss = AverageMeter()
tr_mask_loss = AverageMeter()
tr_sop_loss = AverageMeter()
loss_fct = CrossEntropyLoss(ignore_index=-1)
train_logs = {}
logger.info("***** Running training *****")
logger.info(f" Num examples = {total_train_examples}")
logger.info(f" Batch size = {args.train_batch_size}")
logger.info(f" Num steps = {num_train_optimization_steps}")
logger.info(f" warmup_steps = {args.warmup_steps}")
start_time = time.time()
seed_everything(args.seed) # Added here for reproducibility
for epoch in range(args.epochs):
for idx in range(args.file_num):
epoch_dataset = PregeneratedDataset(
file_id=idx,
training_path=pregenerated_data,
tokenizer=tokenizer,
reduce_memory=args.reduce_memory,
data_name=args.data_name)
if args.local_rank == -1:
train_sampler = RandomSampler(epoch_dataset)
else:
train_sampler = DistributedSampler(epoch_dataset)
train_dataloader = DataLoader(epoch_dataset,
sampler=train_sampler,
batch_size=args.train_batch_size)
model.train()
nb_tr_examples, nb_tr_steps = 0, 0
for step, batch in enumerate(train_dataloader):
batch = tuple(t.to(device) for t in batch)
input_ids, input_mask, segment_ids, lm_label_ids, is_next = batch
outputs = model(input_ids=input_ids,
token_type_ids=segment_ids,
attention_mask=input_mask)
prediction_scores = outputs[0]
seq_relationship_score = outputs[1]
masked_lm_loss = loss_fct(
prediction_scores.view(-1, bert_config.vocab_size),
lm_label_ids.view(-1))
next_sentence_loss = loss_fct(
seq_relationship_score.view(-1, 2), is_next.view(-1))
loss = masked_lm_loss + next_sentence_loss
mask_metric(logits=prediction_scores.view(
-1, bert_config.vocab_size),
target=lm_label_ids.view(-1))
sop_metric(logits=seq_relationship_score.view(-1, 2),
target=is_next.view(-1))
if args.n_gpu > 1:
loss = loss.mean() # mean() to average on multi-gpu.
if args.gradient_accumulation_steps > 1:
loss = loss / args.gradient_accumulation_steps
if args.fp16:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
nb_tr_steps += 1
tr_mask_acc.update(mask_metric.value(), n=input_ids.size(0))
tr_sop_acc.update(sop_metric.value(), n=input_ids.size(0))
tr_loss.update(loss.item(), n=1)
tr_mask_loss.update(masked_lm_loss.item(), n=1)
tr_sop_loss.update(next_sentence_loss.item(), n=1)
if (step + 1) % args.gradient_accumulation_steps == 0:
if args.fp16:
torch.nn.utils.clip_grad_norm_(
amp.master_params(optimizer), args.max_grad_norm)
else:
torch.nn.utils.clip_grad_norm_(model.parameters(),
args.max_grad_norm)
scheduler.step()
optimizer.step()
optimizer.zero_grad()
global_step += 1
if global_step % args.num_eval_steps == 0:
now = time.time()
eta = now - start_time
if eta > 3600:
eta_format = ('%d:%02d:%02d' %
(eta // 3600,
(eta % 3600) // 60, eta % 60))
elif eta > 60:
eta_format = '%d:%02d' % (eta // 60, eta % 60)
else:
eta_format = '%ds' % eta
train_logs['loss'] = tr_loss.avg
train_logs['mask_acc'] = tr_mask_acc.avg
train_logs['sop_acc'] = tr_sop_acc.avg
train_logs['mask_loss'] = tr_mask_loss.avg
train_logs['sop_loss'] = tr_sop_loss.avg
show_info = f'[Training]:[{epoch}/{args.epochs}]{global_step}/{num_train_optimization_steps} ' \
f'- ETA: {eta_format}' + "-".join(
[f' {key}: {value:.4f} ' for key, value in train_logs.items()])
logger.info(show_info)
tr_mask_acc.reset()
tr_sop_acc.reset()
tr_loss.reset()
tr_mask_loss.reset()
tr_sop_loss.reset()
start_time = now
if global_step % args.num_save_steps == 0:
if args.local_rank in [-1, 0] and args.num_save_steps > 0:
# Save model checkpoint
output_dir = args.output_dir / f'lm-checkpoint-{global_step}'
if not output_dir.exists():
output_dir.mkdir()
# save model
model_to_save = model.module if hasattr(
model, 'module'
) else model # Take care of distributed/parallel training
model_to_save.save_pretrained(str(output_dir))
torch.save(args, str(output_dir / 'training_args.bin'))
logger.info("Saving model checkpoint to %s",
output_dir)
# save config
output_config_file = output_dir / CONFIG_NAME
with open(str(output_config_file), 'w') as f:
f.write(model_to_save.config.to_json_string())
# save vocab
tokenizer.save_vocabulary(output_dir)
def __init__(self,
training_path,
file_id,
tokenizer,
data_name,
reduce_memory=False):
self.tokenizer = tokenizer
self.file_id = file_id
data_file = training_path / f"{data_name}_file_{self.file_id}.json"
metrics_file = training_path / f"{data_name}_file_{self.file_id}_metrics.json"
assert data_file.is_file() and metrics_file.is_file()
metrics = json.loads(metrics_file.read_text())
num_samples = metrics['num_training_examples']
seq_len = metrics['max_seq_len']
self.temp_dir = None
self.working_dir = None
if reduce_memory:
self.temp_dir = TemporaryDirectory()
self.working_dir = Path(self.temp_dir.name)
input_ids = np.memmap(filename=self.working_dir /
'input_ids.memmap',
mode='w+',
dtype=np.int32,
shape=(num_samples, seq_len))
input_masks = np.memmap(filename=self.working_dir /
'input_masks.memmap',
shape=(num_samples, seq_len),
mode='w+',
dtype=np.bool)
segment_ids = np.memmap(filename=self.working_dir /
'segment_ids.memmap',
shape=(num_samples, seq_len),
mode='w+',
dtype=np.bool)
lm_label_ids = np.memmap(filename=self.working_dir /
'lm_label_ids.memmap',
shape=(num_samples, seq_len),
mode='w+',
dtype=np.int32)
lm_label_ids[:] = -1
is_nexts = np.memmap(filename=self.working_dir / 'is_nexts.memmap',
shape=(num_samples, ),
mode='w+',
dtype=np.bool)
else:
input_ids = np.zeros(shape=(num_samples, seq_len), dtype=np.int32)
input_masks = np.zeros(shape=(num_samples, seq_len), dtype=np.bool)
segment_ids = np.zeros(shape=(num_samples, seq_len), dtype=np.bool)
lm_label_ids = np.full(shape=(num_samples, seq_len),
dtype=np.int32,
fill_value=-1)
is_nexts = np.zeros(shape=(num_samples, ), dtype=np.bool)
logger.info(f"Loading training examples for {str(data_file)}")
with data_file.open() as f:
for i, line in enumerate(f):
line = line.strip()
example = json.loads(line)
features = convert_example_to_features(example, tokenizer,
seq_len)
input_ids[i] = features.input_ids
segment_ids[i] = features.segment_ids
input_masks[i] = features.input_mask
lm_label_ids[i] = features.lm_label_ids
is_nexts[i] = features.is_next
assert i == num_samples - 1 # Assert that the sample count metric was true
logger.info("Loading complete!")
self.num_samples = num_samples
self.seq_len = seq_len
self.input_ids = input_ids
self.input_masks = input_masks
self.segment_ids = segment_ids
self.lm_label_ids = lm_label_ids
self.is_nexts = is_nexts
def main():
parser = argparse.ArgumentParser()
## Required parameters
parser.add_argument(
"--data_dir",
default=None,
type=str,
required=True,
help=
"The input data dir. Should contain the .tsv files (or other data files) for the task."
)
parser.add_argument("--model_type",
default=None,
type=str,
required=True,
help="Model type selected in the list: ")
parser.add_argument(
"--model_name_or_path",
default=None,
type=str,
required=True,
help="Path to pre-trained model or shortcut name selected in the list")
parser.add_argument(
"--task_name",
default=None,
type=str,
required=True,
help="The name of the task to train selected in the list: " +
", ".join(processors.keys()))
parser.add_argument(
"--output_dir",
default=None,
type=str,
required=True,
help=
"The output directory where the model predictions and checkpoints will be written."
)
parser.add_argument("--vocab_file", default='', type=str)
parser.add_argument("--spm_model_file", default='', type=str)
## Other parameters
parser.add_argument(
"--config_name",
default="",
type=str,
help="Pretrained config name or path if not the same as model_name")
parser.add_argument(
"--tokenizer_name",
default="",
type=str,
help="Pretrained tokenizer name or path if not the same as model_name")
parser.add_argument(
"--cache_dir",
default="",
type=str,
help=
"Where do you want to store the pre-trained models downloaded from s3")
parser.add_argument(
"--max_seq_length",
default=512,
type=int,
help=
"The maximum total input sequence length after tokenization. Sequences longer "
"than this will be truncated, sequences shorter will be padded.")
parser.add_argument("--do_train",
action='store_true',
help="Whether to run training.")
parser.add_argument("--do_eval",
action='store_true',
help="Whether to run eval on the dev set.")
parser.add_argument(
"--do_predict",
action='store_true',
help="Whether to run the model in inference mode on the test set.")
parser.add_argument(
"--do_lower_case",
action='store_true',
help="Set this flag if you are using an uncased model.")
parser.add_argument("--per_gpu_train_batch_size",
default=8,
type=int,
help="Batch size per GPU/CPU for training.")
parser.add_argument("--per_gpu_eval_batch_size",
default=8,
type=int,
help="Batch size per GPU/CPU for evaluation.")
parser.add_argument(
'--gradient_accumulation_steps',
type=int,
default=1,
help=
"Number of updates steps to accumulate before performing a backward/update pass."
)
parser.add_argument("--learning_rate",
default=5e-5,
type=float,
help="The initial learning rate for Adam.")
parser.add_argument("--weight_decay",
default=0.0,
type=float,
help="Weight deay if we apply some.")
parser.add_argument("--adam_epsilon",
default=1e-6,
type=float,
help="Epsilon for Adam optimizer.")
parser.add_argument("--max_grad_norm",
default=1.0,
type=float,
help="Max gradient norm.")
parser.add_argument("--num_train_epochs",
default=3.0,
type=float,
help="Total number of training epochs to perform.")
parser.add_argument(
"--max_steps",
default=-1,
type=int,
help=
"If > 0: set total number of training steps to perform. Override num_train_epochs."
)
parser.add_argument(
"--warmup_proportion",
default=0.1,
type=float,
help=
"Proportion of training to perform linear learning rate warmup for,E.g., 0.1 = 10% of training."
)
parser.add_argument('--logging_steps',
type=int,
default=10,
help="Log every X updates steps.")
parser.add_argument('--save_steps',
type=int,
default=1000,
help="Save checkpoint every X updates steps.")
parser.add_argument(
"--eval_all_checkpoints",
action='store_true',
help=
"Evaluate all checkpoints starting with the same prefix as model_name ending and ending with step number"
)
parser.add_argument("--no_cuda",
action='store_true',
help="Avoid using CUDA when available")
parser.add_argument('--overwrite_output_dir',
action='store_true',
help="Overwrite the content of the output directory")
parser.add_argument(
'--overwrite_cache',
action='store_true',
help="Overwrite the cached training and evaluation sets")
parser.add_argument('--seed',
type=int,
default=42,
help="random seed for initialization")
parser.add_argument(
'--fp16',
action='store_true',
help=
"Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit"
)
parser.add_argument(
'--fp16_opt_level',
type=str,
default='O1',
help=
"For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']."
"See details at https://nvidia.github.io/apex/amp.html")
parser.add_argument("--local_rank",
type=int,
default=-1,
help="For distributed training: local_rank")
parser.add_argument('--server_ip',
type=str,
default='',
help="For distant debugging.")
parser.add_argument('--server_port',
type=str,
default='',
help="For distant debugging.")
os.environ['CUDA_VISIBLE_DEVICES'] = "0"
args = parser.parse_args()
if not os.path.exists(args.output_dir):
os.mkdir(args.output_dir)
args.output_dir = args.output_dir + '{}'.format(args.model_type)
if not os.path.exists(args.output_dir):
os.mkdir(args.output_dir)
init_logger(log_file=args.output_dir +
'/{}-{}.log'.format(args.model_type, args.task_name))
if os.path.exists(args.output_dir) and os.listdir(
args.output_dir
) and args.do_train and not args.overwrite_output_dir:
raise ValueError(
"Output directory ({}) already exists and is not empty. Use --overwrite_output_dir to overcome."
.format(args.output_dir))
# Setup distant debugging if needed
if args.server_ip and args.server_port:
# Distant debugging - see https://code.visualstudio.com/docs/python/debugging#_attach-to-a-local-script
import ptvsd
print("Waiting for debugger attach")
ptvsd.enable_attach(address=(args.server_ip, args.server_port),
redirect_output=True)
ptvsd.wait_for_attach()
# Setup CUDA, GPU & distributed training
if args.local_rank == -1 or args.no_cuda:
device = torch.device("cuda" if torch.cuda.is_available()
and not args.no_cuda else "cpu")
args.n_gpu = torch.cuda.device_count()
else: # Initializes the distributed backend which will take care of sychronizing nodes/GPUs
torch.cuda.set_device(args.local_rank)
device = torch.device("cuda", args.local_rank)
torch.distributed.init_process_group(backend='nccl')
args.n_gpu = 1
args.device = device
# Setup logging
logger.warning(
"Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s",
args.local_rank, device, args.n_gpu, bool(args.local_rank != -1),
args.fp16)
# Set seed
seed_everything(args.seed)
# Prepare GLUE task
args.task_name = args.task_name.lower()
if args.task_name not in processors:
raise ValueError("Task not found: %s" % (args.task_name))
processor = processors[args.task_name]()
args.output_mode = output_modes[args.task_name]
label_list = processor.get_labels()
num_labels = len(label_list)
# Load pretrained model and tokenizer
if args.local_rank not in [-1, 0]:
torch.distributed.barrier(
) # Make sure only the first process in distributed training will download model & vocab
args.model_type = args.model_type.lower()
config = AlbertConfig.from_pretrained(
args.config_name if args.config_name else args.model_name_or_path,
num_labels=num_labels,
finetuning_task=args.task_name)
tokenizer = tokenization_albert.FullTokenizer(
vocab_file=args.vocab_file,
do_lower_case=args.do_lower_case,
spm_model_file=args.spm_model_file)
model = AlbertForSentenceRanking.from_pretrained(
args.model_name_or_path,
from_tf=bool('.ckpt' in args.model_name_or_path),
config=config)
if args.local_rank == 0:
torch.distributed.barrier(
) # Make sure only the first process in distributed training will download model & vocab
model.to(args.device)
logger.info("Training/evaluation parameters %s", args)
# Training
if args.do_train:
train_dataset = load_and_cache_examples(args,
args.task_name,
tokenizer,
data_type='train')
global_step, tr_loss = train(args, train_dataset, model, tokenizer)
logger.info(" global_step = %s, average loss = %s", global_step,
tr_loss)
# Saving best-practices: if you use defaults names for the model, you can reload it using from_pretrained()
if args.do_train and (args.local_rank == -1
or torch.distributed.get_rank() == 0):
# Create output directory if needed
if not os.path.exists(args.output_dir) and args.local_rank in [-1, 0]:
os.makedirs(args.output_dir)
logger.info("Saving model checkpoint to %s", args.output_dir)
# Save a trained model, configuration and tokenizer using `save_pretrained()`.
# They can then be reloaded using `from_pretrained()`
model_to_save = model.module if hasattr(
model,
'module') else model # Take care of distributed/parallel training
model_to_save.save_pretrained(args.output_dir)
# Good practice: save your training arguments together with the trained model
torch.save(args, os.path.join(args.output_dir, 'training_args.bin'))
# Evaluation
results = []
if args.do_eval and args.local_rank in [-1, 0]:
tokenizer = tokenization_albert.FullTokenizer(
vocab_file=args.vocab_file,
do_lower_case=args.do_lower_case,
spm_model_file=args.spm_model_file)
checkpoints = [(0, args.output_dir)]
if args.eval_all_checkpoints:
checkpoints = list(
os.path.dirname(c) for c in sorted(
glob.glob(args.output_dir + '/**/' + WEIGHTS_NAME,
recursive=True)))
checkpoints = [(int(checkpoint.split('-')[-1]), checkpoint)
for checkpoint in checkpoints
if checkpoint.find('checkpoint') != -1]
checkpoints = sorted(checkpoints, key=lambda x: x[0])
logger.info("Evaluate the following checkpoints: %s", checkpoints)
for _, checkpoint in checkpoints:
global_step = checkpoint.split(
'-')[-1] if len(checkpoints) > 1 else ""
prefix = checkpoint.split(
'/')[-1] if checkpoint.find('checkpoint') != -1 else ""
model = AlbertForSentenceRanking.from_pretrained(checkpoint)
model.to(args.device)
result = evaluate(args, model, tokenizer, prefix=prefix)
results.extend([(k + '_{}'.format(global_step), v)
for k, v in result.items()])
output_eval_file = os.path.join(args.output_dir,
"checkpoint_eval_results.txt")
with open(output_eval_file, "w") as writer:
for key, value in results:
writer.write("%s = %s\n" % (key, str(value)))
def train(args, train_dataset, model, tokenizer):
""" Train the model """
args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)
train_sampler = SequentialSampler(
train_dataset) if args.local_rank == -1 else DistributedSampler(
train_dataset)
train_dataloader = DataLoader(train_dataset,
sampler=train_sampler,
batch_size=args.train_batch_size,
collate_fn=collate_fn)
if args.max_steps > 0:
num_training_steps = args.max_steps
args.num_train_epochs = args.max_steps // (
len(train_dataloader) // args.gradient_accumulation_steps) + 1
else:
num_training_steps = len(
train_dataloader
) // args.gradient_accumulation_steps * args.num_train_epochs
args.warmup_steps = int(num_training_steps * args.warmup_proportion)
# Prepare optimizer and schedule (linear warmup and decay)
no_decay = ['bias', 'LayerNorm.weight']
optimizer_grouped_parameters = [{
'params': [
p for n, p in model.named_parameters()
if not any(nd in n for nd in no_decay)
],
'weight_decay':
args.weight_decay
}, {
'params': [
p for n, p in model.named_parameters()
if any(nd in n for nd in no_decay)
],
'weight_decay':
0.0
}]
# optimizer = Lamb(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
optimizer = AdamW(params=optimizer_grouped_parameters,
lr=args.learning_rate,
eps=args.adam_epsilon)
scheduler = get_linear_schedule_with_warmup(
optimizer,
num_warmup_steps=args.warmup_steps,
num_training_steps=num_training_steps)
if args.fp16:
try:
from apex import amp
except ImportError:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to use fp16 training."
)
model, optimizer = amp.initialize(model,
optimizer,
opt_level=args.fp16_opt_level)
# multi-gpu training (should be after apex fp16 initialization)
if args.n_gpu > 1:
model = torch.nn.DataParallel(model)
# Distributed training (should be after apex fp16 initialization)
if args.local_rank != -1:
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[args.local_rank],
output_device=args.local_rank,
find_unused_parameters=True)
# Train!
logger.info("***** Running training *****")
logger.info(" Num examples = %d", len(train_dataset))
logger.info(" Num Epochs = %d", args.num_train_epochs)
logger.info(" Instantaneous batch size per GPU = %d",
args.per_gpu_train_batch_size)
logger.info(
" Total train batch size (w. parallel, distributed & accumulation) = %d",
args.train_batch_size * args.gradient_accumulation_steps *
(torch.distributed.get_world_size() if args.local_rank != -1 else 1))
logger.info(" Gradient Accumulation steps = %d",
args.gradient_accumulation_steps)
logger.info(" Total optimization steps = %d", num_training_steps)
global_step = 0
tr_loss, logging_loss = 0.0, 0.0
model.zero_grad()
seed_everything(
args.seed
) # Added here for reproductibility (even between python 2 and 3)
for _ in range(int(args.num_train_epochs)):
pbar = ProgressBar(n_total=len(train_dataloader), desc='Training')
for step, batch in enumerate(train_dataloader):
model.train()
batch = tuple(t.to(args.device) for t in batch)
inputs = {
'input_ids': batch[0],
'attention_mask': batch[1],
'labels': batch[3]
}
inputs['token_type_ids'] = batch[2]
outputs = model(**inputs)
loss = outputs[
0] # model outputs are always tuple in transformers (see doc)
if args.n_gpu > 1:
loss = loss.mean(
) # mean() to average on multi-gpu parallel training
if args.gradient_accumulation_steps > 1:
loss = loss / args.gradient_accumulation_steps
if args.fp16:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer),
args.max_grad_norm)
else:
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(),
args.max_grad_norm)
tr_loss += loss.item()
if (step + 1) % args.gradient_accumulation_steps == 0:
optimizer.step()
scheduler.step() # Update learning rate schedule
model.zero_grad()
global_step += 1
if args.local_rank in [
-1, 0
] and args.logging_steps > 0 and global_step % args.logging_steps == 0:
#Log metrics
if args.local_rank == -1: # Only evaluate when single GPU otherwise metrics may not average well
evaluate(args, model, tokenizer)
if args.local_rank in [
-1, 0
] and args.save_steps > 0 and global_step % args.save_steps == 0:
# Save model checkpoint
output_dir = os.path.join(args.output_dir,
'checkpoint-{}'.format(global_step))
if not os.path.exists(output_dir):
os.makedirs(output_dir)
model_to_save = model.module if hasattr(
model, 'module'
) else model # Take care of distributed/parallel training
model_to_save.save_pretrained(output_dir)
torch.save(args, os.path.join(output_dir, 'training_args.bin'))
logger.info("Saving model checkpoint to %s", output_dir)
pbar(step, {'loss': loss.item()})
print(" ")
if 'cuda' in str(args.device):
torch.cuda.empty_cache()
return global_step, tr_loss / global_step
def _run_epoch(self):
"""
Run each epoch.
The training steps:
- Get batch and feed them into model
- Get outputs. Caculate all losses and sum them up
- Loss backwards and optimizer steps
- Evaluation
- Update and output result
"""
num_batch = len(self._trainloader)
train_loss = AverageMeter()
with tqdm(enumerate(self._trainloader),
total=num_batch,
disable=not self._verbose) as pbar:
for step, (inputs, target) in pbar:
# Run Train
outputs = self._model(inputs)
## 计算所有的loss并相加
loss = sum([c(outputs, target) for c in self._criterions])
self._backward(loss)
## 更新loss的值
train_loss.update(loss.item())
## 设置Progress Bar
pbar.set_description(f"Epoch {self._epoch}/{self._epochs}")
pbar.set_postfix(loss=f"{loss.item():.3f}")
self._writer.add_scalar("Loss/train", loss.item(),
self._iteration)
# Run Evaluate
self._iteration += 1
if self._iteration % self._validate_interval == 0:
pbar.update(1)
if self._verbose:
pbar.write(f"[Iter-{self._iteration} "
f"Loss-{train_loss.avg:.3f}]")
## 更新validloader评估的结果
result, eval_loss = self.evaluate(self._validloader)
m_string = ""
for metric in result:
res = result[metric]
if not isinstance(metric, str):
metric = metric.ALIAS[0]
self._writer.add_scalar(f"{metric}/eval", res,
self._iteration)
m_string += f"| {metric}: {res} "
logger.info(
f"Epoch: {self._epoch} | Train Loss: {loss.item():.3f} | "
f"Eval loss: {eval_loss.item(): .3f} " + m_string)
if self._verbose:
pbar.write(" Validation: " +
'-'.join(f"{k}: {round(v, 4)}"
for k, v in result.items()))
## Early Stopping
self._early_stopping.update(result)
if self._early_stopping.should_stop_early:
self._save()
pbar.write("Ran out of patience. Stop training...")
break
elif self._early_stopping.is_best_so_far:
self._save()