"""
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()