""" This script provides an example to wrap TencentPretrain for classification. """ import sys import os import random import argparse import torch import torch.nn as nn tencentpretrain_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) sys.path.append(tencentpretrain_dir) from tencentpretrain.embeddings import * from tencentpretrain.encoders import * from tencentpretrain.utils.vocab import Vocab from tencentpretrain.utils.constants import * from tencentpretrain.utils import * from tencentpretrain.utils.optimizers import * from tencentpretrain.utils.config import load_hyperparam from tencentpretrain.utils.seed import set_seed from tencentpretrain.utils.logging import init_logger from tencentpretrain.utils.misc import pooling from tencentpretrain.model_saver import save_model from tencentpretrain.opts import finetune_opts, tokenizer_opts, adv_opts class Classifier(nn.Module): def __init__(self, args): super(Classifier, self).__init__() self.embedding = Embedding(args) for embedding_name in args.embedding: tmp_emb = str2embedding[embedding_name](args, len(args.tokenizer.vocab)) self.embedding.update(tmp_emb, embedding_name) self.encoder = str2encoder[args.encoder](args) self.labels_num = args.labels_num self.pooling_type = args.pooling self.soft_targets = args.soft_targets self.soft_alpha = args.soft_alpha self.output_layer_1 = nn.Linear(args.hidden_size, args.hidden_size) self.output_layer_2 = nn.Linear(args.hidden_size, self.labels_num) def forward(self, src, tgt, seg, soft_tgt=None): """ Args: src: [batch_size x seq_length] tgt: [batch_size] seg: [batch_size x seq_length] """ # Embedding. emb = self.embedding(src, seg) # Encoder. output = self.encoder(emb, seg) # Target. output = pooling(output, seg, self.pooling_type) output = torch.tanh(self.output_layer_1(output)) logits = self.output_layer_2(output) if tgt is not None: if self.soft_targets and soft_tgt is not None: loss = self.soft_alpha * nn.MSELoss()(logits, soft_tgt) + \ (1 - self.soft_alpha) * nn.NLLLoss()(nn.LogSoftmax(dim=-1)(logits), tgt.view(-1)) else: loss = nn.NLLLoss()(nn.LogSoftmax(dim=-1)(logits), tgt.view(-1)) return loss, logits else: return None, logits def count_labels_num(path): labels_set, columns = set(), {} with open(path, mode="r", encoding="utf-8") as f: for line_id, line in enumerate(f): if line_id == 0: for i, column_name in enumerate(line.rstrip("\r\n").split("\t")): columns[column_name] = i continue line = line.rstrip("\r\n").split("\t") label = int(line[columns["label"]]) labels_set.add(label) return len(labels_set) def load_or_initialize_parameters(args, model): if args.pretrained_model_path is not None: # Initialize with pretrained model. model.load_state_dict(torch.load(args.pretrained_model_path, map_location="cpu"), strict=False) else: # Initialize with normal distribution. for n, p in list(model.named_parameters()): if "gamma" not in n and "beta" not in n: p.data.normal_(0, 0.02) def build_optimizer(args, model): param_optimizer = list(model.named_parameters()) no_decay = ["bias", "gamma", "beta"] 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)], "weight_decay": 0.0}, ] if args.optimizer in ["adamw"]: optimizer = str2optimizer[args.optimizer](optimizer_grouped_parameters, lr=args.learning_rate, correct_bias=False) else: optimizer = str2optimizer[args.optimizer](optimizer_grouped_parameters, lr=args.learning_rate, scale_parameter=False, relative_step=False) if args.scheduler in ["constant"]: scheduler = str2scheduler[args.scheduler](optimizer) elif args.scheduler in ["constant_with_warmup"]: scheduler = str2scheduler[args.scheduler](optimizer, args.train_steps*args.warmup) else: scheduler = str2scheduler[args.scheduler](optimizer, args.train_steps*args.warmup, args.train_steps) return optimizer, scheduler def batch_loader(batch_size, src, tgt, seg, soft_tgt=None): instances_num = src.size()[0] for i in range(instances_num // batch_size): src_batch = src[i * batch_size : (i + 1) * batch_size, :] tgt_batch = tgt[i * batch_size : (i + 1) * batch_size] seg_batch = seg[i * batch_size : (i + 1) * batch_size, :] if soft_tgt is not None: soft_tgt_batch = soft_tgt[i * batch_size : (i + 1) * batch_size, :] yield src_batch, tgt_batch, seg_batch, soft_tgt_batch else: yield src_batch, tgt_batch, seg_batch, None if instances_num > instances_num // batch_size * batch_size: src_batch = src[instances_num // batch_size * batch_size :, :] tgt_batch = tgt[instances_num // batch_size * batch_size :] seg_batch = seg[instances_num // batch_size * batch_size :, :] if soft_tgt is not None: soft_tgt_batch = soft_tgt[instances_num // batch_size * batch_size :, :] yield src_batch, tgt_batch, seg_batch, soft_tgt_batch else: yield src_batch, tgt_batch, seg_batch, None def read_dataset(args, path): dataset, columns = [], {} with open(path, mode="r", encoding="utf-8") as f: for line_id, line in enumerate(f): if line_id == 0: for i, column_name in enumerate(line.rstrip("\r\n").split("\t")): columns[column_name] = i continue line = line.rstrip("\r\n").split("\t") tgt = int(line[columns["label"]]) if args.soft_targets and "logits" in columns.keys(): soft_tgt = [float(value) for value in line[columns["logits"]].split(" ")] if "text_b" not in columns: # Sentence classification. text_a = line[columns["text_a"]] src = args.tokenizer.convert_tokens_to_ids([CLS_TOKEN] + args.tokenizer.tokenize(text_a) + [SEP_TOKEN]) seg = [1] * len(src) else: # Sentence-pair classification. text_a, text_b = line[columns["text_a"]], line[columns["text_b"]] src_a = args.tokenizer.convert_tokens_to_ids([CLS_TOKEN] + args.tokenizer.tokenize(text_a) + [SEP_TOKEN]) src_b = args.tokenizer.convert_tokens_to_ids(args.tokenizer.tokenize(text_b) + [SEP_TOKEN]) src = src_a + src_b seg = [1] * len(src_a) + [2] * len(src_b) if len(src) > args.seq_length: src = src[: args.seq_length] seg = seg[: args.seq_length] PAD_ID = args.tokenizer.convert_tokens_to_ids([PAD_TOKEN])[0] while len(src) < args.seq_length: src.append(PAD_ID) seg.append(0) if args.soft_targets and "logits" in columns.keys(): dataset.append((src, tgt, seg, soft_tgt)) else: dataset.append((src, tgt, seg)) return dataset def train_model(args, model, optimizer, scheduler, src_batch, tgt_batch, seg_batch, soft_tgt_batch=None): model.zero_grad() src_batch = src_batch.to(args.device) tgt_batch = tgt_batch.to(args.device) seg_batch = seg_batch.to(args.device) if soft_tgt_batch is not None: soft_tgt_batch = soft_tgt_batch.to(args.device) loss, _ = model(src_batch, tgt_batch, seg_batch, soft_tgt_batch) if torch.cuda.device_count() > 1: loss = torch.mean(loss) if args.fp16: with args.amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward() else: loss.backward() if args.use_adv and args.adv_type == "fgm": args.adv_method.attack(epsilon=args.fgm_epsilon) loss_adv, _ = model(src_batch, tgt_batch, seg_batch, soft_tgt_batch) if torch.cuda.device_count() > 1: loss_adv = torch.mean(loss_adv) loss_adv.backward() args.adv_method.restore() if args.use_adv and args.adv_type == "pgd": K = args.pgd_k args.adv_method.backup_grad() for t in range(K): # apply the perturbation to embedding args.adv_method.attack(epsilon=args.pgd_epsilon, alpha=args.pgd_alpha, is_first_attack=(t == 0)) if t != K - 1: model.zero_grad() else: args.adv_method.restore_grad() loss_adv, _ = model(src_batch, tgt_batch, seg_batch, soft_tgt_batch) if torch.cuda.device_count() > 1: loss_adv = torch.mean(loss_adv) loss_adv.backward() args.adv_method.restore() optimizer.step() scheduler.step() return loss def evaluate(args, dataset): src = torch.LongTensor([sample[0] for sample in dataset]) tgt = torch.LongTensor([sample[1] for sample in dataset]) seg = torch.LongTensor([sample[2] for sample in dataset]) batch_size = args.batch_size correct = 0 # Confusion matrix. confusion = torch.zeros(args.labels_num, args.labels_num, dtype=torch.long) args.model.eval() for i, (src_batch, tgt_batch, seg_batch, _) in enumerate(batch_loader(batch_size, src, tgt, seg)): src_batch = src_batch.to(args.device) tgt_batch = tgt_batch.to(args.device) seg_batch = seg_batch.to(args.device) with torch.no_grad(): _, logits = args.model(src_batch, tgt_batch, seg_batch) pred = torch.argmax(nn.Softmax(dim=1)(logits), dim=1) gold = tgt_batch for j in range(pred.size()[0]): confusion[pred[j], gold[j]] += 1 correct += torch.sum(pred == gold).item() args.logger.info("Confusion matrix:") args.logger.info(confusion) args.logger.info("Report precision, recall, and f1:") eps = 1e-9 for i in range(confusion.size()[0]): p = confusion[i, i].item() / (confusion[i, :].sum().item() + eps) r = confusion[i, i].item() / (confusion[:, i].sum().item() + eps) f1 = 2 * p * r / (p + r + eps) args.logger.info("Label {}: {:.3f}, {:.3f}, {:.3f}".format(i, p, r, f1)) args.logger.info("Acc. (Correct/Total): {:.4f} ({}/{}) ".format(correct / len(dataset), correct, len(dataset))) return correct / len(dataset), confusion def main(): parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) finetune_opts(parser) tokenizer_opts(parser) parser.add_argument("--soft_targets", action='store_true', help="Train model with logits.") parser.add_argument("--soft_alpha", type=float, default=0.5, help="Weight of the soft targets loss.") adv_opts(parser) args = parser.parse_args() # Load the hyperparameters from the config file. args = load_hyperparam(args) # Count the number of labels. args.labels_num = count_labels_num(args.train_path) # Build tokenizer. args.tokenizer = str2tokenizer[args.tokenizer](args) set_seed(args.seed) # Build classification model. model = Classifier(args) # Load or initialize parameters. load_or_initialize_parameters(args, model) # Get logger. args.logger = init_logger(args) args.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = model.to(args.device) # Training phase. trainset = read_dataset(args, args.train_path) instances_num = len(trainset) batch_size = args.batch_size args.train_steps = int(instances_num * args.epochs_num / batch_size) + 1 args.logger.info("Batch size: {}".format(batch_size)) args.logger.info("The number of training instances: {}".format(instances_num)) optimizer, scheduler = build_optimizer(args, model) 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) args.amp = amp if torch.cuda.device_count() > 1: args.logger.info("{} GPUs are available. Let's use them.".format(torch.cuda.device_count())) model = torch.nn.DataParallel(model) args.model = model if args.use_adv: args.adv_method = str2adv[args.adv_type](model) total_loss, result, best_result = 0.0, 0.0, 0.0 args.logger.info("Start training.") for epoch in range(1, args.epochs_num + 1): random.shuffle(trainset) src = torch.LongTensor([example[0] for example in trainset]) tgt = torch.LongTensor([example[1] for example in trainset]) seg = torch.LongTensor([example[2] for example in trainset]) if args.soft_targets: soft_tgt = torch.FloatTensor([example[3] for example in trainset]) else: soft_tgt = None model.train() for i, (src_batch, tgt_batch, seg_batch, soft_tgt_batch) in enumerate(batch_loader(batch_size, src, tgt, seg, soft_tgt)): loss = train_model(args, model, optimizer, scheduler, src_batch, tgt_batch, seg_batch, soft_tgt_batch) total_loss += loss.item() if (i + 1) % args.report_steps == 0: args.logger.info("Epoch id: {}, Training steps: {}, Avg loss: {:.3f}".format(epoch, i + 1, total_loss / args.report_steps)) total_loss = 0.0 result = evaluate(args, read_dataset(args, args.dev_path)) if result[0] > best_result: best_result = result[0] save_model(model, args.output_model_path) # Evaluation phase. if args.test_path is not None: args.logger.info("Test set evaluation.") if torch.cuda.device_count() > 1: args.model.module.load_state_dict(torch.load(args.output_model_path)) else: args.model.load_state_dict(torch.load(args.output_model_path)) evaluate(args, read_dataset(args, args.test_path)) if __name__ == "__main__": main()