def test_pbar_log_message(capsys):
ProgressBar.log_message("test")
captured = capsys.readouterr()
out = captured.out.split('\r')
out = list(map(lambda x: x.strip(), out))
out = list(filter(None, out))
expected = u'test'
assert out[-1] == expected
def test_pbar_log_message_file(tmp_path):
file_path = tmp_path / "temp.txt"
file = open(str(file_path), "w+")
pbar = ProgressBar(file=file)
pbar.log_message("test")
file.close() # Force a flush of the buffer. file.flush() does not work.
file = open(str(file_path), "r")
lines = file.readlines()
expected = "test\n"
assert lines[0] == expected
def add_logging_and_checkpoint_saving(trainer, evaluator, metrics, model, optimizer, args, prefix=""):
""" Add to training engine tensorboard logging, progress bar with average loss, checkpoint saving and save training config. """
# Add progress bar with average loss
RunningAverage(output_transform=lambda x: x).attach(trainer, prefix + "loss")
pbar = ProgressBar(persist=True)
pbar.attach(trainer, metric_names=[prefix + "loss"])
evaluator.add_event_handler(Events.COMPLETED,
lambda _: pbar.log_message("Validation: %s" % pformat(evaluator.state.metrics)))
# Add tensorboard logging with training and evaluation metrics
tb_logger = TensorboardLogger(log_dir=None)
tb_logger.attach(trainer, log_handler=OutputHandler(tag="training", metric_names=[prefix + "loss"]),
event_name=Events.ITERATION_COMPLETED)
tb_logger.attach(trainer, log_handler=OptimizerParamsHandler(optimizer),
event_name=Events.ITERATION_STARTED)
@evaluator.on(Events.COMPLETED)
def tb_log_metrics(engine):
for name in metrics.keys():
tb_logger.writer.add_scalar(name, engine.state.metrics[name], trainer.state.iteration)
# Add checkpoint saving after each epoch - take care of distributed encapsulation ('getattr()')
checkpoint_handler = ModelCheckpoint(tb_logger.writer.log_dir, 'checkpoint', save_interval=1, n_saved=3)
trainer.add_event_handler(Events.EPOCH_COMPLETED, checkpoint_handler, {'mymodel': getattr(model, 'module', model)})
# Save training configuration
torch.save(args, os.path.join(tb_logger.writer.log_dir, CONFIG_NAME))
return checkpoint_handler, tb_logger
def log_iter(engine: Engine, trainer: Engine, pbar: ProgressBar, title: str,
log_interval: int) -> None:
epoch = trainer.state.epoch
iteration = engine.state.iteration
metrics = engine.state.metrics
stats = {k: '%.3f' % v for k, v in metrics.items() if 'loss' in k}
if hasattr(engine.state, 'lr'):
stats['lr'] = ', '.join(['%.1e' % val for val in engine.state.lr])
stats = ', '.join(['{}: {}'.format(*e) for e in stats.items()])
t0 = engine.state.t0
t1 = time.time()
it_time = (t1 - t0) / log_interval
cur_time = humanize_time(t1)
pbar.log_message("[{}][{:.2f} s] {:>5} | ep: {:2d}, it: {:3d}, {}".format(
cur_time, it_time, title, epoch, iteration, stats))
engine.state.t0 = t1
def create_evaluator(args, model):
def inference(engine, batch):
model.eval()
with torch.no_grad():
batch = {
name: input_tensor.to(args.device)
for name, input_tensor in batch.items()
}
lm_logits, mc_logits, *_ = model(
input_ids=batch["input_ids"],
token_type_ids=batch["token_type_ids"],
mc_token_ids=batch["mc_token_ids"],
lm_labels=batch["lm_labels"],
mc_labels=batch["mc_labels"],
persona=batch["persona"],
history=batch["history"],
effects=batch["effects"],
)
lm_logits_flat_shifted = lm_logits[..., :-1, :].contiguous().view(
-1, lm_logits.size(-1))
lm_labels_flat_shifted = batch["lm_labels"][:, 0, :,
1:].contiguous().view(
-1)
return (lm_logits_flat_shifted,
mc_logits), (lm_labels_flat_shifted, batch["mc_labels"])
evaluator = Engine(inference)
metrics = {
"nll":
Loss(torch.nn.CrossEntropyLoss(ignore_index=-100),
output_transform=lambda x: (x[0][0], x[1][0])),
# the accuracy is a filler since multiple-choice is not used.
"accuracy":
Accuracy(
#output_transform=lambda x: (torch.argmax(x[0][1].view((-1,)), dim=0, keepdim=True), x[1][1][:, 0])),
output_transform=lambda x:
(torch.argmax(x[0][1].squeeze(1), dim=-1), x[1][1][:, 0])),
"ppl":
Perplexity(output_transform=lambda x: (x[0][0], None)),
}
for name, metric in metrics.items():
metric.attach(evaluator, name)
if args.local_rank in [-1, 0]:
pbar = ProgressBar(persist=True)
pbar.attach(evaluator)
evaluator.add_event_handler(
Events.COMPLETED, lambda _: pbar.log_message(
"Validation: %s" % pformat(evaluator.state.metrics)))
return evaluator
def train_model(self, n_epochs, train_loader, val_loader, eval_before_start=True):
# Attach evaluation to trainer: we evaluate when we start the training and at the end of each epoch
self.trainer.add_event_handler(Events.EPOCH_COMPLETED, lambda _: self.evaluator.run(val_loader))
self.trainer.add_event_handler(Events.EPOCH_COMPLETED, lambda _: self.update_epoch())
if eval_before_start:
self.trainer.add_event_handler(Events.STARTED, lambda _: self.evaluator.run(val_loader))
# Linearly decrease the learning rate from lr to zero
scheduler = PiecewiseLinear(self.optimizer, "lr", [(0, self.lr), (n_epochs * len(train_loader), 0.0)])
self.trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
# Prepare metrics
RunningAverage(output_transform=lambda x: x).attach(self.trainer, "loss")
metrics = {"nll": Loss(torch.nn.CrossEntropyLoss(ignore_index=-1), output_transform=lambda x: (x[0][0], x[1][0])),
"accuracy": Accuracy(output_transform=lambda x: (x[0][1], x[1][1]))}
metrics["average_ppl"] = MetricsLambda(math.exp, metrics["nll"])
for name, metric in metrics.items():
metric.attach(self.evaluator, name)
# On the main process: add progress bar, tensorboard, checkpoints and save model
pbar = ProgressBar(persist=True)
pbar.attach(self.trainer, metric_names=["loss"])
if not self.verbose:
pbar_eval = ProgressBar(persist=False)
pbar_eval.attach(self.evaluator)
self.evaluator.add_event_handler(Events.STARTED, lambda _: self.logger.info(f'Beginning validation for epoch {self.epoch}...'))
self.evaluator.add_event_handler(Events.COMPLETED, lambda _: pbar.log_message("Validation: %s" % pformat(self.evaluator.state.metrics)))
self.tb_logger.attach(self.trainer, log_handler=OutputHandler(tag="training", metric_names=["loss"]), event_name=Events.ITERATION_COMPLETED)
self.tb_logger.attach(self.trainer, log_handler=OptimizerParamsHandler(self.optimizer), event_name=Events.ITERATION_STARTED)
self.tb_logger.attach(self.evaluator, log_handler=OutputHandler(tag="validation", metric_names=list(metrics.keys()), another_engine=self.trainer),
event_name=Events.EPOCH_COMPLETED)
self.trainer.add_event_handler(Events.EPOCH_COMPLETED, self.checkpoint_handler,
{'mymodel': getattr(self.model, 'module', self.model)}) # "getattr" takes care of distributed encapsulation
# Run the training
self.trainer.run(train_loader, max_epochs=n_epochs)
# On the main process: close tensorboard logger and rename the last checkpoint (for easy re-loading with OpenAIGPTModel.from_pretrained method)
if n_epochs > 0:
os.rename(self.checkpoint_handler._saved[-1][1][-1], os.path.join(cfg.checkpoint_log_folder, self.name, WEIGHTS_NAME))
self.tb_logger.close()
def train():
parser = ArgumentParser()
parser.add_argument(
"--dataset_path",
type=str,
default="",
help="Path or url of the dataset. If empty download from S3.")
parser.add_argument("--dataset_cache",
type=str,
default='./dataset_cache',
help="Path or url of the dataset cache")
parser.add_argument("--model_checkpoint",
type=str,
default="multi-bert",
help="Path, url or short name of the model")
parser.add_argument("--num_candidates",
type=int,
default=2,
help="Number of candidates for training")
parser.add_argument("--max_turns",
type=int,
default=3,
help="Number of previous turns to keep in history")
parser.add_argument("--train_batch_size",
type=int,
default=4,
help="Batch size for training")
parser.add_argument("--valid_batch_size",
type=int,
default=4,
help="Batch size for validation")
parser.add_argument("--gradient_accumulation_steps",
type=int,
default=8,
help="Accumulate gradients on several steps")
parser.add_argument("--lr",
type=float,
default=6.25e-5,
help="Learning rate")
parser.add_argument("--lm_coef",
type=float,
default=1.0,
help="LM loss coefficient")
parser.add_argument("--max_norm",
type=float,
default=1.0,
help="Clipping gradient norm")
parser.add_argument("--n_epochs",
type=int,
default=3,
help="Number of training epochs")
parser.add_argument("--personality_permutations",
type=int,
default=1,
help="Number of permutations of personality sentences")
parser.add_argument(
"--eval_before_start",
action='store_true',
help="If true start with a first evaluation before training")
parser.add_argument("--device",
type=str,
default="cuda" if torch.cuda.is_available() else "cpu",
help="Device (cuda or cpu)")
parser.add_argument(
"--fp16",
type=str,
default="",
help=
"Set to O0, O1, O2 or O3 for fp16 training (see apex documentation)")
parser.add_argument(
"--local_rank",
type=int,
default=-1,
help="Local rank for distributed training (-1: not distributed)")
parser.add_argument("--random_init",
action='store_true',
help="If true random initailze the model")
parser.add_argument(
"--train_lang",
type=str,
default="",
help="train monolingual model, defaul: multilingual model")
args = parser.parse_args()
# logging is set to INFO (resp. WARN) for main (resp. auxiliary) process. logger.info => log main process only, logger.warning => log all processes
logging.basicConfig(
level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN)
logger.warning(
"Running process %d", args.local_rank
) # This is a logger.warning: it will be printed by all distributed processes
logger.info("Arguments: %s", pformat(args))
# Initialize distributed training if needed
args.distributed = (args.local_rank != -1)
if args.distributed:
torch.cuda.set_device(args.local_rank)
args.device = torch.device("cuda", args.local_rank)
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
# Model
logger.info("Prepare tokenizer, pretrained model and optimizer.")
model_path = 'bert-base-multilingual-cased'
if args.train_lang in ["En", "It", "Jp",
"Zh"]: # for Fr Ko Id we use MBERT
model_path = LANG_2_MODEL[args.train_lang]
tokenizer = BertTokenizer.from_pretrained(model_path)
if args.train_lang == "Jp":
tokenizer = BertJapaneseTokenizer.from_pretrained(model_path)
model = Model2Model.from_pretrained(model_path)
# tokenizer = BertTokenizer.from_pretrained('bert-base-multilingual-cased')
# if args.random_init:
# config = BertConfig.from_pretrained('bert-base-multilingual-cased')
# config.is_decoder = True
# bert_decoder = BertForMaskedLM(config)
# model = Model2Model.from_pretrained('bert-base-multilingual-cased', decoder_model=bert_decoder)
# else:
# model = Model2Model.from_pretrained('bert-base-multilingual-cased')
# model_dict = model.state_dict()
# # initialize crossattention with selfattention
# model_dict.update({ name: model_dict[name.replace("crossattention", "attention")] for name in model_dict if "crossattention" in name })
# model.load_state_dict(model_dict)
model.to(args.device)
# Add special tokens if they are not already added
add_special_tokens_(model, tokenizer)
optimizer = AdamW(model.parameters(), lr=args.lr, correct_bias=True)
# Prepare model for FP16 and distributed training if needed (order is important, distributed should be the last)
if args.fp16:
from apex import amp # Apex is only required if we use fp16 training
model, optimizer = amp.initialize(model,
optimizer,
opt_level=args.fp16)
if args.distributed:
model = DistributedDataParallel(model,
device_ids=[args.local_rank],
output_device=args.local_rank,
find_unused_parameters=True)
logger.info("Prepare datasets")
train_loader, val_loader, train_sampler, valid_sampler = get_data_loaders(
args, tokenizer)
# Training function and trainer
def update(engine, batch):
model.train()
batch = tuple(batch[input_name].to(args.device)
for input_name in MODEL_INPUTS)
#batch = tuple(input_tensor.to(args.device) for input_tensor in batch)
encoder_mask, decoder_mask, encoder_input_ids, decoder_input_ids, lm_labels, token_type_ids, decoder_lang_id = batch
model_kwargs = {
"encoder_token_type_ids": token_type_ids,
"decoder_token_type_ids": decoder_lang_id,
"encoder_attention_mask": encoder_mask,
"decoder_attention_mask": decoder_mask,
"decoder_lm_labels": lm_labels
}
lm_loss, prediction_scores, *_ = model(
encoder_input_ids=encoder_input_ids,
decoder_input_ids=decoder_input_ids,
**model_kwargs)
loss = (lm_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_norm)
else:
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_norm)
if engine.state.iteration % args.gradient_accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
return loss.item()
trainer = Engine(update)
# Evaluation function and evaluator (evaluator output is the input of the metrics)
def inference(engine, batch):
model.eval()
with torch.no_grad():
batch = tuple(batch[input_name].to(args.device)
for input_name in MODEL_INPUTS)
#batch = tuple(input_tensor.to(args.device) for input_tensor in batch)
encoder_mask, decoder_mask, encoder_input_ids, decoder_input_ids, lm_labels, token_type_ids, decoder_lang_id = batch
logger.info(tokenizer.decode(encoder_input_ids[0, :].tolist()))
# if we dont send labels to model, it doesnt return losses
model_kwargs = {
"encoder_token_type_ids": token_type_ids,
"decoder_token_type_ids": decoder_lang_id,
"encoder_attention_mask": encoder_mask,
"decoder_attention_mask": decoder_mask
}
lm_logits, *_ = model(encoder_input_ids=encoder_input_ids,
decoder_input_ids=decoder_input_ids,
**model_kwargs)
lm_logits_flat_shifted = lm_logits[..., :-1, :].contiguous().view(
-1, lm_logits.size(-1))
lm_labels_flat_shifted = lm_labels[..., 1:].contiguous().view(-1)
return (lm_logits_flat_shifted, ), (lm_labels_flat_shifted, )
evaluator = Engine(inference)
# Attach evaluation to trainer: we evaluate when we start the training and at the end of each epoch
trainer.add_event_handler(Events.EPOCH_COMPLETED,
lambda _: evaluator.run(val_loader))
if args.n_epochs < 1:
trainer.add_event_handler(Events.COMPLETED,
lambda _: evaluator.run(val_loader))
if args.eval_before_start:
trainer.add_event_handler(Events.STARTED,
lambda _: evaluator.run(val_loader))
# Make sure distributed data samplers split the dataset nicely between the distributed processes
if args.distributed:
trainer.add_event_handler(
Events.EPOCH_STARTED,
lambda engine: train_sampler.set_epoch(engine.state.epoch))
evaluator.add_event_handler(
Events.EPOCH_STARTED,
lambda engine: valid_sampler.set_epoch(engine.state.epoch))
# Linearly decrease the learning rate from lr to zero
scheduler = PiecewiseLinear(optimizer, "lr",
[(0, args.lr),
(args.n_epochs * len(train_loader), 0.0)])
trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
# Prepare metrics - note how we compute distributed metrics
RunningAverage(output_transform=lambda x: x).attach(trainer, "loss")
metrics = {
"nll":
Loss(torch.nn.CrossEntropyLoss(ignore_index=-1),
output_transform=lambda x: (x[0][0], x[1][0]))
}
metrics.update({
"average_nll":
MetricsLambda(average_distributed_scalar, metrics["nll"], args)
})
metrics["average_ppl"] = MetricsLambda(math.exp, metrics["average_nll"])
for name, metric in metrics.items():
metric.attach(evaluator, name)
# On the main process: add progress bar, tensorboard, checkpoints and save model, configuration and tokenizer before we start to train
if args.local_rank in [-1, 0]:
pbar = ProgressBar(persist=True)
pbar.attach(trainer, metric_names=["loss"])
evaluator.add_event_handler(
Events.COMPLETED, lambda _: pbar.log_message(
"Validation: %s" % pformat(evaluator.state.metrics)))
log_dir = make_logdir(args.model_checkpoint)
log_dir += "_lang_id"
if args.random_init:
log_dir = log_dir + "_random_init"
tb_logger = TensorboardLogger(log_dir)
tb_logger.attach(trainer,
log_handler=OutputHandler(tag="training",
metric_names=["loss"]),
event_name=Events.ITERATION_COMPLETED)
tb_logger.attach(trainer,
log_handler=OptimizerParamsHandler(optimizer),
event_name=Events.ITERATION_STARTED)
tb_logger.attach(evaluator,
log_handler=OutputHandler(tag="validation",
metric_names=list(
metrics.keys()),
another_engine=trainer),
event_name=Events.EPOCH_COMPLETED)
checkpoint_handler = ModelCheckpoint(log_dir,
'checkpoint',
save_interval=1,
n_saved=3)
trainer.add_event_handler(
Events.EPOCH_COMPLETED, checkpoint_handler,
{'mymodel': getattr(model, 'module', model)
}) # "getattr" takes care of distributed encapsulation
torch.save(args, log_dir + '/model_training_args.bin')
if args.distributed:
getattr(model.module, 'encoder', model).config.to_json_file(
os.path.join(log_dir, CONFIG_NAME)
) # the config for encoder and decoder should be the same
else:
getattr(model, 'encoder', model).config.to_json_file(
os.path.join(log_dir, CONFIG_NAME)
) # the config for encoder and decoder should be the same
tokenizer.save_pretrained(log_dir)
# Run the training
trainer.run(train_loader, max_epochs=args.n_epochs)
# On the main process: close tensorboard logger and rename the last checkpoint (for easy re-loading with OpenAIGPTModel.from_pretrained method)
if args.local_rank in [-1, 0] and args.n_epochs > 0:
os.rename(
checkpoint_handler._saved[-1][1][-1],
os.path.join(log_dir, WEIGHTS_NAME)
) # TODO: PR in ignite to have better access to saved file paths (cleaner)
tb_logger.close()
class DatasetExperiment(BaseExperiment):
def __init__(self, train, niter, nepoch, eval=None, gpu_id=[0],
sacred_run=None, writers=None, root=None, corruption=None,
**kwargs):
super(DatasetExperiment, self).__init__(**kwargs)
self.train = train
self.eval = eval
self.corruption = corruption
self.niter = niter
self.nepoch = nepoch
self.device = 'cuda:0'
self.sacred_run = sacred_run
self.gpu_id = gpu_id
if root is not None:
self.basedir = make_basedir(os.path.join(root, str(sacred_run._id)))
else:
writers = None
checkpoint = None
if writers is not None:
self.writers = init_writers(*writers, sacred_run=sacred_run, dirname=self.basedir)
else:
self.writers = None
self.create_trainer()
def create_trainer(self):
self.trainer = Engine(self.train_step)
self.trainer.add_event_handler(Events.EPOCH_COMPLETED, self.K_step)
self.trainer.add_event_handler(Events.EPOCH_COMPLETED, self.log_metrics, 'train')
self.trainer.add_event_handler(Events.ITERATION_COMPLETED, self.write_metrics, 'train')
self.pbar = ProgressBar()
self.timer = Timer(average=True)
self.timer.attach(self.trainer, start=Events.EPOCH_STARTED, resume=Events.ITERATION_STARTED,
pause=Events.ITERATION_COMPLETED, step=Events.ITERATION_COMPLETED)
self.trainer.add_event_handler(Events.EPOCH_COMPLETED, self.print_times)
def print_times(self, engine):
iteration = engine.state.iteration
if self.writers is not None:
self.writers.add_scalar('time_per_epoch', self.timer.total, iteration)
self.writers.add_scalar('time_per_batch', self.timer.value(), iteration)
self.timer.reset()
def train_step(self, engine, batch):
self.training()
if isinstance(batch, Mapping):
return self.optimize(**convert_tensor(batch, self.device))
else:
raise NotImplementedError
def infer(self, **kwargs):
raise NotImplementedError
def optimize(self, **kwargs):
raise NotImplementedError
def K_step(self, **kwargs):
raise NotImplementedError
def log_metrics(self, engine, dataset_name):
iteration = self.trainer.state.iteration
epoch = self.trainer.state.epoch
log = f"EPOCH {epoch}" f" - ITER {iteration} - {dataset_name}"
if self.sacred_run is not None:
log = f"ID {self.sacred_run._id} - " + log
for metric, value in engine.state.metrics.items():
if value is None:
value = float('nan')
log += f" {metric}: {value:.6f}"
log += self.extra_log()
self.pbar.log_message(log)
def extra_log(self):
return ""
def write_metrics(self, engine, dataset_name):
if self.writers is not None:
for metric, value in engine.state.metrics.items():
name = dataset_name + '/' + metric
self.writers.add_scalar(name, value, self.trainer.state.iteration)
def run(self):
self.trainer.run(self.train, max_epochs=self.nepoch)
def run_training(model, optimizer, scheduler, output_path,
train_loader, val_loader, epochs, patience,
epochs_pretrain, mixed_precision, classes_weights):
# trainer
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
if classes_weights is not None:
classes_weights = classes_weights.to(device)
crit = nn.CrossEntropyLoss(weight=classes_weights)
metrics = {"accuracy": Accuracy(), "loss": Loss(crit)}
trainer = create_supervised_trainer_with_pretraining(
model, optimizer, crit, device=device, epochs_pretrain=epochs_pretrain,
mixed_precision=mixed_precision)
train_evaluator = create_supervised_evaluator(
model, metrics=metrics, device=device)
val_evaluator = create_supervised_evaluator(
model, metrics=metrics, device=device)
# Out paths
path_ckpt = os.path.join(output_path, "model_ckpt")
log_dir = os.path.join(output_path, "log_dir")
os.makedirs(log_dir, exist_ok=True)
# tensorboard
tb_logger = TensorboardLogger(log_dir=log_dir)
tb_logger.attach(train_evaluator, log_handler=OutputHandler(tag="training", metric_names=[
"accuracy", "loss"], another_engine=trainer), event_name=Events.EPOCH_COMPLETED)
tb_logger.attach(val_evaluator, log_handler=OutputHandler(tag="validation", metric_names=[
"accuracy", "loss"], another_engine=trainer), event_name=Events.EPOCH_COMPLETED)
# training progress
pbar = ProgressBar(persist=True)
pbar.attach(trainer, metric_names="all")
# @trainer.on(Events.EPOCH_COMPLETED)
def log_training_results(engine):
train_evaluator.run(train_loader)
val_evaluator.run(val_loader)
train_loss = train_evaluator.state.metrics["loss"]
val_loss = val_evaluator.state.metrics["loss"]
train_acc = train_evaluator.state.metrics["accuracy"]
val_acc = val_evaluator.state.metrics["accuracy"]
pbar.log_message(
"Training Results - Epoch: {} Loss: {:.6f} Accuracy: {:.6f}".format(engine.state.epoch, train_loss, train_acc))
pbar.log_message(
"Validation Results - Epoch: {} Loss: {:.6f} Accuracy: {:.6f}".format(engine.state.epoch, val_loss, val_acc))
pbar.n = pbar.last_print_n = 0
trainer.add_event_handler(Events.EPOCH_COMPLETED, log_training_results)
# def get_val_loss(engine):
# return -engine.state.metrics['loss']
def get_val_acc(engine):
return engine.state.metrics['accuracy']
# checkpoint and early stopping
checkpointer = ModelCheckpoint(
path_ckpt, "model", score_function=get_val_acc, score_name="accuracy", require_empty=False)
early_stopper = EarlyStopping(patience, get_val_acc, trainer)
to_save = {'optimizer': optimizer, 'model': model}
if scheduler is not None:
to_save["scheduler"] = scheduler
val_evaluator.add_event_handler(Events.COMPLETED, checkpointer, to_save)
val_evaluator.add_event_handler(Events.COMPLETED, early_stopper)
if scheduler is not None:
trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
# free resources
trainer.add_event_handler(
Events.ITERATION_COMPLETED, lambda _: _empty_cache())
train_evaluator.add_event_handler(
Events.ITERATION_COMPLETED, lambda _: _empty_cache())
val_evaluator.add_event_handler(
Events.ITERATION_COMPLETED, lambda _: _empty_cache())
trainer.run(train_loader, max_epochs=epochs)
tb_logger.close()
# Evaluation with best model
model.load_state_dict(torch.load(
glob.glob(os.path.join(path_ckpt, "*.pth"))[0])["model"])
train_evaluator = create_supervised_evaluator(
model, metrics=metrics, device=device)
val_evaluator = create_supervised_evaluator(
model, metrics=metrics, device=device)
train_evaluator.run(train_loader)
val_evaluator.run(val_loader)
_pretty_print("Evaluating best model")
pbar.log_message(
"Best model on training set - Loss: {:.6f} Accuracy: {:.6f}"
.format(train_evaluator.state.metrics["loss"], train_evaluator.state.metrics["accuracy"]))
pbar.log_message(
"Best model on validation set - Loss: {:.6f} Accuracy: {:.6f}"
.format(val_evaluator.state.metrics["loss"], val_evaluator.state.metrics["accuracy"]))
return model, train_evaluator.state.metrics, val_evaluator.state.metrics
def finetune_model(args, model, loader):
optimizer = AdamW(model.parameters(), lr=args.lr, correct_bias=True)
def update(engine, batch):
model.train()
batch = tuple(batch[input_name].to(args.device)
for input_name in MODEL_INPUTS)
input_ids, lm_labels, token_type_ids, nodes_ids, attention_mask = batch
if (not args.graph and not args.edge_list):
nodes_ids = None
if (not args.unilm): attention_mask = None
(lm_loss), *_ = model(input_ids=input_ids,
token_type_ids=token_type_ids,
labels=lm_labels,
nodes=nodes_ids,
attention_mask=attention_mask)
loss = lm_loss / args.gradient_accumulation_steps
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_norm)
if engine.state.iteration % args.gradient_accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
return loss.item()
# Evaluation function and evaluator (evaluator output is the input of the metrics)
def inference(engine, batch):
model.eval()
with torch.no_grad():
batch = tuple(batch[input_name].to(args.device)
for input_name in MODEL_INPUTS)
input_ids, lm_labels, token_type_ids, nodes_ids, attention_mask = batch
if (not args.graph and not args.edge_list):
nodes_ids = None
if (not args.unilm): attention_mask = None
# if we dont send labels to model, it doesnt return losses
lm_logits, *_ = model(input_ids=input_ids,
token_type_ids=token_type_ids,
nodes=nodes_ids,
attention_mask=attention_mask)
lm_logits_flat_shifted = lm_logits[..., :-1, :].contiguous().view(
-1, lm_logits.size(-1))
lm_labels_flat_shifted = lm_labels[..., 1:].contiguous().view(-1)
return (lm_logits_flat_shifted, ), (lm_labels_flat_shifted, )
trainer = Engine(update)
evaluator = Engine(inference)
trainer.add_event_handler(Events.EPOCH_COMPLETED,
lambda _: evaluator.run(loader))
if args.n_epochs < 1:
trainer.add_event_handler(Events.COMPLETED,
lambda _: evaluator.run(loader))
if args.eval_before_start:
trainer.add_event_handler(Events.STARTED,
lambda _: evaluator.run(loader))
# Linearly decrease the learning rate from lr to zero
scheduler = PiecewiseLinear(optimizer, "lr",
[(0, args.lr),
(args.n_epochs * len(loader), 0.0)])
trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
# Prepare metrics - note how we compute distributed metrics
RunningAverage(output_transform=lambda x: x).attach(trainer, "loss")
metrics = {
"nll":
Loss(torch.nn.CrossEntropyLoss(),
output_transform=lambda x: (x[0][0], x[1][0]))
}
metrics.update({
"average_nll":
MetricsLambda(average_distributed_scalar, metrics["nll"], args)
})
metrics["average_ppl"] = MetricsLambda(math.exp, metrics["average_nll"])
for name, metric in metrics.items():
metric.attach(evaluator, name)
pbar = ProgressBar(persist=True)
pbar.attach(trainer, metric_names=["loss"])
evaluator.add_event_handler(
Events.COMPLETED, lambda _: pbar.log_message(
"Validation: %s" % pformat(evaluator.state.metrics)))
trainer.run(loader, max_epochs=args.n_epochs)
return model
scheduler = PiecewiseLinear(optimizer, "lr", [(0, args.lr), (args.n_epochs * len(train_loader), 0.0)])
trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
# Prepare metrics - note how we compute distributed metrics
RunningAverage(output_transform=lambda x: x).attach(trainer, "loss")
metrics = {"nll": Loss(torch.nn.CrossEntropyLoss(), output_transform=lambda x: (x[0][0], x[1][0]))}
metrics.update({"average_nll": MetricsLambda(average_distributed_scalar, metrics["nll"], args)})
metrics["average_ppl"] = MetricsLambda(math.exp, metrics["average_nll"])
for name, metric in metrics.items():
metric.attach(evaluator, name)
# On the main process: add progress bar, tensorboard, checkpoints and save model, configuration and tokenizer before we start to train
if args.local_rank in [-1, 0]:
pbar = ProgressBar(persist=True)
pbar.attach(trainer, metric_names=["loss"])
evaluator.add_event_handler(Events.COMPLETED, lambda _: pbar.log_message("Validation: %s" % pformat(evaluator.state.metrics)))
log_dir = make_logdir(args,args.model_checkpoint)
tb_logger = TensorboardLogger(log_dir)
tb_logger.attach(trainer, log_handler=OutputHandler(tag="training", metric_names=["loss"]), event_name=Events.ITERATION_COMPLETED)
tb_logger.attach(trainer, log_handler=OptimizerParamsHandler(optimizer), event_name=Events.ITERATION_STARTED)
tb_logger.attach(evaluator, log_handler=OutputHandler(tag="validation", metric_names=list(metrics.keys()), another_engine=trainer), event_name=Events.EPOCH_COMPLETED)
checkpoint_handler = ModelCheckpoint(log_dir, 'checkpoint', save_interval=1, n_saved=3)
trainer.add_event_handler(Events.EPOCH_COMPLETED, checkpoint_handler, {'mymodel': getattr(model, 'module', model)}) # "getattr" takes care of distributed encapsulation
torch.save(args, log_dir + '/model_training_args.bin')
getattr(model, 'module', model).config.to_json_file(os.path.join(log_dir, CONFIG_NAME))
tokenizer.save_pretrained(log_dir)
def train(args):
logger.info("Prepare tokenizer, pretrained model and optimizer.")
tokenizer, _, vocab = get_kogpt2_tokenizer()
model = get_kogpt2_model()
model.to(args.device)
optimizer = AdamW(model.parameters(), lr=args.lr, correct_bias=True)
logger.info("Prepare datasets")
train_loader, val_loader = get_data_loaders(args, tokenizer, vocab)
def update(engine, batch):
model.train()
batch = tuple(input_tensor.to(args.device) for input_tensor in batch)
input_ids, labels, token_type_ids = batch
loss, *_ = model(input_ids,
token_type_ids=token_type_ids,
labels=labels)
loss = loss / args.gradient_accumulation_steps
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_norm)
if engine.state.iteration % args.gradient_accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
return loss.item()
trainer = Engine(update)
def inference(engine, batch):
model.eval()
with torch.no_grad():
batch = tuple(
input_tensor.to(args.device) for input_tensor in batch)
input_ids, labels, token_type_ids = batch
# logger.info(tokenizer.decode(input_ids[0, -1, :].tolist()))
# if we dont send labels to model, it doesnt return losses
logits, *_ = model(input_ids, token_type_ids=token_type_ids)
logits_flat_shifted = logits[..., :-1, :].contiguous().view(
-1, logits.size(-1))
labels_flat_shifted = labels[..., 1:].contiguous().view(-1)
return (logits_flat_shifted), (labels_flat_shifted)
evaluator = Engine(inference)
# Attach evaluation to trainer: we evaluate when we start the training and at the end of each epoch
trainer.add_event_handler(Events.EPOCH_COMPLETED,
lambda _: evaluator.run(val_loader))
if args.n_epochs < 1:
trainer.add_event_handler(Events.COMPLETED,
lambda _: evaluator.run(val_loader))
if args.eval_before_start:
trainer.add_event_handler(Events.STARTED,
lambda _: evaluator.run(val_loader))
# Linearly decrease the learning rate from lr to zero
scheduler = PiecewiseLinear(optimizer, "lr",
[(0, args.lr),
(args.n_epochs * len(train_loader), 0.0)])
trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
# Prepare metrics - note how we compute distributed metrics
RunningAverage(output_transform=lambda x: x).attach(trainer, "loss")
metrics = {
"nll":
Loss(torch.nn.CrossEntropyLoss(ignore_index=-100),
output_transform=lambda x: (x[0], x[1])),
"accuracy":
Accuracy(output_transform=lambda x: (x[0], x[1]))
}
for name, metric in metrics.items():
metric.attach(evaluator, name)
# On the main process: add progress bar, tensorboard, checkpoints and save model,
# configuration and tokenizer before we start to train
pbar = ProgressBar(persist=True)
pbar.attach(trainer, metric_names=["loss"])
evaluator.add_event_handler(
Events.COMPLETED, lambda _: pbar.log_message(
"Validation: %s" % pformat(evaluator.state.metrics)))
log_dir = make_logdir("kogpt2_personachat")
tb_logger = TensorboardLogger(log_dir)
tb_logger.attach(trainer,
log_handler=OutputHandler(tag="training",
metric_names=["loss"]),
event_name=Events.ITERATION_COMPLETED)
tb_logger.attach(trainer,
log_handler=OptimizerParamsHandler(optimizer),
event_name=Events.ITERATION_STARTED)
tb_logger.attach(
evaluator,
log_handler=OutputHandler(
tag="validation",
metric_names=list(metrics.keys()),
global_step_transform=global_step_from_engine(trainer)),
event_name=Events.EPOCH_COMPLETED)
checkpoint_handler = ModelCheckpoint(log_dir,
'checkpoint',
save_interval=1,
n_saved=3)
trainer.add_event_handler(
Events.EPOCH_COMPLETED, checkpoint_handler,
{'mymodel': getattr(model, 'module', model)
}) # "getattr" takes care of distributed encapsulation
torch.save(args, log_dir + '/model_training_args.bin')
getattr(model, 'module',
model).config.to_json_file(os.path.join(log_dir, CONFIG_NAME))
# tokenizer.save_pretrained(log_dir)
# Run the training
trainer.run(train_loader, max_epochs=args.n_epochs)
# On the main process: close tensorboard logger and rename the last checkpoint (for easy re-loading with OpenAIGPTModel.from_pretrained method)
# TODO: PR in ignite to have better access to saved file paths (cleaner)
os.rename(os.path.join(log_dir, checkpoint_handler._saved[-1][1]),
os.path.join(log_dir, WEIGHTS_NAME))
tb_logger.close()
def train():
os.environ['CUDA_VISIBLE_DEVICES'] = '7'
parser = ArgumentParser()
parser.add_argument('--gpt2', action='store_true', help="use gpt2")
parser.add_argument("--model_checkpoint", type=str, default="uer/gpt2-chinese-cluecorpussmall", help="Path or URL of the model")
parser.add_argument("--from_step", type=int, default=-1, help="Init learning rate from this step")
parser.add_argument('--pretrained', action='store_true', help="If False train from scratch")
parser.add_argument("--data_path", type=str, default="data/autocloze.json",
help="Path or url of the dataset. ")
parser.add_argument("--train_path", type=str, default="data/toy_train.txt",
help="Path of the train dataset for dist dataset. ")
parser.add_argument("--valid_path", type=str, default="data/toy_valid.txt",
help="Path of the valid dataset for dist dataset. ")
#--------------------------------------------------------------
parser.add_argument("--dataset_cache", type=str, default="dataset_zh",
help="Path or url of the dataset cache")
parser.add_argument('--log_file', '-log_file', type=str, default="", help="Output logs to a file under this path")
parser.add_argument("--num_workers", type=int, default=8, help="Number of subprocesses for data loading")
parser.add_argument("--n_epochs", type=int, default=40, help="Number of training epochs")
parser.add_argument("--train_batch_size", type=int, default=1, help="Batch size for training")
parser.add_argument("--valid_batch_size", type=int, default=1, help="Batch size for validation")
parser.add_argument("--max_history", type=int, default=15, help="Number of previous exchanges to keep in history")
parser.add_argument("--scheduler", type=str, default="noam", choices=['noam', 'linear'], help="method of optim")
parser.add_argument("--n_emd", type=int, default=768, help="Number of n_emd in config file (for noam)")
parser.add_argument("--lr", type=float, default=5e-5, help="Learning rate")
parser.add_argument("--eval_before_start", action='store_true',
help="If true start with a first evaluation before training")
parser.add_argument("--warmup_steps", type=int, default=5000, help="Warm up steps")
parser.add_argument("--valid_steps", type=int, default=5000, help="Perfom validation every X steps")
parser.add_argument("--gradient_accumulation_steps", type=int, default=64,
help="Accumulate gradients on several steps")
parser.add_argument("--max_norm", type=float, default=1.0, help="Clipping gradient norm")
parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu",
help="Device (cuda or cpu)")
parser.add_argument("--fp16", type=str, default="",
help="Set to O0, O1, O2 or O3 for fp16 training (see apex documentation)")
parser.add_argument("--local_rank", type=int, default=-1,
help="Local rank for distributed training (-1: not distributed)")
args = parser.parse_args()
print('cuda ',torch.cuda.is_available())
# logging is set to INFO (resp. WARN) for main (resp. auxiliary) process.
# logger.info => log main process only, logger.warning => log all processes
logging.basicConfig(level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN)
logger.warning("Running process %d", args.local_rank)
logger.info("Arguments: %s", pformat(args))
# Initialize distributed training if needed
args.distributed = (args.local_rank != -1)
'''if args.distributed:
torch.cuda.set_device(args.local_rank)
args.device = torch.device("cuda", args.local_rank)
torch.distributed.init_process_group(backend='nccl', init_method='env://')
'''
args.device = torch.device("cuda")
print('device ',args.device)
logger.info("Prepare tokenizer, pretrained model and optimizer - add special tokens for fine-tuning")
#model_class = OpenAIGPTLMHeadModel if not args.gpt2 else GPT2LMHeadModel
#config_class = OpenAIGPTConfig if not args.gpt2 else GPT2Config
model_class = GPT2LMHeadModel
config_class = GPT2Config
tokenizer_class = BertTokenizer
print('pretrained:',args.pretrained)
if args.pretrained:
print("----------------pretrained")
tokenizer = BertTokenizer.from_pretrained(args.model_checkpoint, do_lower_case=True)
model = GPT2LMHeadModel.from_pretrained(args.model_checkpoint)
else:
tokenizer = BertTokenizer.from_pretrained("uer/gpt2-chinese-cluecorpussmall")
model = GPT2LMHeadModel.from_pretrained("uer/gpt2-chinese-cluecorpussmall",from_tf=True)
#print('generate')
#print(text_generator("这是很久之前的事情了", max_length=100, do_sample=True))
#args.device=torch.device("cuda", 2)
model.to(args.device)
optimizer = AdamW([{'params': model.parameters(), 'initial_lr': args.lr}], lr=args.lr, correct_bias=True)
logger.info("Prepare datasets")
loader_class = build_dist_loaders if not args.data_path else build_dataloaders
train_loader, val_loader, train_sampler, valid_sampler = loader_class(args, tokenizer, logger)
logger.info("Prepare datasets ends")
# Prepare model for FP16 and distributed training if needed (order is important, distributed should be the last)
if args.fp16:
from apex import amp # Apex is only required if we use fp16 training
model, optimizer = amp.initialize(model, optimizer, opt_level=args.fp16)
if args.distributed:
model = DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank)
model=model.module
#if isinstance(model,torch.nn.DataParallel):
#print('params:',params_count(model))
#tokens_embed = model.transformer.get_input_embeddings()
# Training function and trainer
def update(engine, batch):
input_ids, token_type_ids, lm_labels = tuple(input_tensor.to(args.device) for input_tensor in batch)
#for i in range(input_ids.size()[0]):
# for j in range(input_ids.size()[1]):
# if input_ids[i,j]==-1:
# input_ids[i,j]=-100
# if lm_labels[i,j]==-1:
# lm_labels[i,j]=-100
#one=torch.tensor(-100)
#input_ids=torch.where(input_ids==-1,one,input_ids)
#lm_labels=torch.where(lm_labels==-1,one,lm_labels)
#print('traindata',input_ids,lm_labels)
#lm_labels=input_ids
r'''input_shape = input_ids.siz`e`()
input_ids = input_ids.view(-1, input_shape[-1])
inputs_embeds = tokens_embed(input_ids) * math.sqrt(tokens_embed.embedding_dim)'''
model.train()
#(lm_loss), *_ = model(inputs_embeds=inputs_embeds, labels=lm_labels,return_dict=0)
(lm_loss), *_ = model(input_ids=input_ids, labels=lm_labels,return_dict=False)
#print('lm_loss',lm_loss)
loss = lm_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_norm)
else:
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_norm)
if engine.state.iteration % args.gradient_accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
return loss.item(), optimizer.param_groups[0]['lr']
trainer = Engine(update)
# Evaluation function and evaluator (evaluator output is the input of the metrics)
cntepoch=0
def inference(engine, batch):
model.eval()
with torch.no_grad():
input_ids, token_type_ids, lm_labels = tuple(input_tensor.to(args.device) for input_tensor in batch)
# logger.info(tokenizer.decode(input_ids[0, -1, :].tolist()))
#one = torch.tensor(-100)
#input_ids=torch.where(input_ids==-1,one,input_ids)
#print('validdata',input_ids,lm_labels)
#lm_labels=input_ids
r'''input_shape = input_ids.size()
input_ids = input_ids.view(-1, input_shape[-1])
inputs_embeds = tokens_embed(input_ids) * math.sqrt(tokens_embed.embedding_dim)'''
#lm_logits, *_ = model(inputs_embeds=inputs_embeds,return_dict=0)
lm_logits, *_ = model(input_ids=input_ids,return_dict=False)
lm_logits_flat_shifted = lm_logits[..., :-1, :].contiguous().view(-1, lm_logits.size(-1))
lm_labels_flat_shifted = lm_labels[..., 1:].contiguous().view(-1)
return lm_logits_flat_shifted, lm_labels_flat_shifted
cntepoch+=1
torch.save(args, tb_logger.writer.logdir + '_%s/model_training_args.bin'%(str(cntepoch)))
evaluator = Engine(inference)
# Attach evaluation to trainer: we evaluate when we start the training and at the end of each epoch
trainer.add_event_handler(Events.EPOCH_COMPLETED, lambda _: evaluator.run(val_loader))
if args.n_epochs < 1:
trainer.add_event_handler(Events.COMPLETED, lambda _: evaluator.run(val_loader))
if args.eval_before_start:
trainer.add_event_handler(Events.STARTED, lambda _: evaluator.run(val_loader))
# Evaluation during training
@trainer.on(Events.ITERATION_STARTED)
def log_iterations(engine):
# if engine.state.iteration % max(int(0.1 * len(train_loader)), 1) == 0:
if engine.state.iteration % args.valid_steps == 0:
evaluator.run(val_loader)
# Make sure distributed data samplers split the dataset nicely between the distributed processes
if args.distributed:
trainer.add_event_handler(Events.EPOCH_STARTED, lambda engine: train_sampler.set_epoch(engine.state.epoch))
evaluator.add_event_handler(Events.EPOCH_STARTED, lambda engine: valid_sampler.set_epoch(engine.state.epoch))
# noam decrease the learning rate
# model_size = model.config.n_embd
model_size = args.n_emd
noam_lambda = lambda step: (
model_size ** (-0.5) * min((step + 1) ** (-0.5), (step + 1) * args.warmup_steps ** (-1.5)))
noam_scheduler = LambdaLR(optimizer, lr_lambda=noam_lambda, last_epoch=args.from_step)
scheduler = LRScheduler(noam_scheduler)
if args.scheduler == "linear":
scheduler = PiecewiseLinear(optimizer, "lr", [(0, args.lr), (args.n_epochs * len(train_loader), 0.0)])
trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
# Prepare metrics - note how we compute distributed metrics
RunningAverage(output_transform=lambda x: x[0]).attach(trainer, "loss")
RunningAverage(output_transform=lambda x: x[1]).attach(trainer, "lr")
metrics = {"nll": Loss(torch.nn.CrossEntropyLoss(ignore_index=-100), output_transform=lambda x: (x[0], x[1]))}
metrics.update({"average_nll": MetricsLambda(average_distributed_scalar, metrics["nll"], args)})
metrics["average_ppl"] = MetricsLambda(math.exp, metrics["average_nll"])
for name, metric in metrics.items():
metric.attach(evaluator, name)
# On the main process: add progress bar, tensorboard, checkpoints
# And save model, configuration and tokenizer before we start to train
if args.local_rank in [-1, 0]:
pbar = ProgressBar(persist=True, mininterval=2)
pbar.attach(trainer, metric_names=["loss", "lr"])
evaluator.add_event_handler(Events.COMPLETED,
lambda _: pbar.log_message("Validation: %s" % pformat(evaluator.state.metrics)))
tb_logger = TensorboardLogger(log_dir=None)
tb_logger.attach(trainer, log_handler=OutputHandler(tag="training", metric_names=["loss"]),
event_name=Events.ITERATION_COMPLETED)
tb_logger.attach(trainer, log_handler=OptimizerParamsHandler(optimizer), event_name=Events.ITERATION_STARTED)
tb_logger.attach(evaluator, log_handler=OutputHandler(tag="validation", metric_names=list(metrics.keys()),
another_engine=trainer),
event_name=Events.EPOCH_COMPLETED)
checkpoint_handler = ModelCheckpoint(tb_logger.writer.logdir, 'checkpoint', save_interval=1, n_saved=6)
# save model after evaluation
evaluator.add_event_handler(Events.EPOCH_COMPLETED, checkpoint_handler, {
'mymodel': getattr(model, 'module', model)})
trainer.add_event_handler(Events.EPOCH_COMPLETED, checkpoint_handler, {
'mymodel': getattr(model, 'module', model)}) # "getattr" take care of distributed encapsulation
torch.save(args, tb_logger.writer.logdir + '/model_training_args.bin')
getattr(model, 'module', model).config.to_json_file(os.path.join(tb_logger.writer.logdir, CONFIG_NAME))
tokenizer.save_vocabulary(tb_logger.writer.logdir)
# Run the training
trainer.run(train_loader, max_epochs=args.n_epochs)
# On the main process: close tensorboard logger and rename the last checkpoint
# (for easy re-loading with OpenAIGPTModel.from_pretrained method)
if args.local_rank in [-1, 0] and args.n_epochs > 0:
os.rename(checkpoint_handler._saved[-1][1][-1],
os.path.join(tb_logger.writer.logdir,
WEIGHTS_NAME)) # TODO: PR in ignite to have better access to saved file paths (cleaner)
tb_logger.close()
def main():
parser = ArgumentParser()
parser.add_argument(
"--dataset_path",
type=str,
default="",
help="Path or url of the dataset. If empty download from S3.")
parser.add_argument("--dataset_cache",
type=str,
default='./dataset_cache',
help="Path or url of the dataset cache")
parser.add_argument("--model",
type=str,
default="",
help="Model type, one of: %s" %
', '.join(MODELS.keys()))
parser.add_argument("--model_checkpoint",
type=str,
default="",
help="Path, url or short name of a pretrained model")
parser.add_argument("--num_candidates",
type=int,
default=2,
help="Number of candidates for training")
parser.add_argument("--max_history",
type=int,
default=2,
help="Number of previous exchanges to keep in history")
parser.add_argument("--train_batch_size",
type=int,
default=4,
help="Batch size for training")
parser.add_argument("--valid_batch_size",
type=int,
default=4,
help="Batch size for validation")
parser.add_argument("--gradient_accumulation_steps",
type=int,
default=8,
help="Accumulate gradients on several steps")
parser.add_argument("--lr",
type=float,
default=6.25e-5,
help="Learning rate")
parser.add_argument("--lm_coef",
type=float,
default=1.0,
help="LM loss coefficient")
parser.add_argument("--mc_coef",
type=float,
default=1.0,
help="Multiple-choice loss coefficient")
parser.add_argument("--adv_coef",
type=float,
default=1.0,
help="Adversarial dataset prediction loss coefficient")
parser.add_argument("--max_norm",
type=float,
default=1.0,
help="Clipping gradient norm")
parser.add_argument("--weight_decay",
default=0.0,
type=float,
help="Weight decay if we apply some.")
parser.add_argument("--warmup_steps",
default=0,
type=int,
help="Linear warmup over warmup_steps.")
parser.add_argument("--n_epochs",
type=int,
default=3,
help="Number of training epochs")
#parser.add_argument("--personality_permutations", type=int, default=1, help="Number of permutations of personality sentences")
parser.add_argument(
"--eval_before_start",
action='store_true',
help="If true start with a first evaluation before training")
#parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu", help="Device (cuda or cpu)")
parser.add_argument("--no_cuda",
action='store_true',
help="Avoid using CUDA when available")
parser.add_argument(
"--fp16",
type=str,
default="",
help=
"Set to O0, O1, O2 or O3 for fp16 training (see apex documentation)")
parser.add_argument(
"--local_rank",
type=int,
default=-1,
help="Local rank for distributed training (-1: not distributed)")
parser.add_argument(
"--max_sequence_length",
type=int,
default=-1,
help="If set, use this to manually restrict the sequence length. "
"This might be helpful to save resources (memory). "
"If not set, this is looked up from the model config (n_ctx value).")
parser.add_argument(
"--adversarial_dataset_prediction",
action='store_true',
help="Set to train with adversarial dataset prediction")
parser.add_argument("--seed",
type=int,
default=None,
help='set random seed')
args = parser.parse_args()
# logging is set to INFO (resp. WARN) for main (resp. auxiliary) process. logger.info => log main process only, logger.warning => log all processes
logging.basicConfig(
level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN)
logger.warning(
"Running process %d", args.local_rank
) # This is a logger.warning: it will be printed by all distributed processes
logger.info("Arguments: %s", pformat(args))
if args.seed is not None:
torch.manual_seed(args.seed)
args.distributed = (args.local_rank != -1)
logger.info("Prepare tokenizer and data")
if not args.model:
logger.warning(
'"model" parameter is not set! This is deprecated. Please use one of: %s. '
'To mimic deprecated behaviour, "model_checkpoint" will be used as "model"'
% ', '.join(MODELS.keys()))
args.model = args.model_checkpoint
if args.model not in MODELS:
raise NotImplementedError(
'model "%s" not implemented. use one of: %s' %
(args.model, ', '.join(MODELS.keys())))
config_class, tokenizer_class, model_class, _ = MODELS[args.model]
if not args.model_checkpoint:
args.model_checkpoint = args.model
model_config = config_class.from_pretrained(args.model_checkpoint)
tokenizer = tokenizer_class.from_pretrained(args.model_checkpoint)
additional_special_tokens = [TYPE_BACKGROUND, TYPE_BOT, TYPE_USER]
# for adversarial training (dataset prediction)
dataset_labels = None
if args.adversarial_dataset_prediction:
dataset_labels = [
get_dataset_label(dataset_path)
for dataset_path in args.dataset_path.split(',')
]
#additional_special_tokens.extend(dataset_labels)
#if model_class not in ADV_MODELS.values():
assert model_class in ADV_MODELS, f'no adversarial model implemented for model class: {model_class.__name__}'
model_class = ADV_MODELS[model_class]
if not hasattr(model_config, 'cls'):
model_config.cls = {}
if 'dataset_labels' in model_config.cls:
assert all([dl in model_config.cls['dataset_labels']['labels'] for dl in dataset_labels]), \
f'loaded dataset_labels [{model_config.cls["dataset_labels"]["labels"]}] do not contain all ' \
f'current dataset_labels [{dataset_labels}]'
dataset_labels = model_config.cls['dataset_labels']['labels']
else:
model_config.cls['dataset_labels'] = {
'labels': dataset_labels,
'is_adversarial': True
}
model_input_names = [
"input_ids", "mc_token_ids", "lm_labels", "mc_labels",
"dataset_labels", "token_type_ids"
]
# not yet used
model_output_names = [
"lm_loss", "mc_loss", "cl_loss_0", "lm_logits", "mc_logits",
"cl_logits_0", "presents"
]
else:
model_input_names = [
"input_ids", "mc_token_ids", "lm_labels", "mc_labels",
"token_type_ids"
]
# not yet used
model_output_names = [
"lm_loss", "mc_loss", "lm_logits", "mc_logits", "presents"
]
tokenizer.add_special_tokens({
'bos_token':
TYPE_BOS,
'eos_token':
TYPE_EOS,
'pad_token':
TYPE_PAD,
'additional_special_tokens':
additional_special_tokens
})
logger.info("Prepare datasets")
max_sequence_length = model_config.n_ctx if args.max_sequence_length <= 0 else args.max_sequence_length
assert max_sequence_length <= model_config.n_ctx, 'max_sequence_length [%i] was set to a value higher than ' \
'supported by the model (config.n_ctx [%i]). Please use a lower ' \
'value or do not set it [-1] to use the highest supported one.' \
% (max_sequence_length, model_config.n_ctx)
train_loader, val_loader, train_sampler, valid_sampler = get_data_loaders(
args=args,
tokenizer=tokenizer,
model_input_names=model_input_names,
max_sequence_length=max_sequence_length,
dataset_labels=dataset_labels)
logger.info(
"Prepare pretrained model and optimizer - add special tokens for fine-tuning"
)
# Initialize distributed training if needed
# 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
# Load pretrained model and tokenizer
if args.local_rank not in [-1, 0]:
torch.distributed.barrier(
) # Barrier to make sure only the first process in distributed training download model & vocab
#model = model_class.from_pretrained(args.model_checkpoint, num_cl_labels=len(dataset_ids)) # for GPT2DoubleHeadsModelwithAdversarial
model = model_class.from_pretrained(args.model_checkpoint,
config=model_config)
model.resize_token_embeddings(len(tokenizer))
model.to(args.device)
if args.local_rank == 0:
torch.distributed.barrier(
) # End of barrier to make sure only the first process in distributed training download model & vocab
####################################################################################################################
# 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 = OpenAIAdam(model.parameters(), lr=args.lr)
optimizer = AdamW(optimizer_grouped_parameters, lr=args.lr)
# scheduler is set below (see ignite)
#scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=args.warmup_steps,
# num_training_steps=len(train_loader) // args.train_batch_size + 1)
# Check if saved optimizer or scheduler states exist
if os.path.isfile(os.path.join(
args.model_checkpoint, 'optimizer.pt')) and os.path.isfile(
os.path.join(args.model_checkpoint, 'scheduler.pt')):
# Load in optimizer and scheduler states
# TODO: this needs to be dumped somewhere
optimizer.load_state_dict(
torch.load(os.path.join(args.model_checkpoint, 'optimizer.pt')))
#scheduler.load_state_dict(torch.load(os.path.join(args.model_checkpoint, 'scheduler.pt')))
# Prepare model for FP16 and distributed training if needed (order is important, distributed should be the last)
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)
# 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)
# Training function and trainer
def update(engine, batch):
model.train()
batch = {
model_input_names[i]: input_tensor.to(args.device)
for i, input_tensor in enumerate(batch)
}
model_output = model(**batch)
losses = model_output[:
3] if args.adversarial_dataset_prediction else model_output[:
2]
if args.n_gpu > 1: # mean() to average on multi-gpu.
losses = list(losses)
for i in range(len(losses)):
losses[i] = losses[i].mean()
lm_loss, mc_loss = losses[0], losses[1]
loss = (lm_loss * args.lm_coef +
mc_loss * args.mc_coef) / args.gradient_accumulation_steps
# handle adversarial loss
loss_wo_adv = loss.clone()
if args.adversarial_dataset_prediction:
adv_loss = model_output[2]
loss += (adv_loss *
args.adv_coef) / 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_norm)
else:
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_norm)
if engine.state.iteration % args.gradient_accumulation_steps == 0:
optimizer.step()
#scheduler.step() # Update learning rate schedule # already DONE below!
optimizer.zero_grad()
return loss_wo_adv.item(), loss.item()
trainer = Engine(update)
# Evaluation function and evaluator (evaluator output is the input of the metrics)
def inference(engine, batch):
model.eval()
with torch.no_grad():
batch = tuple(
input_tensor.to(args.device) for input_tensor in batch)
if args.adversarial_dataset_prediction:
input_ids, mc_token_ids, lm_labels, mc_labels, dataset_labels, token_type_ids = batch
else:
input_ids, mc_token_ids, lm_labels, mc_labels, token_type_ids = batch
logger.debug(
tokenizer.decode(input_ids[0, -1, :].tolist()).replace(
TYPE_PAD, ''))
model_outputs = model(input_ids=input_ids,
mc_token_ids=mc_token_ids,
token_type_ids=token_type_ids)
lm_logits, mc_logits = model_outputs[0], model_outputs[
1] # So we can also use GPT2 outputs
lm_logits_flat_shifted = lm_logits[..., :-1, :].contiguous().view(
-1, lm_logits.size(-1))
lm_labels_flat_shifted = lm_labels[..., 1:].contiguous().view(-1)
return (lm_logits_flat_shifted,
mc_logits), (lm_labels_flat_shifted, mc_labels)
evaluator = Engine(inference)
# Attach evaluation to trainer: we evaluate when we start the training and at the end of each epoch
trainer.add_event_handler(Events.EPOCH_COMPLETED,
lambda _: evaluator.run(val_loader))
if args.n_epochs < 1:
trainer.add_event_handler(Events.COMPLETED,
lambda _: evaluator.run(val_loader))
if args.eval_before_start:
trainer.add_event_handler(Events.STARTED,
lambda _: evaluator.run(val_loader))
# Make sure distributed data samplers split the dataset nicely between the distributed processes
if args.distributed:
trainer.add_event_handler(
Events.EPOCH_STARTED,
lambda engine: train_sampler.set_epoch(engine.state.epoch))
evaluator.add_event_handler(
Events.EPOCH_STARTED,
lambda engine: valid_sampler.set_epoch(engine.state.epoch))
# Linearly decrease the learning rate from lr to zero (scheduler)
scheduler = PiecewiseLinear(optimizer, "lr",
[(0, args.lr),
(args.n_epochs * len(train_loader), 0.0)])
trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
# Prepare metrics - note how we compute distributed metrics
RunningAverage(output_transform=lambda x: x[0]).attach(trainer, "loss")
if args.adversarial_dataset_prediction:
RunningAverage(output_transform=lambda x: x[1]).attach(
trainer, "loss_w/_adv")
RunningAverage(output_transform=lambda x: x[1] - x[0]).attach(
trainer, "loss_only_adv")
metrics = {
"nll":
Loss(torch.nn.CrossEntropyLoss(ignore_index=-1),
output_transform=lambda x: (x[0][0], x[1][0])),
"accuracy":
Accuracy(output_transform=lambda x: (x[0][1], x[1][1]))
}
metrics.update({
"average_nll":
MetricsLambda(average_distributed_scalar, metrics["nll"], args),
"average_accuracy":
MetricsLambda(average_distributed_scalar, metrics["accuracy"], args)
})
metrics["average_ppl"] = MetricsLambda(math.exp, metrics["average_nll"])
for name, metric in metrics.items():
metric.attach(evaluator, name)
# On the main process: add progress bar, tensorboard, checkpoints and save model, configuration and tokenizer before we start to train
if args.local_rank in [-1, 0]:
pbar = ProgressBar(persist=True)
pbar.attach(trainer, metric_names=["loss"])
evaluator.add_event_handler(
Events.COMPLETED, lambda _: pbar.log_message(
"Validation: %s" % pformat(evaluator.state.metrics)))
tb_logger = TensorboardLogger(log_dir=None)
tb_logger.attach(trainer,
log_handler=OutputHandler(tag="training",
metric_names=["loss"]),
event_name=Events.ITERATION_COMPLETED)
if args.adversarial_dataset_prediction:
tb_logger.attach(trainer,
log_handler=OutputHandler(
tag="training", metric_names=["loss_w/_adv"]),
event_name=Events.ITERATION_COMPLETED)
tb_logger.attach(trainer,
log_handler=OutputHandler(
tag="training",
metric_names=["loss_only_adv"]),
event_name=Events.ITERATION_COMPLETED)
tb_logger.attach(trainer,
log_handler=OptimizerParamsHandler(optimizer),
event_name=Events.ITERATION_STARTED)
tb_logger.attach(evaluator,
log_handler=OutputHandler(tag="validation",
metric_names=list(
metrics.keys()),
another_engine=trainer),
event_name=Events.EPOCH_COMPLETED)
logger.info('save checkpoints to: %s' % tb_logger.writer.log_dir)
checkpoint_handler = ModelCheckpoint(tb_logger.writer.log_dir,
'checkpoint',
save_interval=1,
n_saved=3)
trainer.add_event_handler(
Events.EPOCH_COMPLETED, checkpoint_handler,
{'mymodel': getattr(model, 'module', model)
}) # "getattr" take care of distributed encapsulation
torch.save(args, tb_logger.writer.log_dir + '/model_training_args.bin')
getattr(model, 'module', model).config.to_json_file(
os.path.join(tb_logger.writer.log_dir, CONFIG_NAME))
tokenizer.save_pretrained(tb_logger.writer.log_dir)
#logger.debug("Saving optimizer and scheduler states to %s", tb_logger.writer.log_dir)
#torch.save(optimizer.state_dict(), os.path.join(tb_logger.writer.log_dir, 'optimizer.pt'))
#torch.save(scheduler.state_dict(), os.path.join(tb_logger.writer.log_dir, 'scheduler.pt'))
# Run the training
trainer.run(train_loader, max_epochs=args.n_epochs)
# On the main process: close tensorboard logger and rename the last checkpoint (for easy re-loading with OpenAIGPTModel.from_pretrained method)
if args.local_rank in [-1, 0] and args.n_epochs > 0:
os.rename(
checkpoint_handler._saved[-1][1][-1],
os.path.join(tb_logger.writer.log_dir, WEIGHTS_NAME)
) # TODO: PR in ignite to have better access to saved file paths (cleaner)
tb_logger.close()
def train():
parser = ArgumentParser()
parser.add_argument("--train_batch_size",
type=int,
default=4,
help="Batch size for training")
parser.add_argument("--valid_batch_size",
type=int,
default=4,
help="Batch size for validation")
parser.add_argument("--gradient_accumulation_steps",
type=int,
default=8,
help="Accumulate gradients on several steps")
parser.add_argument("--lr",
type=float,
default=6.25e-5,
help="Learning rate")
parser.add_argument("--lm_coef",
type=float,
default=1.0,
help="LM loss coefficient")
parser.add_argument("--mc_coef",
type=float,
default=1.0,
help="Multiple-choice loss coefficient")
parser.add_argument("--max_norm",
type=float,
default=1.0,
help="Clipping gradient norm")
parser.add_argument("--n_epochs",
type=int,
default=3,
help="Number of training epochs")
parser.add_argument("--device",
type=str,
default="cuda" if torch.cuda.is_available() else "cpu",
help="Device (cuda or cpu)")
args = parser.parse_args()
# logging is set to INFO (resp. WARN) for main (resp. auxiliary) process. logger.info => log main process only, logger.warning => log all processes
logging.basicConfig(level=logging.INFO)
logger.info("Arguments: %s", pformat(args))
logger.info("Prepare tokenizer, pretrained model and optimizer.")
tokenizer_class = GPT2Tokenizer
tokenizer = tokenizer_class.from_pretrained("gpt2")
model_class = GPT2DoubleHeadsModel
model = model_class.from_pretrained("gpt2")
model.to(args.device)
# Add special tokens if they are not already added
add_special_tokens_(model, tokenizer) ### TODO add our own special tokens
optimizer = AdamW(model.parameters(), lr=args.lr, correct_bias=True)
logger.info("Prepare datasets")
train_loader, val_loader, train_sampler, valid_sampler = get_data_loaders(
args, tokenizer) ### TODO load data ourselves
# Training function and trainer
def update(engine, batch):
model.train()
batch = tuple(input_tensor.to(args.device) for input_tensor in batch)
input_ids, mc_token_ids, lm_labels, mc_labels, token_type_ids = batch
(lm_loss), (mc_loss), *_ = model(input_ids,
token_type_ids=token_type_ids,
mc_token_ids=mc_token_ids,
mc_labels=mc_labels,
lm_labels=lm_labels)
loss = (lm_loss * args.lm_coef +
mc_loss * args.mc_coef) / args.gradient_accumulation_steps
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_norm)
if engine.state.iteration % args.gradient_accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
return loss.item()
trainer = Engine(update)
# Evaluation function and evaluator (evaluator output is the input of the metrics)
def inference(engine, batch):
model.eval()
with torch.no_grad():
batch = tuple(
input_tensor.to(args.device) for input_tensor in batch)
input_ids, mc_token_ids, lm_labels, mc_labels, token_type_ids = batch
logger.info(tokenizer.decode(input_ids[0, -1, :].tolist()))
# if we dont send labels to model, it doesnt return losses
lm_logits, mc_logits, *_ = model(
input_ids,
token_type_ids=token_type_ids,
mc_token_ids=mc_token_ids,
)
lm_logits_flat_shifted = lm_logits[..., :-1, :].contiguous().view(
-1, lm_logits.size(-1))
lm_labels_flat_shifted = lm_labels[..., 1:].contiguous().view(-1)
return (lm_logits_flat_shifted,
mc_logits), (lm_labels_flat_shifted, mc_labels)
evaluator = Engine(inference)
# Attach evaluation to trainer: we evaluate when we start the training and at the end of each epoch
trainer.add_event_handler(Events.EPOCH_COMPLETED,
lambda _: evaluator.run(val_loader))
if args.n_epochs < 1:
trainer.add_event_handler(Events.COMPLETED,
lambda _: evaluator.run(val_loader))
trainer.add_event_handler(Events.STARTED,
lambda _: evaluator.run(val_loader))
# Linearly decrease the learning rate from lr to zero
scheduler = PiecewiseLinear(optimizer, "lr",
[(0, args.lr),
(args.n_epochs * len(train_loader), 0.0)])
trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
# Prepare metrics - note how we compute distributed metrics
RunningAverage(output_transform=lambda x: x).attach(trainer, "loss")
metrics = {
"nll":
Loss(torch.nn.CrossEntropyLoss(ignore_index=-1),
output_transform=lambda x: (x[0][0], x[1][0])),
"accuracy":
Accuracy(output_transform=lambda x: (x[0][1], x[1][1]))
}
metrics.update({
"average_nll":
MetricsLambda(average_distributed_scalar, metrics["nll"], args),
"average_accuracy":
MetricsLambda(average_distributed_scalar, metrics["accuracy"], args)
})
metrics["average_ppl"] = MetricsLambda(math.exp, metrics["average_nll"])
for name, metric in metrics.items():
metric.attach(evaluator, name)
# On the main process: add progress bar, tensorboard, checkpoints and save model, configuration and tokenizer before we start to train
pbar = ProgressBar(persist=True)
pbar.attach(trainer, metric_names=["loss"])
evaluator.add_event_handler(
Events.COMPLETED, lambda _: pbar.log_message(
"Validation: %s" % pformat(evaluator.state.metrics)))
log_dir = make_logdir("gpt2")
tb_logger = TensorboardLogger(log_dir)
tb_logger.attach(trainer,
log_handler=OutputHandler(tag="training",
metric_names=["loss"]),
event_name=Events.ITERATION_COMPLETED)
tb_logger.attach(trainer,
log_handler=OptimizerParamsHandler(optimizer),
event_name=Events.ITERATION_STARTED)
tb_logger.attach(evaluator,
log_handler=OutputHandler(tag="validation",
metric_names=list(
metrics.keys()),
another_engine=trainer),
event_name=Events.EPOCH_COMPLETED)
checkpoint_handler = ModelCheckpoint(log_dir,
'checkpoint',
save_interval=1,
n_saved=3)
trainer.add_event_handler(
Events.EPOCH_COMPLETED, checkpoint_handler,
{'mymodel': getattr(model, 'module', model)
}) # "getattr" takes care of distributed encapsulation
torch.save(args, log_dir + '/model_training_args.bin')
getattr(model, 'module',
model).config.to_json_file(os.path.join(log_dir, CONFIG_NAME))
tokenizer.save_pretrained(log_dir)
# Run the training
trainer.run(train_loader, max_epochs=args.n_epochs)
# On the main process: close tensorboard logger and rename the last checkpoint (for easy re-loading with OpenAIGPTModel.from_pretrained method)
if args.n_epochs > 0:
os.rename(
checkpoint_handler._saved[-1][1][-1],
os.path.join(log_dir, WEIGHTS_NAME)
) # TODO: PR in ignite to have better access to saved file paths (cleaner)
tb_logger.close()
def train(dataset_path,
dataset_cache='./dataset_cache',
model_checkpoint='gpt2',
num_candidates=2,
max_history=2,
train_batch_size=4,
valid_batch_size=4,
gradient_accumulation_steps=8,
lr=6.25e-5,
lm_coef=1.0,
mc_coef=1.0,
max_norm=1.0,
n_epochs=3,
personality_permutations=1,
eval_before_start=False,
device="cuda" if torch.cuda.is_available() else "cpu",
fp16='',
path_prefix='',
log_dir='',
local_rank=-1):
args = {**locals()}
# logging is set to INFO (resp. WARN) for main (resp. auxiliary) process. logger.info => log main process only, logger.warning => log all processes
logging.basicConfig(
level=logging.INFO if local_rank in [-1, 0] else logging.WARN)
# This is a logger.warning: it will be printed by all distributed processes
logger.warning("Running process %d", local_rank)
logger.info("Arguments: %s", pformat(args))
# Initialize distributed training if needed
distributed = (local_rank != -1)
args['distributed'] = distributed
if distributed:
torch.cuda.set_device(local_rank)
device = torch.device("cuda", local_rank)
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
logger.info("Prepare tokenizer, pretrained model and optimizer.")
# cant use Autotokenizer because checkpoint could be a Path
tokenizer_class = GPT2Tokenizer if "gpt2" in model_checkpoint else OpenAIGPTTokenizer
tokenizer = tokenizer_class.from_pretrained(model_checkpoint)
model_class = GPT2DoubleHeadsModel if "gpt2" in model_checkpoint else OpenAIGPTDoubleHeadsModel
model = model_class.from_pretrained(model_checkpoint)
model.to(device)
# Add special tokens if they are not already added
add_special_tokens_(model, tokenizer)
optimizer = AdamW(model.parameters(), lr=lr, correct_bias=True)
# Prepare model for FP16 and distributed training if needed (order is important, distributed should be the last)
if fp16:
from apex import amp # Apex is only required if we use fp16 training
model, optimizer = amp.initialize(model, optimizer, opt_level=fp16)
if distributed:
model = DistributedDataParallel(model,
device_ids=[local_rank],
output_device=local_rank)
logger.info("Prepare datasets")
train_loader, val_loader, train_sampler, valid_sampler = get_data_loaders(
dataset_path, dataset_cache, num_candidates, personality_permutations,
max_history, train_batch_size, valid_batch_size, distributed,
tokenizer)
# Training function and trainer
def update(engine, batch):
model.train()
batch = tuple(input_tensor.to(device) for input_tensor in batch)
input_ids, mc_token_ids, lm_labels, mc_labels, token_type_ids = batch
(lm_loss), (mc_loss), *_ = model(input_ids,
token_type_ids=token_type_ids,
mc_token_ids=mc_token_ids,
mc_labels=mc_labels,
lm_labels=lm_labels)
loss = (lm_loss * lm_coef + mc_loss * mc_coef) / \
gradient_accumulation_steps
if fp16:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer),
max_norm)
else:
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
if engine.state.iteration % gradient_accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
return loss.item()
trainer = Engine(update)
# Evaluation function and evaluator (evaluator output is the input of the metrics)
def inference(engine, batch):
model.eval()
with torch.no_grad():
batch = tuple(input_tensor.to(device) for input_tensor in batch)
input_ids, mc_token_ids, lm_labels, mc_labels, token_type_ids = batch
logger.info(tokenizer.decode(input_ids[0, -1, :].tolist()))
# if we dont send labels to model, it doesnt return losses
lm_logits, mc_logits, *_ = model(
input_ids,
token_type_ids=token_type_ids,
mc_token_ids=mc_token_ids,
)
lm_logits_flat_shifted = lm_logits[..., :-1, :].contiguous().view(
-1, lm_logits.size(-1))
lm_labels_flat_shifted = lm_labels[..., 1:].contiguous().view(-1)
return (lm_logits_flat_shifted,
mc_logits), (lm_labels_flat_shifted, mc_labels)
evaluator = Engine(inference)
# Attach evaluation to trainer: we evaluate when we start the training and at the end of each epoch
trainer.add_event_handler(Events.EPOCH_COMPLETED,
lambda _: evaluator.run(val_loader))
if n_epochs < 1:
trainer.add_event_handler(Events.COMPLETED,
lambda _: evaluator.run(val_loader))
if eval_before_start:
trainer.add_event_handler(Events.STARTED,
lambda _: evaluator.run(val_loader))
# Make sure distributed data samplers split the dataset nicely between the distributed processes
if distributed:
trainer.add_event_handler(
Events.EPOCH_STARTED,
lambda engine: train_sampler.set_epoch(engine.state.epoch))
evaluator.add_event_handler(
Events.EPOCH_STARTED,
lambda engine: valid_sampler.set_epoch(engine.state.epoch))
# Linearly decrease the learning rate from lr to zero
scheduler = PiecewiseLinear(optimizer, "lr",
[(0, lr), (n_epochs * len(train_loader), 0.0)])
trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
# Prepare metrics - note how we compute distributed metrics
RunningAverage(output_transform=lambda x: x).attach(trainer, "loss")
metrics = {
"nll":
Loss(torch.nn.CrossEntropyLoss(ignore_index=-1),
output_transform=lambda x: (x[0][0], x[1][0])),
"accuracy":
Accuracy(output_transform=lambda x: (x[0][1], x[1][1]))
}
metrics.update({
"average_nll":
MetricsLambda(average_distributed_scalar, metrics["nll"], local_rank,
device),
"average_accuracy":
MetricsLambda(average_distributed_scalar, metrics["accuracy"],
local_rank, device)
})
metrics["average_ppl"] = MetricsLambda(math.exp, metrics["average_nll"])
for name, metric in metrics.items():
metric.attach(evaluator, name)
# On the main process: add progress bar, tensorboard, checkpoints and save model, configuration and tokenizer before we start to train
if local_rank in [-1, 0]:
pbar = ProgressBar(persist=True)
pbar.attach(trainer, metric_names=["loss"])
evaluator.add_event_handler(
Events.COMPLETED, lambda _: pbar.log_message(
"Validation: %s" % pformat(evaluator.state.metrics)))
log_dir = log_dir if log_dir else make_logdir(model_checkpoint,
path=path_prefix)
tb_logger = TensorboardLogger(log_dir)
tb_logger.attach(trainer,
log_handler=OutputHandler(tag="training",
metric_names=["loss"]),
event_name=Events.ITERATION_COMPLETED)
tb_logger.attach(trainer,
log_handler=OptimizerParamsHandler(optimizer),
event_name=Events.ITERATION_STARTED)
tb_logger.attach(evaluator,
log_handler=OutputHandler(tag="validation",
metric_names=list(
metrics.keys()),
another_engine=trainer),
event_name=Events.EPOCH_COMPLETED)
checkpoint_handler = ModelCheckpoint(log_dir,
'checkpoint',
save_interval=1,
n_saved=3)
trainer.add_event_handler(
Events.EPOCH_COMPLETED, checkpoint_handler,
{'mymodel': getattr(model, 'module', model)
}) # "getattr" takes care of distributed encapsulation
torch.save(args, log_dir + '/model_training_bin')
getattr(model, 'module',
model).config.to_json_file(os.path.join(log_dir, CONFIG_NAME))
tokenizer.save_pretrained(log_dir)
# Run the training
trainer.run(train_loader, max_epochs=n_epochs)
# On the main process: close tensorboard logger and rename the last checkpoint (for easy re-loading with OpenAIGPTModel.from_pretrained method)
if local_rank in [-1, 0] and n_epochs > 0:
# TODO: PR in ignite to have better access to saved file paths (cleaner)
os.rename(checkpoint_handler._saved[-1][1][-1],
os.path.join(log_dir, WEIGHTS_NAME))
tb_logger.close()
def train():
parser = ArgumentParser()
parser.add_argument(
"--dataset_path",
type=str,
default="",
help="Path or url of the dataset. If empty download from S3.")
parser.add_argument("--dataset_cache",
type=str,
default='./dataset_cache',
help="Path or url of the dataset cache")
parser.add_argument("--model_checkpoint",
type=str,
default="gpt2",
help="Path, url or short name of the model")
parser.add_argument(
"--task",
type=str,
default="dialogue",
help="one of task from [dialogue, qa, mt, nlg, summarization]")
parser.add_argument("--emb_only",
action='store_true',
help="fine tune only task embeddings")
parser.add_argument("--linear_perturb",
action='store_true',
help="fine tune only task embeddings")
parser.add_argument("--max_history",
type=int,
default=2,
help="Number of previous exchanges to keep in history")
parser.add_argument("--train_batch_size",
type=int,
default=4,
help="Batch size for training")
parser.add_argument("--valid_batch_size",
type=int,
default=4,
help="Batch size for validation")
parser.add_argument("--gradient_accumulation_steps",
type=int,
default=8,
help="Accumulate gradients on several steps")
parser.add_argument("--lr",
type=float,
default=6.25e-5,
help="Learning rate")
parser.add_argument("--lm_coef",
type=float,
default=1.0,
help="LM loss coefficient")
parser.add_argument("--max_norm",
type=float,
default=1.0,
help="Clipping gradient norm")
parser.add_argument("--n_epochs",
type=int,
default=1,
help="Number of training epochs")
parser.add_argument("--personality_permutations",
type=int,
default=1,
help="Number of permutations of personality sentences")
parser.add_argument(
"--eval_before_start",
action='store_true',
help="If true start with a first evaluation before training")
parser.add_argument("--device",
type=str,
default="cuda" if torch.cuda.is_available() else "cpu",
help="Device (cuda or cpu)")
parser.add_argument(
"--fp16",
type=str,
default="",
help=
"Set to O0, O1, O2 or O3 for fp16 training (see apex documentation)")
parser.add_argument(
"--local_rank",
type=int,
default=-1,
help="Local rank for distributed training (-1: not distributed)")
parser.add_argument("--perturbation_layers",
type=int,
default=0,
help="number of perturbation layers")
parser.add_argument("--self_copy",
action='store_true',
help="add self copy ")
parser.add_argument("--adapter_bottleneck",
type=int,
default=0,
help="adapter layer bottleneck")
parser.add_argument("--random_init",
action='store_true',
help="don't use GPT-2 pre-trained model ")
parser.add_argument("--distillation", action='store_true')
parser.add_argument("--outputlayer_only", action='store_true')
args = parser.parse_args()
# logging is set to INFO (resp. WARN) for main (resp. auxiliary) process. logger.info => log main process only, logger.warning => log all processes
logging.basicConfig(
level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN)
logger.warning(
"Running process %d", args.local_rank
) # This is a logger.warning: it will be printed by all distributed processes
logger.info("Arguments: %s", pformat(args))
# Initialize distributed training if needed
args.distributed = (args.local_rank != -1)
if args.distributed:
torch.cuda.set_device(args.local_rank)
args.device = torch.device("cuda", args.local_rank)
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
logger.info("Prepare tokenizer, pretrained model and optimizer.")
tokenizer_class = GPT2Tokenizer if "gpt2" in args.model_checkpoint else OpenAIGPTTokenizer # cant use Autotokenizer because checkpoint could be a Path
tokenizer = tokenizer_class.from_pretrained(args.model_checkpoint)
model_class = GPT2LMHeadModel if "gpt2" in args.model_checkpoint else OpenAIGPTDoubleHeadsModel
if not args.random_init:
model = model_class.from_pretrained(
args.model_checkpoint,
perturbation_layers=args.perturbation_layers,
self_copy=args.self_copy,
adapter_bottleneck=args.adapter_bottleneck)
else:
config = GPT2Config()
model = model_class(config,
perturbation_layers=args.perturbation_layers,
self_copy=args.self_copy,
adapter_bottleneck=args.adapter_bottleneck)
model.to(args.device)
# Add special tokens if they are not already added
add_special_tokens_(model, tokenizer)
if args.adapter_bottleneck > 0:
parameters_to_update = [
p for n, p in model.named_parameters() if "adapter" in str(n)
] + [model.transformer.wte.weight]
optimizer = AdamW(parameters_to_update, lr=args.lr, correct_bias=True)
elif args.outputlayer_only:
parameters_to_update = [model.transformer.wte.weight]
optimizer = AdamW(parameters_to_update, lr=args.lr, correct_bias=True)
else:
optimizer = AdamW(model.parameters(), lr=args.lr, correct_bias=True)
# Prepare model for FP16 and distributed training if needed (order is important, distributed should be the last)
if args.fp16:
from apex import amp # Apex is only required if we use fp16 training
model, optimizer = amp.initialize(model,
optimizer,
opt_level=args.fp16)
if args.distributed:
model = DistributedDataParallel(model,
device_ids=[args.local_rank],
output_device=args.local_rank)
logger.info("Prepare datasets")
train_loader, val_loader, train_sampler, valid_sampler = get_data_loaders(
args, tokenizer)
# Training function and trainer
def update_emb(engine, batch):
model.train()
batch = tuple(input_tensor.to(args.device) for input_tensor in batch)
input_ids, lm_labels, token_type_ids = batch
(lm_loss), *_ = model(input_ids,
token_type_ids=token_type_ids,
labels=lm_labels)
loss = lm_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_norm)
else:
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_norm)
if engine.state.iteration % args.gradient_accumulation_steps == 0:
param_to_check = []
for n, p in model.named_parameters():
if (n != "transformer.wte.weight"):
param_to_check.append(p)
a = list(param_to_check)[0].clone()
model.transformer.wte.weight.grad[:50257, :] = 0
model.transformer.wte.weight.data.add_(
-args.lr, model.transformer.wte.weight.grad.data)
optimizer.zero_grad()
param_to_check = []
for n, p in model.named_parameters():
if (n != "transformer.wte.weight"):
param_to_check.append(p)
b = list(param_to_check)[0].clone()
assert torch.equal(a.data, b.data)
return loss.item()
# Training function and trainer
def update_linear_perturbation(engine, batch):
model.train()
batch = tuple(input_tensor.to(args.device) for input_tensor in batch)
input_ids, lm_labels, token_type_ids = batch
(lm_loss), *_ = model(input_ids,
token_type_ids=token_type_ids,
labels=lm_labels)
loss = lm_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_norm)
else:
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_norm)
if engine.state.iteration % args.gradient_accumulation_steps == 0:
model.transformer.wte.weight.grad[:50257, :] = 0
# model.transformer.wte.weight.data.add_(-args.lr,model.transformer.wte.weight.grad.data)
optimizer.step()
optimizer.zero_grad()
return loss.item()
# Training function and trainer
def update_all(engine, batch):
model.train()
batch = tuple(input_tensor.to(args.device) for input_tensor in batch)
input_ids, lm_labels, token_type_ids = batch
(lm_loss), *_ = model(input_ids,
token_type_ids=token_type_ids,
labels=lm_labels,
self_copy=args.self_copy)
loss = lm_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_norm)
else:
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_norm)
if engine.state.iteration % args.gradient_accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
return loss.item()
if args.emb_only:
trainer = Engine(update_emb)
elif (args.linear_perturb or args.adapter_bottleneck > 0):
trainer = Engine(update_linear_perturbation)
else:
trainer = Engine(update_all)
# Evaluation function and evaluator (evaluator output is the input of the metrics)
def inference(engine, batch):
model.eval()
with torch.no_grad():
batch = tuple(
input_tensor.to(args.device) for input_tensor in batch)
input_ids, lm_labels, token_type_ids = batch
logger.info(tokenizer.decode(input_ids[0, -1, :].tolist()))
# if we dont send labels to model, it doesnt return losses
lm_logits, *_ = model(
input_ids,
token_type_ids=token_type_ids,
)
lm_logits_flat_shifted = lm_logits[..., :-1, :].contiguous().view(
-1, lm_logits.size(-1))
lm_labels_flat_shifted = lm_labels[..., 1:].contiguous().view(-1)
return (lm_logits_flat_shifted, ), (lm_labels_flat_shifted, )
evaluator = Engine(inference)
# Attach evaluation to trainer: we evaluate when we start the training and at the end of each epoch
trainer.add_event_handler(Events.EPOCH_COMPLETED,
lambda _: evaluator.run(val_loader))
if args.n_epochs < 1:
trainer.add_event_handler(Events.COMPLETED,
lambda _: evaluator.run(val_loader))
if args.eval_before_start:
trainer.add_event_handler(Events.STARTED,
lambda _: evaluator.run(val_loader))
# Make sure distributed data samplers split the dataset nicely between the distributed processes
if args.distributed:
trainer.add_event_handler(
Events.EPOCH_STARTED,
lambda engine: train_sampler.set_epoch(engine.state.epoch))
evaluator.add_event_handler(
Events.EPOCH_STARTED,
lambda engine: valid_sampler.set_epoch(engine.state.epoch))
# Linearly decrease the learning rate from lr to zero
scheduler = PiecewiseLinear(optimizer, "lr",
[(0, args.lr),
(args.n_epochs * len(train_loader), 0.0)])
trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
# Prepare metrics - note how we compute distributed metrics
RunningAverage(output_transform=lambda x: x).attach(trainer, "loss")
metrics = {
"nll":
Loss(torch.nn.CrossEntropyLoss(ignore_index=-1),
output_transform=lambda x: (x[0][0], x[1][0]))
}
metrics.update({
"average_nll":
MetricsLambda(average_distributed_scalar, metrics["nll"], args)
})
metrics["average_ppl"] = MetricsLambda(math.exp, metrics["average_nll"])
for name, metric in metrics.items():
metric.attach(evaluator, name)
# On the main process: add progress bar, tensorboard, checkpoints and save model, configuration and tokenizer before we start to train
if args.local_rank in [-1, 0]:
pbar = ProgressBar(persist=True)
pbar.attach(trainer, metric_names=["loss"])
evaluator.add_event_handler(
Events.COMPLETED, lambda _: pbar.log_message(
"Validation: %s" % pformat(evaluator.state.metrics)))
log_dir = make_logdir(args.model_checkpoint,
task=args.task,
lr=args.lr,
layer=args.perturbation_layers,
self_copy=args.self_copy,
n_epochs=args.n_epochs,
adapter=args.adapter_bottleneck,
random_init=args.random_init)
if args.distillation:
log_dir += "_distillation"
if args.outputlayer_only:
log_dir += "_outputlayer_only"
tb_logger = TensorboardLogger(log_dir)
tb_logger.attach(trainer,
log_handler=OutputHandler(tag="training",
metric_names=["loss"]),
event_name=Events.ITERATION_COMPLETED)
tb_logger.attach(trainer,
log_handler=OptimizerParamsHandler(optimizer),
event_name=Events.ITERATION_STARTED)
tb_logger.attach(evaluator,
log_handler=OutputHandler(tag="validation",
metric_names=list(
metrics.keys()),
another_engine=trainer),
event_name=Events.EPOCH_COMPLETED)
checkpoint_handler = ModelCheckpoint(log_dir,
'checkpoint',
save_interval=1,
n_saved=3)
trainer.add_event_handler(
Events.EPOCH_COMPLETED, checkpoint_handler,
{'mymodel': getattr(model, 'module', model)
}) # "getattr" takes care of distributed encapsulation
torch.save(args, log_dir + '/model_training_args.bin')
getattr(model, 'module',
model).config.to_json_file(os.path.join(log_dir, CONFIG_NAME))
tokenizer.save_pretrained(log_dir)
# Run the training
trainer.run(train_loader, max_epochs=args.n_epochs)
# On the main process: close tensorboard logger and rename the last checkpoint (for easy re-loading with OpenAIGPTModel.from_pretrained method)
if args.local_rank in [-1, 0] and args.n_epochs > 0:
os.rename(
checkpoint_handler._saved[-1][1][-1],
os.path.join(log_dir, WEIGHTS_NAME)
) # TODO: PR in ignite to have better access to saved file paths (cleaner)
tb_logger.close()
def train():
parser = ArgumentParser()
parser.add_argument("--dataset_path", type=str, default=DATA_FOLDER, help="Path of the dataset.")
parser.add_argument("--image_path", type=str, default=IMG_FOLDER, help="Path of the images.")
parser.add_argument("--images_feature_path", type=str, default=IMG_FEATURE_FOLDER, help="Path of the images.")
parser.add_argument("--dataset_cache", type=str, default=DATA_CACHE, help="Path of the dataset cache_no_pretrained")
parser.add_argument("--model_checkpoint", type=str, default="gpt2", help="Path, url or short name of the model")
parser.add_argument('--dhead_gpt2', action='store_true', default=False, help="use double head gpt2")
parser.add_argument("--from_step", type=int, default=-1, help="Init learning rate from this step")
parser.add_argument('--pretrained', action='store_true', default=True, help="If False train from scratch")
parser.add_argument("--num_candidates", type=int, default=1, help="Number of candidates for training")
parser.add_argument("--max_history", type=int, default=3, help="Number of previous turns to keep in history")
parser.add_argument("--max_length", type=int, default=256, help="Max length of input sentence")
parser.add_argument("--train_batch_size", type=int, default=58, help="Batch size for training")
parser.add_argument("--valid_batch_size", type=int, default=32, help="Batch size for validation")
parser.add_argument("--gradient_accumulation_steps", type=int, default=9, help="Accumulate gradients on several steps")
parser.add_argument("--lr", type=float, default=6.25e-5, help="Learning rate")
parser.add_argument("--scheduler", type=str, default="linear", choices=['noam', 'linear'], help="method of optim")
parser.add_argument("--n_emd", type=int, default=768, help="Number of n_emd in config file (for noam)")
parser.add_argument("--warmup_steps", type=int, default=5000, help="Warm up steps")
parser.add_argument("--lm_coef", type=float, default=2.0, help="LM loss coefficient")
parser.add_argument("--mc_coef", type=float, default=1.0, help="Multiple-choice loss coefficient")
parser.add_argument("--max_norm", type=float, default=1.0, help="Clipping gradient norm")
parser.add_argument("--n_epochs", type=int, default=50, help="Number of training epochs")
parser.add_argument("--num_workers", type=int, default=0, help="Number of subprocesses for data loading")
parser.add_argument("--personality_permutations", type=int, default=1, help="Number of permutations of personality sentences")
parser.add_argument("--eval_before_start", action='store_true', help="If true start with a first evaluation before training")
parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu", help="Device (cuda or cpu)")
parser.add_argument("--fp16", type=str, default="O1", help="Set to O0, O1, O2 or O3 for fp16 training (see apex documentation)")
parser.add_argument("--local_rank", type=int, default=-1, help="Local rank for distributed training (-1: not distributed)")
args = parser.parse_args()
# logging is set to INFO (resp. WARN) for main (resp. auxiliary) process. logger.info => log main process only, logger.warning => log all processes
logging.basicConfig(level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN)
logger.warning("Running process %d", args.local_rank) # This is a logger.warning: it will be printed by all distributed processes
logger.info("Arguments: %s", pformat(args))
# Initialize distributed training if needed
args.distributed = (args.local_rank != -1)
if args.distributed:
torch.cuda.set_device(args.local_rank)
args.device = torch.device("cuda", args.local_rank)
torch.distributed.init_process_group(backend='nccl', init_method='env://')
logger.info("Prepare tokenizer, pretrained model and optimizer.")
tokenizer_class = BertTokenizer
config_class = GPT2Config # GPT2Config if "gpt2" in args.model_checkpoint else OpenAIGPTConfig
model_class = GPT2LMHeadModel # GPT2DoubleHeadsModel if "gpt2" in args.model_checkpoint else OpenAIGPTDoubleHeadsModel
if args.pretrained:
tokenizer = tokenizer_class.from_pretrained(MODEL_CHECKPOINT, do_lower_case=False)
# tokenizer = tokenizer_class(vocab_file=VOCAB_PATH, do_lower_case=True)
model = model_class.from_pretrained(MODEL_CHECKPOINT)
else:
tokenizer = tokenizer_class(vocab_file=VOCAB_PATH, do_lower_case=False)
tokenizer.add_special_tokens(ATTR_TO_SPECIAL_TOKEN)
config = config_class.from_json_file(CONFIG_PATH)
model = model_class(config)
model.to(args.device)
# Add special tokens if they are not already added
# add_special_tokens_(model, tokenizer)
# optimizer = AdamW(model.parameters(), lr=args.lr, correct_bias=True)
optimizer = AdamW([{'params': model.parameters(), 'initial_lr': args.lr}], lr=args.lr, correct_bias=True)
# Prepare model for FP16 and distributed training if needed (order is important, distributed should be the last)
if args.fp16:
from apex import amp # Apex is only required if we use fp16 training
model, optimizer = amp.initialize(model, optimizer, opt_level=args.fp16)
if args.distributed:
model = DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank)
logger.info("Prepare datasets")
train_loader, val_loader, train_sampler, valid_sampler = build_dataloader(args, tokenizer, logger)
def update(engine, batch):
model.train()
batch = tuple(torch.tensor(input_data).to(args.device) if idx not in [2, 3] else input_data for idx, input_data in enumerate(batch))
input_ids, token_type_ids, input_images, image_ids, lm_labels, mc_token_ids, mc_labels = batch
if args.dhead_gpt2:
(lm_loss), (mc_loss), *_ = model(input_ids,
token_type_ids=token_type_ids,
mc_token_ids=mc_token_ids,
mc_labels=mc_labels,
lm_labels=lm_labels)
loss = (lm_loss * args.lm_coef + mc_loss * args.mc_coef) / args.gradient_accumulation_steps
else:
(lm_loss), *_ = model(input_ids,
labels=lm_labels,
token_type_ids=token_type_ids,
input_images=input_images,
image_ids=image_ids)
loss = lm_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_norm)
else:
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_norm)
if engine.state.iteration % args.gradient_accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
return loss.item() #, optimizer.param_groups[0]['lr']
trainer = Engine(update)
# Evaluation function and evaluator (evaluator output is the input of the metrics)
def inference(engine, batch):
model.eval()
with torch.no_grad():
batch = tuple(input_tensor.to(args.device) for input_tensor in batch)
input_ids, mc_token_ids, lm_labels, mc_labels, token_type_ids = batch
# logger.info(tokenizer.decode(input_ids[0, -1, :].tolist()))
# if we dont send labels to model, it doesnt return losses
if args.dhead_gpt2:
lm_logits, mc_logits, *_ = model(
input_ids, token_type_ids=token_type_ids, mc_token_ids=mc_token_ids,
)
lm_logits_flat_shifted = lm_logits[..., :-1, :].contiguous().view(-1, lm_logits.size(-1))
lm_labels_flat_shifted = lm_labels[..., 1:].contiguous().view(-1)
return (lm_logits_flat_shifted, mc_logits), (lm_labels_flat_shifted, mc_labels)
else:
lm_logits, *_ = model(input_ids, token_type_ids=token_type_ids)
lm_logits_flat_shifted = lm_logits[..., :-1, :].contiguous().view(-1, lm_logits.size(-1))
lm_labels_flat_shifted = lm_labels[..., 1:].contiguous().view(-1)
return lm_logits_flat_shifted, lm_labels_flat_shifted
evaluator = Engine(inference)
# Attach evaluation to trainer: we evaluate when we start the training and at the end of each epoch
# trainer.add_event_handler(Events.EPOCH_COMPLETED, lambda _: evaluator.run(val_loader))
if args.n_epochs < 1:
trainer.add_event_handler(Events.COMPLETED, lambda _: evaluator.run(val_loader))
if args.eval_before_start:
trainer.add_event_handler(Events.STARTED, lambda _: evaluator.run(val_loader))
# Make sure distributed data samplers split the dataset nicely between the distributed processes
if args.distributed:
trainer.add_event_handler(Events.EPOCH_STARTED, lambda engine: train_sampler.set_epoch(engine.state.epoch))
evaluator.add_event_handler(Events.EPOCH_STARTED, lambda engine: valid_sampler.set_epoch(engine.state.epoch))
# Linearly decrease the learning rate from lr to zero
model_size = args.n_emd
noam_lambda = lambda step: (
model_size ** (-0.5) * min((step + 1) ** (-0.5), (step + 1) * args.warmup_steps ** (-1.5)))
noam_scheduler = LambdaLR(optimizer, lr_lambda=noam_lambda, last_epoch=args.from_step)
scheduler = LRScheduler(noam_scheduler)
if args.scheduler == "linear":
scheduler = PiecewiseLinear(optimizer, "lr", [(0, args.lr), (args.n_epochs * len(train_loader), 0.0)])
trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
# Prepare metrics - note how we compute distributed metrics
RunningAverage(output_transform=lambda x: x).attach(trainer, "loss")
metrics = {"nll": Loss(torch.nn.CrossEntropyLoss(ignore_index=-100), output_transform=lambda x: (x[0][0], x[1][0])),
"accuracy": Accuracy(output_transform=lambda x: (x[0][1], x[1][1]))}
metrics.update({"average_nll": MetricsLambda(average_distributed_scalar, metrics["nll"], args),
"average_accuracy": MetricsLambda(average_distributed_scalar, metrics["accuracy"], args)})
metrics["average_ppl"] = MetricsLambda(math.exp, metrics["average_nll"])
for name, metric in metrics.items():
metric.attach(evaluator, name)
# On the main process: add progress bar, tensorboard, checkpoints and save model, configuration and tokenizer before we start to train
if args.local_rank in [-1, 0]:
pbar = ProgressBar(persist=True)
pbar.attach(trainer, metric_names=["loss"])
evaluator.add_event_handler(Events.COMPLETED, lambda _: pbar.log_message("Validation: %s" % pformat(evaluator.state.metrics)))
log_dir = make_logdir(args.model_checkpoint)
tb_logger = TensorboardLogger(log_dir)
tb_logger.attach(trainer, log_handler=OutputHandler(tag="training", metric_names=["loss"]), event_name=Events.ITERATION_COMPLETED)
tb_logger.attach(trainer, log_handler=OptimizerParamsHandler(optimizer), event_name=Events.ITERATION_STARTED)
# tb_logger.attach(evaluator, log_handler=OutputHandler(tag="validation", metric_names=list(metrics.keys()), another_engine=trainer), event_name=Events.EPOCH_COMPLETED)
checkpoint_handler = ModelCheckpoint(log_dir, 'checkpoint', n_saved=None)
trainer.add_event_handler(Events.EPOCH_COMPLETED(every=1), checkpoint_handler, {'mymodel': getattr(model, 'module', model)}) # "getattr" takes care of distributed encapsulation
torch.save(args, log_dir + '/model_training_args.bin')
getattr(model, 'module', model).config.to_json_file(os.path.join(log_dir, CONFIG_NAME))
tokenizer.save_pretrained(log_dir)
# Run the training
trainer.run(train_loader, max_epochs=args.n_epochs)
# On the main process: close tensorboard logger and rename the last checkpoint (for easy re-loading with OpenAIGPTModel.from_pretrained method)
if args.local_rank in [-1, 0] and args.n_epochs > 0:
os.rename(os.path.join(log_dir, checkpoint_handler._saved[-1][1]), os.path.join(log_dir, WEIGHTS_NAME)) # TODO: PR in ignite to have better access to saved file paths (cleaner)
tb_logger.close()
def train():
parser = ArgumentParser()
parser.add_argument(
"--dataset_path",
type=str,
default='wikitext-2',
help="One of ('wikitext-103', 'wikitext-2') or a dict of splits paths."
)
parser.add_argument("--dataset_cache",
type=str,
default='./dataset_cache',
help="Path or url of the dataset cache")
parser.add_argument("--embed_dim",
type=int,
default=410,
help="Embeddings dim")
parser.add_argument("--hidden_dim",
type=int,
default=2100,
help="Hidden dimension")
parser.add_argument("--num_max_positions",
type=int,
default=256,
help="Max input length")
parser.add_argument("--num_heads",
type=int,
default=10,
help="Number of heads")
parser.add_argument("--num_layers",
type=int,
default=16,
help="NUmber of layers")
parser.add_argument("--dropout", type=float, default=0.1, help="Dropout")
parser.add_argument("--initializer_range",
type=float,
default=0.02,
help="Dropout")
parser.add_argument("--train_batch_size",
type=int,
default=8,
help="Batch size for training")
parser.add_argument("--valid_batch_size",
type=int,
default=8,
help="Batch size for validation")
parser.add_argument("--lr",
type=float,
default=2.5e-4,
help="Learning rate")
parser.add_argument("--max_norm",
type=float,
default=0.25,
help="Clipping gradient norm")
parser.add_argument("--weight_decay",
type=float,
default=0.0,
help="Weight decay")
parser.add_argument("--n_epochs",
type=int,
default=200,
help="Number of training epochs")
parser.add_argument("--n_warmup",
type=float,
default=1000,
help="Number of warmup iterations")
parser.add_argument("--eval_every",
type=int,
default=-1,
help="Evaluate every X steps (-1 => end of epoch)")
parser.add_argument("--gradient_accumulation_steps",
type=int,
default=1,
help="Accumulate gradient")
parser.add_argument("--device",
type=str,
default="cuda" if torch.cuda.is_available() else "cpu",
help="Device (cuda or cpu)")
parser.add_argument(
"--local_rank",
type=int,
default=-1,
help="Local rank for distributed training (-1: not distributed)")
args = parser.parse_args()
# logging is set to INFO (resp. WARN) for main (resp. auxiliary) process. logger.info => log on main process only, logger.warning => log on all processes
logging.basicConfig(
level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN)
logger.warning(
"Running process %d", args.local_rank
) # This is a logger.warning: it will be printed by all distributed processes
logger.info("Arguments: %s", pformat(
args)) # This is a logger.info: only printed on the first process
# Initialize distributed training if needed
args.distributed = (args.local_rank != -1)
if args.distributed:
torch.cuda.set_device(args.local_rank)
args.device = torch.device("cuda", args.local_rank)
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
logger.info("Prepare tokenizer, model and optimizer")
tokenizer = BertTokenizer.from_pretrained(
'bert-base-cased',
do_lower_case=False) # Let's use a pre-defined tokenizer
args.num_embeddings = len(
tokenizer.vocab
) # We need this to create the model at next line (number of embeddings to use)
model = TransformerWithLMHead(args)
model.to(args.device)
optimizer = Adam(model.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
logger.info("Model has %s parameters",
sum(p.numel() for p in model.parameters() if p.requires_grad))
# Prepare model for distributed training if needed
if args.distributed:
model = DistributedDataParallel(model,
device_ids=[args.local_rank],
output_device=args.local_rank)
logger.info("Prepare datasets")
train_loader, val_loader, train_sampler, valid_sampler, train_num_words, valid_num_words = get_data_loaders(
args, tokenizer)
# Training function and trainer
def update(engine, batch):
model.train()
batch = batch.transpose(0, 1).contiguous().to(
args.device) # to shape [seq length, batch]
logits, loss = model(batch, labels=batch)
loss = loss / args.gradient_accumulation_steps
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_norm)
if engine.state.iteration % args.gradient_accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
return loss.item()
trainer = Engine(update)
# Evaluation function and evaluator (evaluator output is the input of the metrics)
def inference(engine, batch):
model.eval()
with torch.no_grad():
batch = batch.transpose(0, 1).contiguous().to(
args.device) # to shape [seq length, batch]
logits = model(batch)
shift_logits = logits[:-1].view(-1, logits.size(-1))
shift_labels = batch[1:].view(-1)
return shift_logits, shift_labels
evaluator = Engine(inference)
# Attach evaluation to trainer: we evaluate at the end of each epoch and every 'eval_every' iterations if needed
trainer.add_event_handler(Events.EPOCH_COMPLETED,
lambda _: evaluator.run(val_loader))
if args.eval_every > 0:
trainer.add_event_handler(
Events.ITERATION_COMPLETED,
lambda engine: evaluator.run(val_loader)
if engine.state.iteration % args.eval_every == 0 else None)
if args.n_epochs < 1:
trainer.add_event_handler(Events.COMPLETED,
lambda _: evaluator.run(val_loader))
# Make sure distributed data samplers split the dataset nicely between the distributed processes
if args.distributed:
trainer.add_event_handler(
Events.EPOCH_STARTED,
lambda engine: train_sampler.set_epoch(engine.state.epoch))
evaluator.add_event_handler(
Events.EPOCH_STARTED,
lambda engine: valid_sampler.set_epoch(engine.state.epoch))
# Learning rate schedule: linearly warm-up to lr and then decrease the learning rate to zero with cosine schedule
cos_scheduler = CosineAnnealingScheduler(optimizer, 'lr', args.lr, 0.0,
len(train_loader) * args.n_epochs)
scheduler = create_lr_scheduler_with_warmup(cos_scheduler, 0.0, args.lr,
args.n_warmup)
trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
# Prepare metrics - note how we average distributed metrics using average_distributed_scalar
RunningAverage(output_transform=lambda x: x).attach(trainer, "loss")
metrics = {"nll": Loss(torch.nn.CrossEntropyLoss(ignore_index=-1))}
metrics.update({
"average_nll":
MetricsLambda(average_distributed_scalar, metrics["nll"], args)
})
metrics["average_ppl"] = MetricsLambda(math.exp, metrics["average_nll"])
metrics["average_word_ppl"] = MetricsLambda(
lambda x: math.exp(x * val_loader.dataset.numel() / valid_num_words),
metrics["average_nll"])
for name, metric in metrics.items():
metric.attach(evaluator, name)
# On the main process: add progress bar, tensorboard, checkpoints and save model and configuration before we start to train
if args.local_rank in [-1, 0]:
pbar = ProgressBar(persist=True)
pbar.attach(trainer, metric_names=["loss"])
evaluator.add_event_handler(
Events.COMPLETED, lambda _: pbar.log_message(
"Validation: %s" % pformat(evaluator.state.metrics)))
tb_logger = TensorboardLogger(log_dir=None)
tb_logger.attach(trainer,
log_handler=OutputHandler(tag="training",
metric_names=["loss"]),
event_name=Events.ITERATION_COMPLETED)
tb_logger.attach(trainer,
log_handler=OptimizerParamsHandler(optimizer),
event_name=Events.ITERATION_STARTED)
@evaluator.on(Events.COMPLETED) # Log evaluator metrics on tensorboard
def tb_log_metrics(engine):
for name in metrics.keys():
tb_logger.writer.add_scalar(name, engine.state.metrics[name],
trainer.state.iteration)
checkpoint_handler = ModelCheckpoint(tb_logger.writer.log_dir,
'checkpoint',
save_interval=1,
n_saved=3)
trainer.add_event_handler(
Events.EPOCH_COMPLETED, checkpoint_handler,
{'mymodel': getattr(model, 'module', model)
}) # "getattr" take care of distributed encapsulation
torch.save(args, os.path.join(tb_logger.writer.log_dir, CONFIG_NAME))
# Run the training
trainer.run(train_loader, max_epochs=args.n_epochs)
# On the main process: close tensorboard logger and rename the last checkpoint for easy re-loading
if args.local_rank in [-1, 0] and args.n_epochs > 0:
os.rename(
checkpoint_handler._saved[-1][1][-1],
os.path.join(tb_logger.writer.log_dir, WEIGHTS_NAME)
) # TODO: PR in ignite to have better access to saved file paths (cleaner)
tb_logger.close()
def train():
parser = ArgumentParser()
parser.add_argument("--dataset_path",
type=str,
default="",
help="Path or url of the dataset.")
parser.add_argument("--use_adapter",
default=False,
action='store_true',
help="Use adapter or not")
parser.add_argument("--keyword_Module",
type=str,
default="",
help="add, attention, ")
parser.add_argument("--model_checkpoint",
type=str,
default="bertGpt",
help="Path, url or short name of the model")
parser.add_argument("--train_batch_size",
type=int,
default=8,
help="Batch size for training")
parser.add_argument("--valid_batch_size",
type=int,
default=8,
help="Batch size for validation")
parser.add_argument("--gradient_accumulation_steps",
type=int,
default=8,
help="Accumulate gradients on several steps")
parser.add_argument("--lr",
type=float,
default=6.25e-5,
help="Learning rate")
parser.add_argument("--max_norm",
type=float,
default=1.0,
help="Clipping gradient norm")
parser.add_argument("--n_epochs",
type=int,
default=3,
help="Number of training epochs")
parser.add_argument(
"--eval_before_start",
action='store_true',
help="If true start with a first evaluation before training")
parser.add_argument("--device",
type=str,
default="cuda" if torch.cuda.is_available() else "cpu",
help="Device (cuda or cpu)")
parser.add_argument(
"--fp16",
type=str,
default="",
help=
"Set to O0, O1, O2 or O3 for fp16 training (see apex documentation)")
parser.add_argument(
"--local_rank",
type=int,
default=-1,
help="Local rank for distributed training (-1: not distributed)")
parser.add_argument("--bert_model_path",
default="./",
type=str,
help="Bert pre-trained model path")
parser.add_argument(
"--vocab_file",
default="./vocab.korean.rawtext.list",
type=str,
help="The vocabulary file that the BERT model was trained on.")
parser.add_argument(
"--do_lower_case",
action='store_true',
help="Set this flag if you are using an uncased model.")
parser.add_argument('--seed',
type=int,
default=42,
help="random seed for initialization")
args = parser.parse_args()
# logging is set to INFO (resp. WARN) for main (resp. auxiliary) process. logger.info => log main process only, logger.warning => log all processes
logging.basicConfig(
format='%(asctime)s - %(levelname)s - %(name)s - %(message)s',
datefmt='%m/%d/%Y %H:%M:%S',
level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN)
logger.warning(
"Running process %d", args.local_rank
) # This is a logger.warning: it will be printed by all distributed processes
logger.info("Arguments: %s", pformat(args))
# Initialize distributed training if needed
args.distributed = (args.local_rank != -1)
if args.distributed:
torch.cuda.set_device(args.local_rank)
args.device = torch.device("cuda", args.local_rank)
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
logger.info("Prepare tokenizer, pretrained model and optimizer.")
#tokenizer_class = GPT2Tokenizer if "gpt2" in args.model_checkpoint else OpenAIGPTTokenizer # cant use Autotokenizer because checkpoint could be a Path
#tokenizer = tokenizer_class.from_pretrained(args.model_checkpoint)
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
# Load KoBERT model and tokenizer
bert_tokenizer = BertTokenizer.from_pretrained(
args.vocab_file, do_lower_case=args.do_lower_case)
bert_model = BertModel.from_pretrained(args.bert_model_path)
bert_model.to(args.device)
# Load KoGPT2 model and tokenizer
tok_path = get_tokenizer()
gpt_model, gpt_vocab = get_pytorch_conkogpt2_model2(
keyword_Module=args.keyword_Module, use_adapter=args.use_adapter)
gpt_tokenizer = SentencepieceTokenizer(tok_path)
gpt_model.to(args.device)
model = Seq2Seq(bert_model, gpt_model, gpt_vocab, args)
optimizer = AdamW(model.parameters(), lr=args.lr, correct_bias=True)
# Prepare model for FP16 and distributed training if needed (order is important, distributed should be the last)
#if args.fp16:
#from apex import amp # Apex is only required if we use fp16 training
#model, optimizer = amp.initialize(model, optimizer, opt_level=args.fp16)
if args.distributed:
model = DistributedDataParallel(model,
device_ids=[args.local_rank],
output_device=args.local_rank)
logger.info("Prepare datasets")
train_loader, val_loader, train_sampler, valid_sampler = get_data_loaders(
args, bert_tokenizer, gpt_tokenizer, gpt_vocab)
# Training function and trainer
def update(engine, batch):
model.train()
batch = tuple(input_tensor.to(args.device) for input_tensor in batch)
source_ids, target_ids, lm_labels = batch
#(lm_loss), *_ = model(input_ids, token_type_ids=token_type_ids, labels=lm_labels)
(lm_loss), *_ = model(source_ids, target_ids, lm_labels=lm_labels)
loss = lm_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_norm)
else:
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_norm)
if engine.state.iteration % args.gradient_accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
return loss.item()
trainer = Engine(update)
# Evaluation function and evaluator (evaluator output is the input of the metrics)
def inference(engine, batch):
model.eval()
with torch.no_grad():
batch = tuple(
input_tensor.to(args.device) for input_tensor in batch)
source_ids, target_ids, lm_labels = batch
#lm_logits, *_ = model(input_ids, token_type_ids=token_type_ids,)
lm_logits, *_ = model(source_ids, target_ids)
lm_logits_flat_shifted = lm_logits[..., :-1, :].contiguous().view(
-1, lm_logits.size(-1))
lm_labels_flat_shifted = lm_labels[..., 1:].contiguous().view(-1)
return (lm_logits_flat_shifted), (lm_labels_flat_shifted)
evaluator = Engine(inference)
# Attach evaluation to trainer: we evaluate when we start the training and at the end of each epoch
trainer.add_event_handler(Events.EPOCH_COMPLETED,
lambda _: evaluator.run(val_loader))
if args.n_epochs < 1:
trainer.add_event_handler(Events.COMPLETED,
lambda _: evaluator.run(val_loader))
if args.eval_before_start:
trainer.add_event_handler(Events.STARTED,
lambda _: evaluator.run(val_loader))
# Make sure distributed data samplers split the dataset nicely between the distributed processes
if args.distributed:
trainer.add_event_handler(
Events.EPOCH_STARTED,
lambda engine: train_sampler.set_epoch(engine.state.epoch))
evaluator.add_event_handler(
Events.EPOCH_STARTED,
lambda engine: valid_sampler.set_epoch(engine.state.epoch))
# Linearly decrease the learning rate from lr to zero
scheduler = PiecewiseLinear(optimizer, "lr",
[(0, args.lr),
(args.n_epochs * len(train_loader), 0.0)])
trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
# Prepare metrics - note how we compute distributed metrics
RunningAverage(output_transform=lambda x: x).attach(trainer, "loss")
metrics = {
"nll":
Loss(torch.nn.CrossEntropyLoss(ignore_index=-100),
output_transform=lambda x: (x[0], x[1]))
}
metrics.update({
"average_nll":
MetricsLambda(average_distributed_scalar, metrics["nll"], args)
})
metrics["average_ppl"] = MetricsLambda(math.exp, metrics["average_nll"])
for name, metric in metrics.items():
metric.attach(evaluator, name)
# On the main process: add progress bar, tensorboard, checkpoints and save model, configuration and tokenizer before we start to train
if args.local_rank in [-1, 0]:
pbar = ProgressBar(persist=True)
pbar.attach(trainer, metric_names=["loss"])
evaluator.add_event_handler(
Events.COMPLETED, lambda _: pbar.log_message(
"Validation: %s" % pformat(evaluator.state.metrics)))
log_dir = make_logdir(args.model_checkpoint, args.dataset_path,
args.keyword_Module)
tb_logger = TensorboardLogger(log_dir)
tb_logger.attach(trainer,
log_handler=OutputHandler(tag="training",
metric_names=["loss"]),
event_name=Events.ITERATION_COMPLETED)
tb_logger.attach(trainer,
log_handler=OptimizerParamsHandler(optimizer),
event_name=Events.ITERATION_STARTED)
tb_logger.attach(evaluator,
log_handler=OutputHandler(tag="validation",
metric_names=list(
metrics.keys()),
another_engine=trainer),
event_name=Events.EPOCH_COMPLETED)
checkpoint_handler = ModelCheckpoint(log_dir,
'checkpoint',
save_interval=1,
n_saved=2)
trainer.add_event_handler(
Events.EPOCH_COMPLETED, checkpoint_handler,
{'mymodel': model
}) # "getattr" takes care of distributed encapsulation
torch.save(args, log_dir + '/model_training_args.bin')
#getattr(model, 'module', model).config.to_json_file(os.path.join(log_dir, CONFIG_NAME))
#tokenizer.save_pretrained(log_dir)
# Run the training
trainer.run(train_loader, max_epochs=args.n_epochs)
# On the main process: close tensorboard logger and rename the last checkpoint (for easy re-loading with OpenAIGPTModel.from_pretrained method)
if args.local_rank in [-1, 0] and args.n_epochs > 0:
os.rename(
os.path.join(log_dir, checkpoint_handler._saved[-1][1]),
os.path.join(log_dir, WEIGHTS_NAME)
) # TODO: PR in ignite to have better access to saved file paths (cleaner)
tb_logger.close()
def train():
parser = ArgumentParser()
parser.add_argument(
"--dataset_path", type=str, default="", help="Path or url of the dataset. If empty download from S3."
)
parser.add_argument(
"--logdir", type=str, default=None, help="If provided, the model will be output to this folder."
)
parser.add_argument("--dataset_cache", type=str, default="./dataset_cache", help="Path or url of the dataset cache")
parser.add_argument("--use_mlflow", action="store_true", help="If true we enable mlflow")
parser.add_argument("--lm_coef", type=float, default=1.0, help="LM loss coefficient")
parser.add_argument("--mc_coef", type=float, default=1.0, help="Multiple-choice loss coefficient")
parser.add_argument(
"--tracking_uri", type=str, default="http://localhost:5000", help="url for mlflow tracking server"
)
parser.add_argument("--num_candidates", type=int, default=5, help="Number of candidates for training")
parser.add_argument("--experiment", type=str, help="experiment name for mlflow")
parser.add_argument("--task_config", type=str, help="Path to the tokenization config file")
parser.add_argument("--special_tokens_file", type=str, default=None, help="Path to the special tokens file")
parser.add_argument(
"--model_checkpoint", type=str, default="distilgpt2", help="Path, url or short name of the model"
)
parser.add_argument("--model_type", type=str, default=None, help="gpt or gpt2")
parser.add_argument("--train_batch_size", type=int, default=4, help="Batch size for training")
parser.add_argument("--valid_batch_size", type=int, default=1, help="Batch size for validation")
parser.add_argument(
"--gradient_accumulation_steps", type=int, default=8, help="Accumulate gradients on several steps"
)
parser.add_argument("--lr", type=float, default=1e-4, help="Learning rate")
parser.add_argument("--adam_epsilon", type=float, default=1e-6, help="Learning rate")
parser.add_argument("--max_norm", type=float, default=1.0, help="Clipping gradient norm")
parser.add_argument("--patience", type=int, default=1, help="patience parameter for early stopping")
parser.add_argument("--n_epochs", type=int, default=10, help="Number of training epochs")
parser.add_argument("--max_data", type=int, default=0, help="Number of data items (0 includes everything)")
parser.add_argument(
"--val_max_data", type=int, default=0, help="Number of validation data items (0 includes everything)"
)
parser.add_argument(
"--eval_before_start", action="store_true", help="If true start with a first evaluation before training"
)
parser.add_argument(
"--overwrite_output_dir",
action="store_true",
help="If true, and the logdir is explictly passed, it will be overwritten",
)
parser.add_argument("--ul", action="store_true", help="If true use unlikelihood sampling")
parser.add_argument("--freeze", action="store_true", help="If true freeze layers")
parser.add_argument("--smoothing", type=float, default=0.0, help="label smoothing epsilon")
parser.add_argument("--ignore_cache", action="store_true", help="If true ignore the dataset cache")
parser.add_argument(
"--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu", help="Device (cuda or cpu)"
)
parser.add_argument(
"--fp16", type=str, default="", help="Set to O0, O1, O2 or O3 for fp16 training (see apex documentation)"
)
parser.add_argument(
"--local_rank", type=int, default=-1, help="Local rank for distributed training (-1: not distributed)"
)
parser.add_argument("--warmup-steps", default=0, type=int, help="Linear warmup over warmup_steps.")
# custom training
parser.add_argument("--sequence-tune-rate", type=float, default=0.5)
parser.add_argument("--sequence-ngram-n", type=int, default=4)
parser.add_argument(
"--multitask", action="store_true", help="If true use multitask training with multiple choice loss"
)
parser.add_argument(
"--retrain_base",
type=str,
default=None,
help="JSON file with training parameters or MLflow run_id from which to get the parameters for retraining",
)
parser.add_argument(
"--training_args_file",
type=str,
default=None,
help="File with the training arguments generated by a previous run to use as parameters",
)
parser.add_argument("--scheduler", type=str, default="piecewiselinear", help="scheduler choice")
parser.add_argument("--optimizer", type=str, default="AdamW", help="optimizer choice")
parser.add_argument(
"--max_block_size", type=int, default=None, help="If set, data is truncated to fit this max size"
)
args = parser.parse_args()
if args.retrain_base:
try:
logger.info(f"reading the arguments from {args.retrain_base}")
model_training_args = json.load(open(args.retrain_base))
except:
model_training_args = load_training_args(args.retrain_base)
passed_args = [x[2:] for x in sys.argv if x.startswith("--")]
# this is set by pytorch
passed_args.extend(["ignore_cache", "local_rank"])
for key, value in model_training_args.items():
# we only update an argument if it's not passed explicitly
if key not in passed_args:
if value:
args.__setattr__(key, value)
logger.info(vars(args))
if args.logdir is None:
args.logdir = Path(f"runs/{get_curr_time()}")
else:
args.logdir = Path(args.logdir)
if not is_empty_or_absent_dir(args.logdir) and not args.overwrite_output_dir:
logger.error(f"Error: {args.logdir} is not empty and you did not pass --overwrite_output_dir as True")
exit()
else:
if args.local_rank in [-1, 0]:
logger.info(f"deleting the existing folder {args.logdir}")
try:
rmtree(args.logdir)
except:
pass
logger.info(f"outputting model to {args.logdir}")
try:
def finalize():
if args.local_rank not in [-1, 0,]:
# Make sure only the first process in distributed training will download model & vocab
torch.distributed.barrier()
if args.local_rank in [-1, 0] and args.n_epochs > 0:
try:
# On the main process: rename the last checkpoint
# (for easy re-loading with from_pretrained method)
os.rename(checkpoint_handler._saved[-1][1][-1], os.path.join(args.logdir, WEIGHTS_NAME))
if args.use_mlflow:
mlflow.log_artifact(args.logdir / WEIGHTS_NAME, "training")
logger.info("ending mlflow run")
logger.info(f"run_id: {mlflow.active_run().info.run_id}")
mlflow.end_run()
rmtree(args.logdir)
except:
logger.info("No checkpoint to finalize the model. Deleting run")
# TODO: fix issue in mlflow trying to delete the experiment multiple times
mlflow.delete_run(mlflow.active_run().info.run_id)
rmtree(args.logdir)
if args.local_rank == 0:
torch.distributed.barrier()
args.logdir.mkdir(parents=True, exist_ok=True)
TRAINING_ARGS_FILE = args.logdir / "model_training_args.json"
args_dict = deepcopy(vars(args))
args_dict["logdir"] = str(args_dict["logdir"])
json.dump(args_dict, open(TRAINING_ARGS_FILE, "w"), indent=2)
if args.use_mlflow:
if args.local_rank in [-1, 0]:
assert args.tracking_uri
assert args.experiment
mlflow.set_tracking_uri(args.tracking_uri)
mlflow.set_experiment(args.experiment)
mlflow.start_run()
# Log parameters
mlflow.log_params(vars(args))
# Log training arguments into a file
mlflow.log_artifact(TRAINING_ARGS_FILE, "training")
# The validation maximum number of items shouldn't be more than the training (used during debugging)
if args.val_max_data == 0 and args.max_data > 0:
args.val_max_data = args.max_data
# Logging is set to INFO (resp. WARN) for main (resp. auxiliary)
# process. logger.info => log main process only, logger.warning => log all processes
logging.basicConfig(level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN)
# This is a logger.warning: it will be printed by all distributed processes
logger.warning("Running process %d", args.local_rank)
# Initialize distributed training if needed
args.distributed = args.local_rank != -1
if args.distributed:
torch.cuda.set_device(args.local_rank)
args.device = torch.device("cuda", args.local_rank)
torch.distributed.init_process_group(backend="nccl", init_method="env://")
logger.info(f"Reading the task configuration: {args.task_config}")
copyfile(args.task_config, args.logdir / "task_config.json")
task_config = load_task_config(args.task_config)
logger.info("Prepare tokenizer, pretrained model and optimizer - add special tokens for fine-tuning")
model_directory, is_local = get_model_directory(args.model_checkpoint)
model, tokenizer = load_pretrained(
model_directory,
model_type=args.model_type,
smoothing=args.smoothing,
multitask=args.multitask,
special_tokens_file=args.special_tokens_file,
task_config=task_config,
dataset_path=args.dataset_path,
)
special_tokens = read_special_tokens(
task_config=task_config,
special_tokens_file=args.special_tokens_file,
dataset_path=args.dataset_path
)
logger.info(f"adding {len(special_tokens)}")
tokenizer.add_tokens(special_tokens)
model.resize_token_embeddings(len(tokenizer))
model.to(args.device)
if args.freeze:
transformer = list(model.children())[0]
i = 0
for param in transformer.parameters():
param.requires_grad = False
i += 1
if i >= len(list(transformer.parameters())) // 2:
break
if args.optimizer.lower() == "rmsprop":
optimizer = RMSprop(model.parameters(), lr=args.lr)
elif args.optimizer.lower() == "adam":
optimizer = Adam(model.parameters(), lr=args.lr)
elif args.optimizer.lower() == "adafactor":
optimizer = Adafactor(model.parameters(), lr=args.lr, warmup_init=False)
elif args.optimizer.lower() == "sgd":
optimizer = SGD(model.parameters(), lr=args.lr)
elif args.optimizer.lower() == "novograd":
optimizer = Novograd(model.parameters(), lr=args.lr)
else:
optimizer = AdamW(model.parameters(), lr=args.lr)
# Prepare model for FP16 and distributed training if needed (order is important, distributed should be the last)
if args.fp16:
from apex import amp # Apex is only required if we use fp16 training
model, optimizer = amp.initialize(model, optimizer, opt_level=args.fp16)
if args.distributed:
model = DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank)
logger.info("Prepare datasets")
train_loader, val_loader, train_sampler, valid_sampler = get_data_loaders(args, task_config, tokenizer)
def named_batch(batch, with_labels=True):
"""Helper function so that we get a dictionary with key as the input name and the value as the input value.
This makes it easier to pass parameters to the model by their name, without caring about the order
"""
named_batch = {}
# The components in the batch are ordered as in MODEL_INPUTS
i = 0
for input_name in MODEL_INPUTS:
if not with_labels and "labels" in input_name:
continue
key = input_name
if not args.multitask:
if "mc_" in input_name:
continue
# the field is called `lm_labels` in the DoubleHeads and `labels` in single head model
if input_name == "lm_labels":
key = "labels"
named_batch[key] = batch[i]
i += 1
return named_batch
# Training function and trainer
def update(engine, batch):
model.train()
n_batch = named_batch(tuple(input_tensor.to(args.device) for input_tensor in batch))
outputs = model(**n_batch)
lm_loss = outputs[0]
if args.multitask:
mc_loss = outputs[1]
else:
mc_loss = 0
loss = (lm_loss * args.lm_coef + mc_loss * args.mc_coef) / 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_norm)
else:
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_norm)
if engine.state.iteration % args.gradient_accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
return loss.item()
trainer = Engine(update)
# Evaluation function and evaluator (evaluator output is the input of the metrics)
def inference(engine, batch):
model.eval()
with torch.no_grad():
n_batch = named_batch(tuple(input_tensor.to(args.device) for input_tensor in batch))
outputs = model(**{key: n_batch[key] for key in n_batch if "labels" not in key})
lm_logits = outputs[0]
lm_labels = n_batch["lm_labels"] if args.multitask else n_batch["labels"]
lm_logits_flat_shifted = lm_logits[..., :-1, :].contiguous().view(-1, lm_logits.size(-1))
lm_labels_flat_shifted = lm_labels[..., 1:].contiguous().view(-1)
if args.multitask:
mc_logits = outputs[1]
mc_labels = n_batch["mc_labels"]
return (lm_logits_flat_shifted, mc_logits), (lm_labels_flat_shifted, mc_labels)
else:
return lm_logits_flat_shifted, lm_labels_flat_shifted
evaluator = Engine(inference)
def checkpointing_score_function(engine):
""""""
val_metric = engine.state.metrics["average_ppl"]
logger.info(val_metric)
return -val_metric
def score_function(engine):
""""""
val_ppl = engine.state.metrics["average_ppl"]
return -val_ppl
# Attach evaluation to trainer: we evaluate when we start the training and at the end of each epoch
trainer.add_event_handler(Events.EPOCH_COMPLETED, lambda _: evaluator.run(val_loader))
if args.n_epochs < 1:
trainer.add_event_handler(Events.COMPLETED, lambda _: evaluator.run(val_loader))
if args.eval_before_start:
trainer.add_event_handler(Events.STARTED, lambda _: evaluator.run(val_loader))
# Attach mlflow logger
# trainer.add_event_handler(Events.EPOCH_COMPLETED, lambda _: evaluator.run(val_loader))
# Make sure distributed data samplers split the dataset nicely between the distributed processes
if args.distributed:
trainer.add_event_handler(Events.EPOCH_STARTED, lambda engine: train_sampler.set_epoch(engine.state.epoch))
evaluator.add_event_handler(
Events.EPOCH_STARTED, lambda engine: valid_sampler.set_epoch(engine.state.epoch)
)
if args.scheduler.lower() == "piecewiselinear":
# Linearly decrease the learning rate from lr to zero
scheduler = PiecewiseLinear(optimizer, "lr", [(0, args.lr), (args.n_epochs * len(train_loader), 0.0)])
elif args.scheduler.lower() == "linearcyclical":
scheduler = LinearCyclicalScheduler(optimizer, "lr", args.lr / 10, args.lr, len(train_loader))
elif args.scheduler.lower() == "cosine":
scheduler = CosineAnnealingLR(optimizer, args.n_epochs * len(train_loader), 1e-4)
elif args.warmup_steps > 0:
t_total = len(train_loader) // args.gradient_accumulation_steps * args.n_epochs
scheduler = get_linear_schedule_with_warmup(optimizer, args.warmup_steps, t_total)
trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
# Prepare metrics - note how we compute distributed metrics
RunningAverage(output_transform=lambda x: x).attach(trainer, "loss")
if args.multitask:
metrics = {
"nll": Loss(torch.nn.CrossEntropyLoss(ignore_index=-1), output_transform=lambda x: (x[0][0], x[1][0])),
"accuracy": Accuracy(output_transform=lambda x: (x[0][1], x[1][1])),
}
metrics.update(
{
"average_nll": MetricsLambda(average_distributed_scalar, metrics["nll"], args),
"average_accuracy": MetricsLambda(average_distributed_scalar, metrics["accuracy"], args),
}
)
else:
metrics = {"nll": Loss(torch.nn.CrossEntropyLoss(ignore_index=-1, reduction="mean"))}
metrics.update({"average_nll": MetricsLambda(average_distributed_scalar, metrics["nll"], args)})
metrics["average_ppl"] = MetricsLambda(math.exp, metrics["average_nll"])
for name, metric in metrics.items():
metric.attach(evaluator, name)
# On the main process: add progress bar, tensorboard, checkpoints and save model,
# configuration and tokenizer before we start to train
if args.local_rank in [-1, 0]:
pbar = ProgressBar(persist=True)
pbar.attach(trainer, metric_names=["loss"])
evaluator.add_event_handler(
Events.COMPLETED, lambda _: pbar.log_message("Validation: %s" % pformat(evaluator.state.metrics))
)
checkpoint_handler = ModelCheckpoint(
args.logdir,
filename_prefix="checkpoint",
score_function=checkpointing_score_function,
create_dir=True,
n_saved=2,
)
evaluator.add_event_handler(
Events.COMPLETED, checkpoint_handler, {"mymodel": getattr(model, "module", model)}
) # "getattr" takes care of distributed encapsulation
getattr(model, "module", model).config.to_json_file(os.path.join(args.logdir, CONFIG_NAME))
tokenizer.save_pretrained(args.logdir)
early_handler = EarlyStopping(patience=args.patience, score_function=score_function, trainer=trainer)
evaluator.add_event_handler(Events.COMPLETED, early_handler)
if args.use_mlflow and args.local_rank in [-1, 0]:
class MLflowTracker:
def __init__(self):
self.iteration = 1
def eval_metric_logger(self, engine):
mlflow.log_metric("last_epoch", self.iteration)
for metric in engine.state.metrics:
mlflow.log_metric(f"eval_{metric}", engine.state.metrics[metric], step=self.iteration)
self.iteration += 1
def train_metric_logger(self, engine):
for metric in engine.state.metrics:
mlflow.log_metric(f"train_{metric}", engine.state.metrics[metric], step=engine.state.epoch)
def finish_experiment(self, engine):
mlflow.log_metric("finished", True)
def start_experiment(self, engine):
# log the initial artifacts in the dir
mlflow.log_artifacts(args.logdir, "training")
mlflow.log_metric("finished", False)
mlflow_tracker = MLflowTracker()
trainer.add_event_handler(Events.STARTED, mlflow_tracker.start_experiment)
# Log the train and validation metrics
trainer.add_event_handler(Events.EPOCH_COMPLETED, mlflow_tracker.train_metric_logger)
evaluator.add_event_handler(Events.COMPLETED, mlflow_tracker.eval_metric_logger)
# Log the model
trainer.add_event_handler(Events.COMPLETED, mlflow_tracker.finish_experiment)
# Run the training
trainer.run(train_loader, max_epochs=args.n_epochs)
except KeyboardInterrupt:
finalize()
logger.info("training about to finish")
finalize()
logger.info("finalized training")
def train():
config_file = "configs/train_daily_dialog_emotion_action_config.json"
config = Config.from_json_file(config_file)
# logging is set to INFO (resp. WARN) for main (resp. auxiliary) process. logger.info => log main process only, logger.warning => log all processes
logging.basicConfig(
level=logging.INFO if config.local_rank in [-1, 0] else logging.WARN)
logger.warning(
"Running process %d", config.local_rank
) # This is a logger.warning: it will be printed by all distributed processes
logger.info("Arguments: %s", pformat(config))
# Initialize distributed training if needed
config.distributed = (config.local_rank != -1)
if config.distributed:
torch.cuda.set_device(config.local_rank)
config.device = torch.device("cuda", config.local_rank)
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
logger.info(
"Prepare tokenizer, pretrained model and optimizer - add special tokens for fine-tuning"
)
tokenizer_class = GPT2Tokenizer if "gpt2" in config.model_checkpoint else OpenAIGPTTokenizer
tokenizer = tokenizer_class.from_pretrained(config.model_checkpoint)
model_class = GPT2DoubleHeadsModel if "gpt2" in config.model_checkpoint else OpenAIGPTDoubleHeadsModel
model = model_class.from_pretrained(config.model_checkpoint)
tokenizer.set_special_tokens(SPECIAL_TOKENS)
model.set_num_special_tokens(len(SPECIAL_TOKENS))
model.to(config.device)
optimizer = OpenAIAdam(model.parameters(), lr=config.lr)
# Prepare model for FP16 and distributed training if needed (order is important, distributed should be the last)
if config.fp16:
from apex import amp # Apex is only required if we use fp16 training
model, optimizer = amp.initialize(model,
optimizer,
opt_level=config.fp16)
if config.distributed:
model = DistributedDataParallel(model,
device_ids=[config.local_rank],
output_device=config.local_rank)
logger.info("Prepare datasets")
train_loader, val_loader, train_sampler, valid_sampler = get_data_loaders(
config, tokenizer)
# Training function and trainer
def update(engine, batch):
model.train()
input_ids, mc_token_ids, lm_labels, mc_labels, token_type_ids, token_emotion_ids, token_action_ids = tuple(
input_tensor.to(config.device) for input_tensor in batch)
lm_loss, mc_loss = model(input_ids, mc_token_ids, lm_labels, mc_labels,
token_type_ids, token_emotion_ids,
token_action_ids)
loss = (lm_loss * config.lm_coef +
mc_loss * config.mc_coef) / config.gradient_accumulation_steps
if config.fp16:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer),
config.max_norm)
else:
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), config.max_norm)
if engine.state.iteration % config.gradient_accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
return loss.item()
trainer = Engine(update)
# Evaluation function and evaluator (evaluator output is the input of the metrics)
def inference(engine, batch):
model.eval()
with torch.no_grad():
batch = tuple(
input_tensor.to(config.device) for input_tensor in batch)
input_ids, mc_token_ids, lm_labels, mc_labels, token_type_ids, token_emotion_ids, token_action_ids = batch
#logger.info(tokenizer.decode(input_ids[0, -1, :].tolist()))
model_outputs = model(input_ids,
mc_token_ids,
token_type_ids=token_type_ids,
token_emotion_ids=token_emotion_ids,
token_action_ids=token_action_ids)
lm_logits, mc_logits = model_outputs[0], model_outputs[
1] # So we can also use GPT2 outputs
lm_logits_flat_shifted = lm_logits[..., :-1, :].contiguous().view(
-1, lm_logits.size(-1))
lm_labels_flat_shifted = lm_labels[..., 1:].contiguous().view(-1)
return (lm_logits_flat_shifted,
mc_logits), (lm_labels_flat_shifted, mc_labels)
evaluator = Engine(inference)
# Attach evaluation to trainer: we evaluate when we start the training and at the end of each epoch
trainer.add_event_handler(Events.EPOCH_COMPLETED,
lambda _: evaluator.run(val_loader))
if config.n_epochs < 1:
trainer.add_event_handler(Events.COMPLETED,
lambda _: evaluator.run(val_loader))
if config.eval_before_start:
trainer.add_event_handler(Events.STARTED,
lambda _: evaluator.run(val_loader))
# Make sure distributed data samplers split the dataset nicely between the distributed processes
if config.distributed:
trainer.add_event_handler(
Events.EPOCH_STARTED,
lambda engine: train_sampler.set_epoch(engine.state.epoch))
evaluator.add_event_handler(
Events.EPOCH_STARTED,
lambda engine: valid_sampler.set_epoch(engine.state.epoch))
# Linearly decrease the learning rate from lr to zero
scheduler = PiecewiseLinear(optimizer, "lr",
[(0, config.lr),
(config.n_epochs * len(train_loader), 0.0)])
trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
# Prepare metrics - note how we compute distributed metrics
RunningAverage(output_transform=lambda x: x).attach(trainer, "loss")
metrics = {
"nll":
Loss(torch.nn.CrossEntropyLoss(ignore_index=-1),
output_transform=lambda x: (x[0][0], x[1][0])),
"accuracy":
Accuracy(output_transform=lambda x: (x[0][1], x[1][1]))
}
metrics.update({
"average_nll":
MetricsLambda(average_distributed_scalar, metrics["nll"], config),
"average_accuracy":
MetricsLambda(average_distributed_scalar, metrics["accuracy"], config)
})
metrics["average_ppl"] = MetricsLambda(math.exp, metrics["average_nll"])
for name, metric in metrics.items():
metric.attach(evaluator, name)
# On the main process: add progress bar, tensorboard, checkpoints and save model, configuration and tokenizer before we start to train
if config.local_rank in [-1, 0]:
pbar = ProgressBar(persist=True)
pbar.attach(trainer, metric_names=["loss"])
evaluator.add_event_handler(
Events.COMPLETED, lambda _: pbar.log_message(
"Validation: %s" % pformat(evaluator.state.metrics)))
tb_logger = TensorboardLogger(log_dir=config.log_dir)
tb_logger.attach(trainer,
log_handler=OutputHandler(tag="training",
metric_names=["loss"]),
event_name=Events.ITERATION_COMPLETED)
tb_logger.attach(trainer,
log_handler=OptimizerParamsHandler(optimizer),
event_name=Events.ITERATION_STARTED)
tb_logger.attach(evaluator,
log_handler=OutputHandler(tag="validation",
metric_names=list(
metrics.keys()),
another_engine=trainer),
event_name=Events.EPOCH_COMPLETED)
checkpoint_handler = ModelCheckpoint(tb_logger.writer.log_dir,
'checkpoint',
save_interval=1,
n_saved=3)
trainer.add_event_handler(
Events.EPOCH_COMPLETED, checkpoint_handler,
{'mymodel': getattr(model, 'module', model)
}) # "getattr" take care of distributed encapsulation
torch.save(config,
tb_logger.writer.log_dir + '/model_training_args.bin')
getattr(model, 'module', model).config.to_json_file(
os.path.join(tb_logger.writer.log_dir, CONFIG_NAME))
tokenizer.save_vocabulary(tb_logger.writer.log_dir)
# Run the training
trainer.run(train_loader, max_epochs=config.n_epochs)
# On the main process: close tensorboard logger and rename the last checkpoint (for easy re-loading with OpenAIGPTModel.from_pretrained method)
if config.local_rank in [-1, 0] and config.n_epochs > 0:
os.rename(
checkpoint_handler._saved[-1][1][-1],
os.path.join(tb_logger.writer.log_dir, WEIGHTS_NAME)
) # TODO: PR in ignite to have better access to saved file paths (cleaner)
tb_logger.close()
def main():
parser = argparse.ArgumentParser()
# Required parameters
parser.add_argument("--model", type=str, default='ffn', help="model's name")
parser.add_argument("--mode", type=int, choices=[0, 1, 2], default=None)
parser.add_argument("--SNRdb", type=float, default=None)
parser.add_argument("--pilot_version", type=int, choices=[1, 2], default=1)
parser.add_argument("--loss_type", type=str, default="BCELoss")
parser.add_argument("--train_batch_size", type=int, default=128)
parser.add_argument("--valid_batch_size", type=int, default=128)
parser.add_argument("--gradient_accumulation_steps", type=int, default=1)
parser.add_argument("--max_norm", type=float, default=-1)
parser.add_argument("--lr", type=float, default=1e-3)
parser.add_argument("--noise_lambda", type=float, default=1.0)
parser.add_argument("--lr_scheduler", type=str, choices=["linear", "cycle", "cosine"], default="linear")
parser.add_argument("--reset_lr_scheduler", type=str, choices=["linear", "cycle", "cosine"], default=None)
parser.add_argument("--reset_trainer", action='store_true')
parser.add_argument("--modify_model", action='store_true')
parser.add_argument("--wd", type=float, default=1e-4, help="weight decay")
parser.add_argument("--eval_iter", type=int, default=10)
parser.add_argument("--save_iter", type=int, default=10)
parser.add_argument("--n_epochs", type=int, default=10)
parser.add_argument("--flush_dataset", type=int, default=0)
parser.add_argument("--no_cache", action='store_true')
parser.add_argument("--with_pure_y", action='store_true')
parser.add_argument("--with_h", action='store_true')
parser.add_argument("--only_l1", action='store_true', help="Only loss 1")
parser.add_argument("--interpolation", action='store_true', help="if interpolate between pure and reconstruction.")
parser.add_argument("--data_dir", type=str, default="data")
parser.add_argument("--cache_dir", type=str, default="train_cache")
parser.add_argument("--output_path", type=str, default="runs", help="model save")
parser.add_argument("--resume_from", type=str, default=None, help="resume training.")
parser.add_argument("--first_cache_index", type=int, default=0)
parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu",
help="Device (cuda or cpu)")
parser.add_argument("--local_rank", type=int, default=-1,
help="Local rank for distributed training (-1: not distributed)")
parser.add_argument("--seed", type=int, default=43)
parser.add_argument("--debug", action='store_true')
args = parser.parse_args()
args.output_path = os.path.join(args.output_path, f'pilot_{args.pilot_version}')
args.cache_dir = os.path.join(args.data_dir, args.cache_dir)
# Setup CUDA, GPU & distributed training
args.distributed = (args.local_rank != -1)
if not args.distributed:
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
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", init_method='env://')
args.n_gpu = torch.cuda.device_count() if not args.distributed else 1
args.device = device
# Set seed
set_seed(args)
logger = setup_logger("trainer", distributed_rank=args.local_rank)
# Model construction
model = getattr(models, args.model)(args)
model = model.to(device)
optimizer = AdamW(model.parameters(), lr = args.lr, weight_decay=args.wd)
if args.loss_type == "MSELoss":
criterion = nn.MSELoss(reduction='sum').to(device)
else:
criterion = getattr(nn, args.loss_type, getattr(auxiliary, args.loss_type, None))().to(device)
criterion2 = nn.MSELoss(reduction='sum').to(device)
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_dataset = SIGDataset(args, data_type="train")
valid_dataset = SIGDataset(args, data_type="valid")
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset) if args.distributed else None
valid_sampler = torch.utils.data.distributed.DistributedSampler(valid_dataset) if args.distributed else None
train_loader = DataLoader(train_dataset, sampler=train_sampler, batch_size=args.train_batch_size, pin_memory=True, shuffle=(not args.distributed))
valid_loader = DataLoader(valid_dataset, sampler=valid_sampler, batch_size=args.valid_batch_size, pin_memory=True, shuffle=False)
lr_scheduler = None
if args.lr_scheduler == "linear":
lr_scheduler = PiecewiseLinear(optimizer, "lr", [(0, args.lr), (args.n_epochs * len(train_loader), 0.0)])
elif args.lr_scheduler == "cycle":
lr_scheduler = LinearCyclicalScheduler(optimizer, 'lr', 0.0, args.lr, args.eval_iter * len(train_loader))
elif args.lr_scheduler == "cosine":
lr_scheduler = CosineAnnealingScheduler(optimizer, 'lr', args.lr, 0.0, args.eval_iter * len(train_loader))
# Training function and trainer
def update(engine, batch):
model.train()
y, x_label, y_pure, H = train_dataset.prepare_batch(batch, device=args.device)
if args.with_pure_y and args.with_h:
x_pred, y_pure_pred, H_pred = model(y, pure=y_pure, H=H, opp=True)
loss_1 = criterion(x_pred, x_label) / args.gradient_accumulation_steps
if args.loss_type == "MSELoss":
loss_1 = loss_1 / x_pred.size(0)
loss_noise = criterion2(y_pure_pred, y_pure) / y.size(0) / args.gradient_accumulation_steps
loss_noise_h = criterion2(H_pred, H) / H.size(0) / args.gradient_accumulation_steps
if args.only_l1:
loss = loss_1
else:
loss = loss_1 + loss_noise * args.noise_lambda + loss_noise_h
output = (loss.item(), loss_1.item(), loss_noise.item(), loss_noise_h.item())
elif args.with_pure_y:
x_pred, y_pure_pred = model(y, pure=y_pure if args.interpolation else None, opp=True)
loss_1 = criterion(x_pred, x_label) / args.gradient_accumulation_steps
loss_noise = criterion2(y_pure_pred, y_pure) / y.size(0) / args.gradient_accumulation_steps
loss = loss_1 + loss_noise * args.noise_lambda
output = (loss.item(), loss_1.item(), loss_noise.item())
elif args.with_h:
x_pred, H_pred = model(y, opp=True)
loss_1 = criterion(x_pred, x_label) / args.gradient_accumulation_steps
loss_noise = criterion2(H_pred, H) / H.size(0) / args.gradient_accumulation_steps
loss = loss_1 + loss_noise * args.noise_lambda
output = (loss.item(), loss_1.item(), loss_noise.item())
else:
x_pred = model(y)
loss_1 = criterion(x_pred, x_label) / args.gradient_accumulation_steps
loss = loss_1
output = (loss.item(), loss_1.item(), torch.zeros_like(loss_1).item())
loss.backward()
if args.max_norm > 0:
torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_norm)
if engine.state.iteration % args.gradient_accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
return output
trainer = Engine(update)
to_save = {"trainer": trainer, "model": model, "optimizer": optimizer, "lr_scheduler": lr_scheduler}
metric_names = ["loss", "l1", "ln"]
if args.with_pure_y and args.with_h:
metric_names.append("lnH")
common.setup_common_training_handlers(
trainer=trainer,
train_sampler=train_loader.sampler,
to_save=to_save,
save_every_iters=len(train_loader) * args.save_iter,
lr_scheduler=lr_scheduler,
output_names=metric_names,
with_pbars=False,
clear_cuda_cache=False,
output_path=args.output_path,
n_saved=2,
)
resume_from = args.resume_from
if resume_from is not None:
checkpoint_fp = Path(resume_from)
assert checkpoint_fp.exists(), "Checkpoint '{}' is not found".format(checkpoint_fp.as_posix())
logger.info("Resume from a checkpoint: {}".format(checkpoint_fp.as_posix()))
checkpoint = torch.load(checkpoint_fp.as_posix(), map_location="cpu")
if args.reset_trainer:
to_save.pop("trainer")
checkpoint_to_load = to_save if 'validation' not in resume_from else {"model": model}
Checkpoint.load_objects(to_load=checkpoint_to_load, checkpoint=checkpoint)
if args.reset_lr_scheduler is not None:
if args.reset_lr_scheduler == "linear":
lr_scheduler = PiecewiseLinear(optimizer, "lr", [(0, args.lr), (args.n_epochs * len(train_loader), 0.0)])
elif args.reset_lr_scheduler == "cycle":
lr_scheduler = LinearCyclicalScheduler(optimizer, 'lr', 0.0, args.lr, args.eval_iter * len(train_loader))
elif args.reset_lr_scheduler == "cosine":
lr_scheduler = CosineAnnealingScheduler(optimizer, 'lr', args.lr, 0.0, args.eval_iter * len(train_loader))
metrics = {
"accuracy": Accuracy(lambda output: (torch.round(output[0][0]), output[1][0])),
"loss_1": Loss(criterion, output_transform=lambda output: (output[0][0], output[1][0])),
"loss_noise": Loss(criterion2, output_transform=lambda output: (output[0][1], output[1][1]))
}
if args.with_pure_y and args.with_h:
metrics["loss_noise_h"] = Loss(criterion2, output_transform=lambda output: (output[0][2], output[1][2]))
def _inference(engine: Engine, batch: Sequence[torch.Tensor]) -> Union[Any, Tuple[torch.Tensor]]:
model.eval()
with torch.no_grad():
x, y, x_pure, H = valid_dataset.prepare_batch(batch, device=args.device, non_blocking=True)
if args.with_pure_y and args.with_h:
y_pred, x_pure_pred, h_pred = model(x, opp=True)
outputs = (y_pred, x_pure_pred, h_pred), (y, x_pure, H)
elif args.with_pure_y:
y_pred, x_pure_pred = model(x, opp=True)
outputs = (y_pred, x_pure_pred), (y, x_pure)
elif args.with_h:
y_pred, h_pred = model(x, opp=True)
outputs = (y_pred, h_pred), (y, H)
else:
y_pred = model(x)
x_pure_pred = x_pure
outputs = (y_pred, x_pure_pred), (y, x_pure)
return outputs
evaluator = Engine(_inference)
for name, metric in metrics.items():
metric.attach(evaluator, name)
trainer.add_event_handler(Events.EPOCH_COMPLETED(every=args.eval_iter), lambda _: evaluator.run(valid_loader))
if args.flush_dataset > 0:
trainer.add_event_handler(Events.EPOCH_COMPLETED(every=args.n_epochs//args.flush_dataset),
lambda _: train_loader.dataset.reset() if args.no_cache else train_loader.dataset.reload())
# On the main process: add progress bar, tensorboard, checkpoints and save model, configuration and tokenizer before we start to train
if args.local_rank in [-1, 0]:
pbar = ProgressBar(persist=True)
pbar.attach(trainer, metric_names=metric_names, output_transform=lambda _: {"lr": f"{optimizer.param_groups[0]['lr']:.2e}"})
evaluator.add_event_handler(Events.COMPLETED, lambda _: pbar.log_message("Validation: %s" % pformat(evaluator.state.metrics)))
tb_logger = common.setup_tb_logging(args.output_path, trainer, optimizer, evaluators={'validation': evaluator}, log_every_iters=1)
# Store 3 best models by validation accuracy:
common.gen_save_best_models_by_val_score(
save_handler=DiskSaver(args.output_path, require_empty=False),
evaluator=evaluator,
models={"model": model},
metric_name="accuracy",
n_saved=3,
trainer=trainer,
tag="validation"
)
# Run the training
trainer.run(train_loader, max_epochs=args.n_epochs)
if args.local_rank in [-1, 0]:
tb_logger.close()
def train(self):
""" """
loader_train, sampler_train = self.dataset_obj.get_torch_features(
split="train",
batch_size=self.hps.train_batch_size,
num_wrong_utts=self.hps.num_candidates - 1)
loader_valid, sampler_valid = self.dataset_obj.get_torch_features(
split="valid",
batch_size=self.hps.valid_batch_size,
num_wrong_utts=self.hps.num_candidates - 1)
# Print (readable) values of the first three samples, if debug mode.
if self.hps.debug:
print("*** SAMPLES ***")
input_ids = loader_train.dataset.tensors[0].tolist()
token_type_ids = loader_train.dataset.tensors[4].tolist()
mc_labels = loader_train.dataset.tensors[3].tolist()
for sample in range(3):
print("SAMPLE {}:".format(sample))
input_ids_txt = [
self.tokenizer.decode(input_ids[sample][i],
clean_up_tokenization_spaces=False)
for i in range(self.hps.num_candidates)
]
token_type_ids_txt = [
self.tokenizer.decode(token_type_ids[sample][i])
for i in range(self.hps.num_candidates)
]
for num in range(self.hps.num_candidates):
print("{}. candidate:".format(num + 1))
print("INPUT IDS: {}".format(input_ids_txt[num]))
print("TOKEN TYPE IDS: {}".format(token_type_ids_txt[num]))
print("Correct utterance: {}".format(mc_labels[sample] + 1))
trainer = Engine(self.update)
evaluator = Engine(self.inference)
# Evaluation at the end of each epoch, as well as at the beginning of the training.
trainer.add_event_handler(Events.EPOCH_COMPLETED,
lambda _: evaluator.run(loader_valid))
if self.hps.n_epochs < 1:
trainer.add_event_handler(Events.COMPLETED,
lambda _: evaluator.run(loader_valid))
if self.hps.eval_before_start:
trainer.add_event_handler(Events.STARTED,
lambda _: evaluator.run(loader_valid))
# Adding learning rate.
args = {"loader_train": loader_train}
self.learningrate(trainer=trainer, args=args)
# Adding metrics to the trainer
RunningAverage(output_transform=lambda x: x).attach(trainer, "loss")
metrics = self.metrics()
for name, metric in metrics.items():
metric.attach(evaluator, name)
# Printing evaluation results at the end of each epoch.
pbar = ProgressBar(persist=True)
pbar.attach(trainer, metric_names=["loss"])
evaluator.add_event_handler(
Events.COMPLETED, lambda _: pbar.log_message(
"Validation: %s" % pformat(evaluator.state.metrics)))
# Save model weights after each epoch.
log_dir = self.mk_logdir(self.hps.model_checkpoint)
checkpoint_handler = ModelCheckpoint(log_dir,
"checkpoint",
save_interval=1,
n_saved=3)
trainer.add_event_handler(
Events.EPOCH_COMPLETED, checkpoint_handler,
{"mymodel": getattr(self.model, "module", self.model)})
torch.save(self.hps, log_dir / "model_training_args.bin")
getattr(self.model, "module", self.model).config.to_json_file(
os.path.join(log_dir, CONFIG_NAME))
self.tokenizer.save_pretrained(log_dir)
# Run the training
trainer.run(loader_train, max_epochs=self.hps.n_epochs)
os.rename(checkpoint_handler._saved[-1][1][-1],
os.path.join(log_dir, WEIGHTS_NAME))
def train():
parser = ArgumentParser()
parser.add_argument("--train_path", type=str, default='data/spolin-train-acl.json', help="Set data path")
parser.add_argument("--valid_path", type=str, default='data/spolin-valid.json', help="Set data path")
parser.add_argument("--correct_bias", type=bool, default=False, help="Set to true to correct bias for Adam optimizer")
parser.add_argument("--lr", type=float, default=2e-5, help="Set learning rate")
parser.add_argument("--n_epochs", type=int, default=4, help="Set number of epochs")
parser.add_argument("--num_warmup_steps", type=float, default=1000, help="Set number of warm-up steps")
parser.add_argument("--num_total_steps", type=float, default=10000, help="Set number of total steps")
parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu", help="Device (cuda or cpu)")
parser.add_argument("--max_grad_norm", type=float, default=1.0, help="Set maximum gradient normalization.")
parser.add_argument("--pretrained_path", type=str, default='bert-base-uncased', help="Choose which pretrained model to use (bert-base-uncased, roberta-base, roberta-large, roberta-large-mnli)")
parser.add_argument("--batch_size", type=int, default=32, help="Provide the batch size")
parser.add_argument("--random_seed", type=int, default=42, help="Set the random seed")
parser.add_argument("--test", action='store_true', help="If true, run with small dataset for testing code")
parser.add_argument("--base", action='store_true', help="If true, run with base experiment configuration (training with spont only) for comparison")
args = parser.parse_args()
logging.basicConfig(level=logging.INFO)
logger.info("Arguments: {}".format(pformat(args)))
if 'roberta' in args.pretrained_path:
# initialize tokenizer and model
logger.info("Initialize model and tokenizer.")
tokenizer = RobertaTokenizer.from_pretrained(args.pretrained_path, cache_dir = '../pretrained_models')
model = RobertaForSequenceClassification.from_pretrained(args.pretrained_path, cache_dir='../pretrained_models')
### START MODEL MODIFICATION
# Pretrained model was not trained with token type ids.
# fix token type embeddings for finetuning. Without this, the model can only take 0s as valid input for token_type_ids
model.config.type_vocab_size = 2
model.roberta.embeddings.token_type_embeddings = torch.nn.Embedding(2, model.config.hidden_size)
model.roberta.embeddings.token_type_embeddings.weight.data.normal_(mean=0.0, std=model.config.initializer_range)
### END MOD
elif 'bert' in args.pretrained_path:
model = BertForSequenceClassification.from_pretrained(args.pretrained_path, cache_dir='../pretrained_models')
tokenizer = BertTokenizer.from_pretrained(args.pretrained_path, cache_dir='../pretrained_models')
model.to(args.device)
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_rate': 0.01},
{'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay_rate': 0.0}
]
optimizer = AdamW(optimizer_grouped_parameters,
lr=args.lr,
correct_bias = args.correct_bias)
scheduler = WarmupLinearSchedule(optimizer, warmup_steps=args.num_warmup_steps, t_total=args.num_total_steps)
logger.info("Prepare datasets")
logger.info("Loading train set...")
train_data = get_data(args.train_path)
valid_data = get_data(args.valid_path)
cornell_valid_data = {k: {'cornell': valid_data[k]['cornell']} for k in valid_data.keys()}
spont_valid_data = {k: {'spont': valid_data[k]['spont']} for k in valid_data.keys()}
train_loader, train_sampler = get_data_loaders(args, train_data, args.train_path, tokenizer)
logger.info("Loading validation set...")
valid_p = Path(args.valid_path)
cornell_valid_loader, cornell_valid_sampler = get_data_loaders(args, cornell_valid_data, f"{str(valid_p.parent)}/cornell_{valid_p.name}", tokenizer)
spont_valid_loader, spont_valid_sampler = get_data_loaders(args, spont_valid_data, f"{str(valid_p.parent)}/spont_{valid_p.name}", tokenizer)
# Training function and trainer
def update(engine, batch):
model.train()
batch = tuple(input_tensor.to(args.device) for input_tensor in batch)
b_input_ids, b_input_mask, b_input_segment, b_labels = batch
optimizer.zero_grad()
#roberta has issues with token_type_ids
loss, logits = model(b_input_ids, token_type_ids=b_input_segment, attention_mask=b_input_mask, labels=b_labels)
# loss, logits = model(b_input_ids, token_type_ids=None, attention_mask=b_input_mask, labels=b_labels)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
optimizer.step()
scheduler.step()
return loss.item(), logits, b_labels
trainer = Engine(update)
# Evaluation function and evaluator
def inference(engine, batch):
model.eval()
batch = tuple(input_tensor.to(args.device) for input_tensor in batch)
b_input_ids, b_input_mask, b_input_segment, b_labels = batch
with torch.no_grad():
#roberta has issues with token_type_ids
# loss, logits = model(b_input_ids, token_type_ids = None, attention_mask=b_input_mask, labels=b_labels)
loss, logits = model(b_input_ids, token_type_ids = b_input_segment, attention_mask=b_input_mask, labels=b_labels)
label_ids = b_labels
return logits, label_ids, loss.item()
cornell_evaluator = Engine(inference)
spont_evaluator = Engine(inference)
trainer.add_event_handler(Events.EPOCH_COMPLETED, lambda _: cornell_evaluator.run(cornell_valid_loader))
trainer.add_event_handler(Events.EPOCH_COMPLETED, lambda _: spont_evaluator.run(spont_valid_loader))
RunningAverage(output_transform=lambda x: x[0]).attach(trainer, "loss")
RunningAverage(Accuracy(output_transform=lambda x: (x[1], x[2]))).attach(trainer, "accuracy")
if torch.cuda.is_available():
GpuInfo().attach(trainer, name='gpu')
recall = Recall(output_transform=lambda x: (x[0], x[1]))
precision = Precision(output_transform=lambda x: (x[0], x[1]))
F1 = (precision * recall * 2 / (precision + recall)).mean()
accuracy = Accuracy(output_transform=lambda x: (x[0], x[1]))
metrics = {"recall": recall, "precision": precision, "f1": F1, "accuracy": accuracy, "loss": Average(output_transform=lambda x: x[2])}
for name, metric in metrics.items():
metric.attach(cornell_evaluator, name)
metric.attach(spont_evaluator, name)
pbar = ProgressBar(persist=True)
pbar.attach(trainer, metric_names=['loss', 'accuracy'])
pbar.attach(trainer, metric_names=['gpu:0 mem(%)', 'gpu:0 util(%)'])
cornell_evaluator.add_event_handler(Events.COMPLETED, lambda _: pbar.log_message("Cornell validation metrics:\n %s" % pformat(cornell_evaluator.state.metrics)))
spont_evaluator.add_event_handler(Events.COMPLETED, lambda _: pbar.log_message("Spont validation metrics:\n %s" % pformat(spont_evaluator.state.metrics)))
tb_logger = TensorboardLogger(log_dir=None)
tb_logger.attach(trainer, log_handler=OutputHandler(tag="training", metric_names=["loss"]), event_name=Events.ITERATION_COMPLETED)
tb_logger.attach(trainer, log_handler=OptimizerParamsHandler(optimizer), event_name=Events.ITERATION_STARTED)
tb_logger.attach(cornell_evaluator, log_handler=OutputHandler(tag="valid", metric_names=list(metrics.keys()), another_engine=trainer), event_name=Events.EPOCH_COMPLETED)
tb_logger.attach(spont_evaluator, log_handler=OutputHandler(tag="valid", metric_names=list(metrics.keys()), another_engine=trainer), event_name=Events.EPOCH_COMPLETED)
# tb_logger.writer.log_dir -> tb_logger.writer.logdir (this is the correct attribute name as seen in: https://tensorboardx.readthedocs.io/en/latest/_modules/tensorboardX/writer.html#SummaryWriter)
checkpoint_handler = ModelCheckpoint(tb_logger.writer.logdir, 'checkpoint', save_interval=1, n_saved=5)
trainer.add_event_handler(Events.EPOCH_COMPLETED, checkpoint_handler, {'mymodel': getattr(model, 'module', model)}) # "getattr" take care of distributed encapsulation
torch.save(args, tb_logger.writer.logdir + '/model_training_args.bin')
getattr(model, 'module', model).config.to_json_file(os.path.join(tb_logger.writer.logdir, CONFIG_NAME))
tokenizer.save_vocabulary(tb_logger.writer.logdir)
trainer.run(train_loader, max_epochs = args.n_epochs)
if args.n_epochs > 0:
os.rename(checkpoint_handler._saved[-1][1][-1], os.path.join(tb_logger.writer.logdir, WEIGHTS_NAME)) # TODO: PR in ignite to have better access to saved file paths (cleaner)
tb_logger.close()
def train():
parser = ArgumentParser()
parser.add_argument("--dataset_path",
type=str,
default="",
help="Path or url of the dataset.")
parser.add_argument("--train_batch_size",
type=int,
default=64,
help="Batch size for training")
parser.add_argument("--valid_batch_size",
type=int,
default=64,
help="Batch size for validation")
parser.add_argument("--gradient_accumulation_steps",
type=int,
default=8,
help="Accumulate gradients on several steps")
parser.add_argument("--lr",
type=float,
default=6.25e-4,
help="Learning rate")
parser.add_argument("--max_norm",
type=float,
default=1.0,
help="Clipping gradient norm")
parser.add_argument("--n_epochs",
type=int,
default=15,
help="Number of training epochs")
parser.add_argument(
"--eval_before_start",
action='store_true',
help="If true start with a first evaluation before training")
parser.add_argument("--device",
type=str,
default="cuda" if torch.cuda.is_available() else "cpu",
help="Device (cuda or cpu)")
parser.add_argument(
"--fp16",
type=str,
default="",
help=
"Set to O0, O1, O2 or O3 for fp16 training (see apex documentation)")
parser.add_argument(
"--local_rank",
type=int,
default=-1,
help="Local rank for distributed training (-1: not distributed)")
parser.add_argument("--gpt2_model_name",
type=str,
default="gpt2",
help="Path, url or short name of the model")
parser.add_argument('-d_model', type=int, default=512)
parser.add_argument('-d_inner_hid', type=int, default=2048)
parser.add_argument('-d_k', type=int, default=64)
parser.add_argument('-d_v', type=int, default=64)
parser.add_argument('-n_head', type=int, default=8)
parser.add_argument('-n_layers', type=int, default=6)
parser.add_argument('-warmup', '--n_warmup_steps', type=int, default=4000)
parser.add_argument('-dropout', type=float, default=0.1)
parser.add_argument('-embs_share_weight', action='store_true')
parser.add_argument('-proj_share_weight', action='store_true')
parser.add_argument('-label_smoothing', action='store_true')
args = parser.parse_args()
args.d_word_vec = args.d_model
# logging is set to INFO (resp. WARN) for main (resp. auxiliary) process. logger.info => log main process only, logger.warning => log all processes
logging.basicConfig(
level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN)
logger.warning(
"Running process %d", args.local_rank
) # This is a logger.warning: it will be printed by all distributed processes
logger.info("Arguments: %s", pformat(args))
# Initialize distributed training if needed
args.distributed = (args.local_rank != -1)
if args.distributed:
torch.cuda.set_device(args.local_rank)
args.device = torch.device("cuda", args.local_rank)
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
logger.info("Prepare tokenizer, pretrained model and optimizer.")
tokenizer_class = GPT2Tokenizer if "gpt2" in args.gpt2_model_name else OpenAIGPTTokenizer # cant use Autotokenizer because checkpoint could be a Path
tokenizer = tokenizer_class.from_pretrained(args.gpt2_model_name)
num_tokens = len(tokenizer.encoder)
num_added_tokens = tokenizer.add_special_tokens(
ATTR_TO_SPECIAL_TOKEN) # doesn't add if they are already there
model = Transformer(
num_tokens + num_added_tokens,
num_tokens + num_added_tokens,
src_pad_idx=tokenizer.convert_tokens_to_ids(SPECIAL_TOKENS[-1]),
trg_pad_idx=tokenizer.convert_tokens_to_ids(SPECIAL_TOKENS[-1]),
trg_emb_prj_weight_sharing=args.proj_share_weight,
emb_src_trg_weight_sharing=args.embs_share_weight,
d_k=args.d_k,
d_v=args.d_v,
d_model=args.d_model,
d_word_vec=args.d_word_vec,
d_inner=args.d_inner_hid,
n_layers=args.n_layers,
n_head=args.n_head,
dropout=args.dropout).to(args.device)
optimizer = AdamW(model.parameters(), lr=args.lr, correct_bias=True)
# Prepare model for FP16 and distributed training if needed (order is important, distributed should be the last)
if args.fp16:
from apex import amp # Apex is only required if we use fp16 training
model, optimizer = amp.initialize(model,
optimizer,
opt_level=args.fp16)
if args.distributed:
model = DistributedDataParallel(model,
device_ids=[args.local_rank],
output_device=args.local_rank)
logger.info("Prepare datasets")
train_loader, val_loader, train_sampler, valid_sampler = get_data_loaders(
args, tokenizer, tokenizer)
# Training function and trainer
def update(engine, batch):
model.train()
batch = tuple(input_tensor.to(args.device) for input_tensor in batch)
source_ids, target_ids, lm_labels = batch
(lm_loss), *_ = model(source_ids, target_ids, labels=lm_labels)
loss = lm_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_norm)
else:
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_norm)
if engine.state.iteration % args.gradient_accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
return loss.item()
trainer = Engine(update)
# Evaluation function and evaluator (evaluator output is the input of the metrics)
def inference(engine, batch):
model.eval()
with torch.no_grad():
batch = tuple(
input_tensor.to(args.device) for input_tensor in batch)
source_ids, target_ids, lm_labels = batch
#logger.info(tokenizer.decode(target_ids[0].tolist()))
lm_logits, *_ = model(source_ids, target_ids)
lm_logits_flat_shifted = lm_logits[..., :-1, :].contiguous().view(
-1, lm_logits.size(-1))
lm_labels_flat_shifted = lm_labels[..., 1:].contiguous().view(-1)
return (lm_logits_flat_shifted, ), (lm_labels_flat_shifted, )
evaluator = Engine(inference)
# Attach evaluation to trainer: we evaluate when we start the training and at the end of each epoch
trainer.add_event_handler(Events.EPOCH_COMPLETED,
lambda _: evaluator.run(val_loader))
if args.n_epochs < 1:
trainer.add_event_handler(Events.COMPLETED,
lambda _: evaluator.run(val_loader))
if args.eval_before_start:
trainer.add_event_handler(Events.STARTED,
lambda _: evaluator.run(val_loader))
# Make sure distributed data samplers split the dataset nicely between the distributed processes
if args.distributed:
trainer.add_event_handler(
Events.EPOCH_STARTED,
lambda engine: train_sampler.set_epoch(engine.state.epoch))
evaluator.add_event_handler(
Events.EPOCH_STARTED,
lambda engine: valid_sampler.set_epoch(engine.state.epoch))
# Linearly decrease the learning rate from lr to zero
scheduler = PiecewiseLinear(optimizer, "lr",
[(0, args.lr),
(args.n_epochs * len(train_loader), 0.0)])
trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
# Prepare metrics - note how we compute distributed metrics
RunningAverage(output_transform=lambda x: x).attach(trainer, "loss")
metrics = {
"nll":
Loss(torch.nn.CrossEntropyLoss(ignore_index=-100),
output_transform=lambda x: (x[0][0], x[1][0]))
}
metrics.update({
"average_nll":
MetricsLambda(average_distributed_scalar, metrics["nll"], args)
})
metrics["average_ppl"] = MetricsLambda(math.exp, metrics["average_nll"])
for name, metric in metrics.items():
metric.attach(evaluator, name)
# On the main process: add progress bar, tensorboard, checkpoints and save model, configuration and tokenizer before we start to train
if args.local_rank in [-1, 0]:
pbar = ProgressBar(persist=True)
pbar.attach(trainer, metric_names=["loss"])
evaluator.add_event_handler(
Events.COMPLETED, lambda _: pbar.log_message(
"Validation: %s" % pformat(evaluator.state.metrics)))
log_dir = make_logdir(args.gpt2_model_name, args.dataset_path)
tb_logger = TensorboardLogger(log_dir)
tb_logger.attach(trainer,
log_handler=OutputHandler(tag="training",
metric_names=["loss"]),
event_name=Events.ITERATION_COMPLETED)
tb_logger.attach(trainer,
log_handler=OptimizerParamsHandler(optimizer),
event_name=Events.ITERATION_STARTED)
tb_logger.attach(evaluator,
log_handler=OutputHandler(tag="validation",
metric_names=list(
metrics.keys()),
another_engine=trainer),
event_name=Events.EPOCH_COMPLETED)
checkpoint_handler = ModelCheckpoint(log_dir,
'checkpoint',
save_interval=1,
n_saved=4)
trainer.add_event_handler(
Events.EPOCH_COMPLETED, checkpoint_handler,
{'mymodel': getattr(model, 'module', model)
}) # "getattr" takes care of distributed encapsulation
torch.save(args, log_dir + '/model_training_args.bin')
#getattr(model, 'module', model).config.to_json_file(os.path.join(log_dir, CONFIG_NAME))
tokenizer.save_pretrained(log_dir)
# Run the training
trainer.run(train_loader, max_epochs=args.n_epochs)
# On the main process: close tensorboard logger and rename the last checkpoint (for easy re-loading with OpenAIGPTModel.from_pretrained method)
if args.local_rank in [-1, 0] and args.n_epochs > 0:
os.rename(
os.path.join(log_dir, checkpoint_handler._saved[-1][1]),
os.path.join(log_dir, WEIGHTS_NAME)
) # TODO: PR in ignite to have better access to saved file paths (cleaner)
tb_logger.close()
def train():
parser = ArgumentParser()
parser.add_argument(
"--dataset_path",
type=str,
default="",
help="Path or url of the dataset. If empty download from S3.")
parser.add_argument("--dataset_cache",
type=str,
default='persona_comet_weak_label_preprocessed',
help="Path or url of the dataset cache")
parser.add_argument("--model_checkpoint",
type=str,
default="openai-gpt",
help="Path, url or short name of the model")
parser.add_argument("--num_candidates",
type=int,
default=2,
help="Number of candidates for training")
parser.add_argument("--max_history",
type=int,
default=2,
help="Number of previous exchanges to keep in history")
parser.add_argument("--train_batch_size",
type=int,
default=4,
help="Batch size for training")
parser.add_argument("--valid_batch_size",
type=int,
default=4,
help="Batch size for validation")
parser.add_argument("--gradient_accumulation_steps",
type=int,
default=8,
help="Accumulate gradients on several steps")
parser.add_argument("--lr",
type=float,
default=6.25e-5,
help="Learning rate")
parser.add_argument("--lm_coef",
type=float,
default=1.0,
help="LM loss coefficient")
parser.add_argument("--mc_coef",
type=float,
default=1.0,
help="Multiple-choice loss coefficient")
parser.add_argument("--max_norm",
type=float,
default=1.0,
help="Clipping gradient norm")
parser.add_argument("--n_epochs",
type=int,
default=3,
help="Number of training epochs")
parser.add_argument("--personality_permutations",
type=int,
default=1,
help="Number of permutations of personality sentences")
parser.add_argument(
"--eval_before_start",
action='store_true',
help="If true start with a first evaluation before training")
parser.add_argument("--device",
type=str,
default="cuda" if torch.cuda.is_available() else "cpu",
help="Device (cuda or cpu)")
parser.add_argument(
"--fp16",
type=str,
default="",
help=
"Set to O0, O1, O2 or O3 for fp16 training (see apex documentation)")
parser.add_argument(
"--local_rank",
type=int,
default=-1,
help="Local rank for distributed training (-1: not distributed)")
parser.add_argument("--num_beams",
type=int,
default=5,
help="Number of beams for comet expansion")
parser.add_argument("--test_run_num",
type=int,
default=-1,
help="Datapoints to run with in a test run")
parser.add_argument("--exp_name",
type=str,
default="",
required=True,
help="Provide an experiment name")
parser.add_argument("--do_train", action='store_true', help="Do training")
parser.add_argument("--do_eval", action='store_true', help="Do Evaluation")
parser.add_argument("--no_persona",
action='store_true',
help="No Persona Evaluation")
parser.add_argument("--no_comet_persona",
action='store_true',
help="No Persona Evaluation")
parser.add_argument("--uniform_prior",
action='store_true',
help="Uniform prior")
parser.add_argument("--entropy_regularize_prior_wt",
type=float,
default=0.0,
help="entropy regularize prior")
parser.add_argument("--training_type",
type=str,
default="",
help="Marginalize or Reinforce")
parser.add_argument("--use_baseline",
action='store_true',
help="Use baseline")
parser.add_argument("--moving_avg_ratio",
type=float,
default=0.99,
help="Moving avg ratio for running mean baseline")
parser.add_argument("--reinforce_loss_coef",
type=float,
default=0.99,
help="Loss coef for reinforce")
parser.add_argument("--prior_model",
type=str,
default="bow",
help="Prior model selection")
parser.add_argument("--log_dir",
type=str,
default="",
required=True,
help="Provide a log dir")
parser.add_argument("--use_structured_prior",
action='store_true',
default=False,
help="Use effect type as feature")
parser.add_argument("--use_structured_prior_binarypotential",
action='store_true',
default=False,
help="")
parser.add_argument("--effect_emb_dim",
type=int,
default=6,
help="Embedding type while computing effect feature")
args = parser.parse_args()
# logging is set to INFO (resp. WARN) for main (resp. auxiliary) process. logger.info => log main process only, logger.warning => log all processes
logging.basicConfig(
level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN)
print(
"Running process {}".format(args.local_rank)
) # This is a logger.warning: it will be printed by all distributed processes
print("Arguments: {}".format(pformat(args)))
# Initialize distributed training if needed
args.distributed = (args.local_rank != -1)
if args.distributed:
torch.cuda.set_device(args.local_rank)
args.device = torch.device("cuda", args.local_rank)
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
print("Prepare tokenizer, pretrained model and optimizer.")
tokenizer_class = GPT2Tokenizer if "gpt2" in args.model_checkpoint else OpenAIGPTTokenizer # cant use Autotokenizer because checkpoint could be a Path
tokenizer = tokenizer_class.from_pretrained(args.model_checkpoint)
model_class = GPT2DoubleHeadsModel if "gpt2" in args.model_checkpoint else OpenAIGPTDoubleHeadsModel
if args.do_eval and not args.do_train:
print('Loading model from checkpoint {}'.format(args.model_checkpoint))
# model = model_class.from_pretrained(args.model_checkpoint)
# model.to(args.device)
model = LatentMarginalizedModel(args, generator_class=model_class)
print('Num parameters: {}'.format(count_parameters(model)))
model.to(args.device)
# Add special tokens if they are not already added
add_special_tokens_(model, tokenizer)
optimizer = AdamW(model.parameters(), lr=args.lr, correct_bias=True)
# Prepare model for FP16 and distributed training if needed (order is important, distributed should be the last)
if args.fp16:
from apex import amp # Apex is only required if we use fp16 training
model, optimizer = amp.initialize(model,
optimizer,
opt_level=args.fp16)
if args.distributed:
model = DistributedDataParallel(model,
device_ids=[args.local_rank],
output_device=args.local_rank)
print("Prepare datasets")
start = datetime.now()
train_dataset = PersonaChatDataset(args, tokenizer, split='train')
if args.do_eval:
val_dataset = PersonaChatDataset(args, tokenizer, split='valid')
train_sampler = torch.utils.data.sampler.RandomSampler(train_dataset)
if args.no_comet_persona:
max_num_persona = MAX_NUM_PERSONA
else:
max_num_persona = MAX_NUM_COMET_PERSONA
train_loader = DataLoader(train_dataset,
sampler=train_sampler,
batch_size=args.train_batch_size,
collate_fn=partial(
collate_dialog,
max_num_persona=max_num_persona),
pin_memory=True)
if args.do_eval:
val_loader = DataLoader(val_dataset,
shuffle=False,
batch_size=args.valid_batch_size,
collate_fn=partial(
collate_dialog,
max_num_persona=max_num_persona),
pin_memory=True)
print('{} - Data loaded. Starting training'.format(datetime.now() - start))
# train_loader, val_loader, train_sampler, valid_sampler = get_data_loaders(args, tokenizer)
# Training function and trainer
def update(engine, batch):
model.train()
batch = tuple(input_tensor.to(args.device) for input_tensor in batch)
input_ids, token_type_ids, lm_labels, mc_token_ids, mc_labels, persona, history, effects = batch
(lm_loss), (mc_loss), (loss_prior), (conditional_lm_loss), (
num_labels), (track_rewards) = model(input_ids=input_ids,
token_type_ids=token_type_ids,
mc_token_ids=mc_token_ids,
lm_labels=lm_labels,
mc_labels=mc_labels,
persona=persona,
history=history,
effects=effects)
loss = (lm_loss * args.lm_coef +
mc_loss * args.mc_coef) / 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_norm)
else:
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_norm)
if engine.state.iteration % args.gradient_accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
return loss.item(), lm_loss.item(), mc_loss.item(), loss_prior.item(
), conditional_lm_loss.item(), track_rewards.item()
trainer = Engine(update)
# Evaluation function and evaluator (evaluator output is the input of the metrics)
def inference(engine, batch):
model.eval()
with torch.no_grad():
batch = tuple(
input_tensor.to(args.device) for input_tensor in batch)
input_ids, mc_token_ids, lm_labels, mc_labels, token_type_ids = batch
# print(tokenizer.decode(input_ids[0, -1, :].tolist()))
# if we dont send labels to model, it doesnt return losses
lm_logits, mc_logits, *_ = model(
input_ids,
token_type_ids=token_type_ids,
mc_token_ids=mc_token_ids,
)
lm_logits_flat_shifted = lm_logits[..., :-1, :].contiguous().view(
-1, lm_logits.size(-1))
lm_labels_flat_shifted = lm_labels[..., 1:].contiguous().view(-1)
return (lm_logits_flat_shifted,
mc_logits), (lm_labels_flat_shifted, mc_labels)
evaluator = Engine(inference)
# Make sure distributed data samplers split the dataset nicely between the distributed processes
# if args.distributed:
# trainer.add_event_handler(Events.EPOCH_STARTED, lambda engine: train_sampler.set_epoch(engine.state.epoch))
# evaluator.add_event_handler(Events.EPOCH_STARTED, lambda engine: valid_sampler.set_epoch(engine.state.epoch))
# Linearly decrease the learning rate from lr to zero
scheduler = PiecewiseLinear(optimizer, "lr",
[(0, args.lr),
(args.n_epochs * len(train_loader), 0.0)])
trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
# Prepare metrics - note how we compute distributed metrics
RunningAverage(output_transform=lambda x: x[0]).attach(trainer, "loss")
RunningAverage(output_transform=lambda x: x[1]).attach(trainer, "lm_loss")
RunningAverage(output_transform=lambda x: x[2]).attach(trainer, "mc_loss")
RunningAverage(output_transform=lambda x: x[3]).attach(
trainer, "prior_loss")
RunningAverage(output_transform=lambda x: x[4]).attach(
trainer, "cond_lm_loss")
RunningAverage(output_transform=lambda x: x[5]).attach(trainer, "rewards")
metrics = {
"nll":
Loss(torch.nn.CrossEntropyLoss(ignore_index=-100),
output_transform=lambda x: (x[0][0], x[1][0])),
"accuracy":
Accuracy(output_transform=lambda x: (x[0][1], x[1][1]))
}
metrics.update({
"average_nll":
MetricsLambda(average_distributed_scalar, metrics["nll"], args),
"average_accuracy":
MetricsLambda(average_distributed_scalar, metrics["accuracy"], args)
})
metrics["average_ppl"] = MetricsLambda(math.exp, metrics["average_nll"])
for name, metric in metrics.items():
metric.attach(evaluator, name)
# On the main process: add progress bar, tensorboard, checkpoints and save model, configuration and tokenizer before we start to train
def print_model_save(engine):
print("Training complete. Saving Model.")
def print_validation(engine):
print("Model saved. Starting validation.")
if args.local_rank in [-1, 0]:
pbar = ProgressBar(persist=True)
pbar.attach(trainer,
metric_names=[
"loss", "lm_loss", "mc_loss", "prior_loss",
"cond_lm_loss", "rewards"
])
evaluator.add_event_handler(
Events.COMPLETED, lambda _: pbar.log_message(
"Validation: %s" % pformat(evaluator.state.metrics)))
log_dir = make_logdir(args.model_checkpoint, args.exp_name)
log_dir = os.path.join(args.log_dir, log_dir)
print("Logging at log dir: {}".format(log_dir))
# tb stuff
# tb_logger = TensorboardLogger(log_dir)
# tb_logger.attach(trainer, log_handler=OutputHandler(tag="training", metric_names=["loss"]), event_name=Events.ITERATION_COMPLETED)
# tb_logger.attach(trainer, log_handler=OptimizerParamsHandler(optimizer), event_name=Events.ITERATION_STARTED)
# tb_logger.attach(evaluator, log_handler=OutputHandler(tag="validation", metric_names=list(metrics.keys()), another_engine=trainer), event_name=Events.EPOCH_COMPLETED)
# save model checkpoints
checkpoint_handler = ModelCheckpoint(log_dir,
'checkpoint',
save_interval=1,
n_saved=None)
trainer.add_event_handler(Events.EPOCH_COMPLETED, print_model_save)
trainer.add_event_handler(
Events.EPOCH_COMPLETED, checkpoint_handler,
{'mymodel': getattr(model, 'module', model)
}) # "getattr" takes care of distributed encapsulation
torch.save(args, log_dir + '/model_training_args.bin')
# getattr(model, 'module', model).config.to_json_file(os.path.join(log_dir, CONFIG_NAME))
tokenizer.save_pretrained(log_dir)
# Attach evaluation to trainer: we evaluate when we start the training and at the end of each epoch
trainer.add_event_handler(Events.EPOCH_COMPLETED, print_validation)
if args.do_eval:
trainer.add_event_handler(Events.EPOCH_COMPLETED,
lambda _: evaluator.run(val_loader))
if args.n_epochs < 1:
trainer.add_event_handler(Events.COMPLETED,
lambda _: evaluator.run(val_loader))
if args.eval_before_start:
trainer.add_event_handler(Events.STARTED,
lambda _: evaluator.run(val_loader))
# Run the training
if args.do_train:
trainer.run(train_loader, max_epochs=args.n_epochs)
if args.do_eval and not args.do_train:
print('Running only Evaluation. No Training.')
evaluator.run(val_loader)
# On the main process: close tensorboard logger and rename the last checkpoint (for easy re-loading with OpenAIGPTModel.from_pretrained method)
if args.local_rank in [-1, 0] and args.n_epochs > 0 and args.do_train:
os.rename(
os.path.join(log_dir, checkpoint_handler._saved[-1][1]),
os.path.join(log_dir, WEIGHTS_NAME)
) # TODO: PR in ignite to have better access to saved file paths (cleaner)
def train():
parser = ArgumentParser()
parser.add_argument(
"--dataset_path",
type=str,
default="",
help="Path or url of the dataset. If empty download from S3.")
parser.add_argument("--model_checkpoint",
type=str,
default="t5-small",
help="Path, url or short name of the model")
parser.add_argument("--max_history",
type=int,
default=7,
help="Number of previous exchanges to keep in history")
parser.add_argument("--train_batch_size",
type=int,
default=10,
help="Batch size for training")
parser.add_argument("--valid_batch_size",
type=int,
default=10,
help="Batch size for validation")
parser.add_argument("--gradient_accumulation_steps",
type=int,
default=12,
help="Accumulate gradients on several steps")
parser.add_argument("--lr", type=float, default=6e-4, help="Learning rate")
parser.add_argument("--max_norm",
type=float,
default=1.0,
help="Clipping gradient norm")
parser.add_argument("--n_epochs",
type=int,
default=3,
help="Number of training epochs")
parser.add_argument("--device",
type=str,
default="cuda" if torch.cuda.is_available() else "cpu",
help="Device (cuda or cpu)")
parser.add_argument(
"--fp16",
type=str,
default="",
help=
"Set to O0, O1, O2 or O3 for fp16 training (see apex documentation)")
parser.add_argument(
"--local_rank",
type=int,
default=-1,
help="Local rank for distributed training (-1: not distributed)")
parser.add_argument("--save_name", type=str, default="")
parser.add_argument("--mask_ratio", type=float, default=0.15)
parser.add_argument("--objective",
type=str,
default="span_denosing",
help="response_generation, span_denosing, both")
args = parser.parse_args()
# logging is set to INFO (resp. WARN) for main (resp. auxiliary) process. logger.info => log main process only, logger.warning => log all processes
logging.basicConfig(
level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN)
logger.warning(
"Running process %d", args.local_rank
) # This is a logger.warning: it will be printed by all distributed processes
logger.info("Arguments: %s", pformat(args))
# Initialize distributed training if needed
args.distributed = (args.local_rank != -1)
if args.distributed:
torch.cuda.set_device(args.local_rank)
args.device = torch.device("cuda", args.local_rank)
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
logger.info("Prepare tokenizer, pretrained model and optimizer.")
tokenizer = T5Tokenizer.from_pretrained(args.model_checkpoint)
model = T5ForConditionalGeneration.from_pretrained(args.model_checkpoint)
model.to(args.device)
# Add special tokens if they are not already added
add_special_tokens_(model, tokenizer)
optimizer = AdamW(model.parameters(), lr=args.lr, correct_bias=True)
# Prepare model for FP16 and distributed training if needed (order is important, distributed should be the last)
if args.fp16:
from apex import amp # Apex is only required if we use fp16 training
model, optimizer = amp.initialize(model,
optimizer,
opt_level=args.fp16)
if args.distributed:
model = DistributedDataParallel(model,
device_ids=[args.local_rank],
output_device=args.local_rank)
def collate_fn(data):
batch = {
"corrupted_context": [],
"context": [],
"target": [],
"response": []
}
padded_dataset = {}
batch_size = len(data)
resp_sos, context_sos = tokenizer.convert_tokens_to_ids([
"<go_r>",
"<go_b>",
])
for x in data:
corrupted_context = ["fill : "]
target = []
length = len(x["context_words"])
mask_bool = random_spans_noise_mask(length=length,
noise_density=args.mask_ratio,
mean_noise_span_length=3.0)
mask_id = 0
#print(mask_bool)
for i in range(length):
if mask_bool[i]:
if i > 0 and mask_bool[i - 1]:
target.append(x["context_words"][i])
else:
target.append(f"<extra_id_{mask_id}>")
target.append(x["context_words"][i])
corrupted_context.append(f"<extra_id_{mask_id}>")
mask_id += 1
else:
corrupted_context.append(x["context_words"][i])
target.append("<eos_b>")
batch["context"].append(
tokenizer.encode("response : " + " ".join(x["context_words"])))
batch["corrupted_context"].append(
tokenizer.encode(" ".join(corrupted_context)))
batch["target"].append(tokenizer.encode(" ".join(target)))
batch["response"].append(tokenizer.encode(x["response"]))
# print(" ".join(x["context_words"]))
# print(" ".join(corrupted_context))
# print(" ".join(target))
# print("")
# print(tokenizer.decode(batch["corrupted_context"][-1]))
# print(tokenizer.decode(batch["target"][-1]))
# print(tokenizer.decode(batch["response"][-1]))
# print("")
context_ids, context_masks = padInput(batch["context"])
input_ids, masks = padInput(batch["corrupted_context"])
target_ids, target_inputs = padOutput(batch["target"])
response_ids, response_inputs = padOutput(batch["response"])
#inputs
padded_dataset["input_ids"] = torch.tensor(input_ids, dtype=torch.long)
padded_dataset["masks"] = torch.tensor(masks, dtype=torch.long)
padded_dataset["context_ids"] = torch.tensor(context_ids,
dtype=torch.long)
padded_dataset["context_masks"] = torch.tensor(context_masks,
dtype=torch.long)
padded_dataset["target_ids"] = torch.tensor(target_ids,
dtype=torch.long)
padded_dataset["response_ids"] = torch.tensor(response_ids,
dtype=torch.long)
padded_dataset["target_inputs"] = torch.tensor(np.concatenate((np.ones(
(batch_size, 1)) * context_sos, target_inputs[:, :-1]),
axis=1),
dtype=torch.long)
padded_dataset["response_inputs"] = torch.tensor(np.concatenate(
(np.ones((batch_size, 1)) * resp_sos, response_inputs[:, :-1]),
axis=1),
dtype=torch.long)
return padded_dataset
logger.info("Prepare datasets")
train_dataset, valid_dataset, train_sampler, valid_sampler = get_data(
args, tokenizer)
train_loader = DataLoader(train_dataset,
sampler=train_sampler,
batch_size=args.train_batch_size,
shuffle=(not args.distributed),
collate_fn=collate_fn,
num_workers=4)
val_loader = DataLoader(valid_dataset,
sampler=valid_sampler,
batch_size=args.valid_batch_size,
shuffle=False,
collate_fn=collate_fn,
num_workers=4)
logger.info("Train dataset length: {}".format(len(train_dataset)))
logger.info("Valid dataset length: {}".format(len(valid_dataset)))
# for batch in train_loader:
# #print(batch)
# exit(0)
# Training function and trainer
def update(engine, batch):
model.train()
batch = tuple(batch[input_name].to(args.device)
for input_name in MODEL_INPUTS)
input_ids, masks, context_ids, context_masks, target_ids, target_inputs, response_ids, response_inputs = batch
# print("input")
# print(tokenizer.decode(input_ids[0, :].tolist()))
# print("context_ids")
# print(tokenizer.decode(context_ids[0, :].tolist()))
# print("target")
# print(tokenizer.decode(target_ids[0, :].tolist()))
# print("target In")
# print(tokenizer.decode(target_inputs[0, :].tolist()))
# print("response_ids")
# print(tokenizer.decode(response_ids[0, :].tolist()))
# print("response_inputs")
# print(tokenizer.decode(response_inputs[0, :].tolist()))
#exit(0)
outputs = model(input_ids,
attention_mask=masks,
decoder_input_ids=target_inputs,
lm_labels=target_ids)
context_loss = outputs[0]
outputs = model(context_ids,
attention_mask=context_masks,
decoder_input_ids=response_inputs,
lm_labels=response_ids)
resp_loss = outputs[0]
loss = (context_loss + resp_loss) / args.gradient_accumulation_steps
loss = (context_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_norm)
else:
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_norm)
if engine.state.iteration % args.gradient_accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
return loss.item()
trainer = Engine(update)
# Evaluation function and evaluator (evaluator output is the input of the metrics)
def inference(engine, batch):
model.eval()
with torch.no_grad():
batch = tuple(batch[input_name].to(args.device)
for input_name in MODEL_INPUTS)
input_ids, masks, context_ids, context_masks, target_ids, target_inputs, response_ids, response_inputs = batch
outputs = model(
input_ids,
attention_mask=masks,
decoder_input_ids=target_inputs #, lm_labels=target_ids
)
context_logits = outputs[0]
outputs = model(
context_ids,
attention_mask=context_masks,
decoder_input_ids=response_inputs,
#lm_labels=response_ids
)
resp_logits = outputs[0]
context_logits_flat_shifted = context_logits.view(
-1, context_logits.size(-1))
context_labels_flat_shifted = target_ids.view(-1)
resp_logits_flat_shifted = resp_logits.view(
-1, resp_logits.size(-1))
resp_labels_flat_shifted = response_ids.view(-1)
return (context_logits_flat_shifted,
resp_logits_flat_shifted), (context_labels_flat_shifted,
resp_labels_flat_shifted)
#return (context_logits_flat_shifted, context_logits_flat_shifted), (context_labels_flat_shifted, context_labels_flat_shifted)
evaluator = Engine(inference)
# Attach evaluation to trainer: we evaluate when we start the training and at the end of each epoch
trainer.add_event_handler(Events.EPOCH_COMPLETED,
lambda _: evaluator.run(val_loader))
if args.n_epochs < 1:
trainer.add_event_handler(Events.COMPLETED,
lambda _: evaluator.run(val_loader))
# if args.eval_before_start:
# trainer.add_event_handler(Events.STARTED, lambda _: evaluator.run(val_loader))
# Make sure distributed data samplers split the dataset nicely between the distributed processes
if args.distributed:
trainer.add_event_handler(
Events.EPOCH_STARTED,
lambda engine: train_sampler.set_epoch(engine.state.epoch))
evaluator.add_event_handler(
Events.EPOCH_STARTED,
lambda engine: valid_sampler.set_epoch(engine.state.epoch))
# Linearly decrease the learning rate from lr to zero
scheduler = PiecewiseLinear(optimizer, "lr",
[(0, args.lr),
(args.n_epochs * len(train_loader), 0.0)])
trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
# Prepare metrics - note how we compute distributed metrics
RunningAverage(output_transform=lambda x: x).attach(trainer, "loss")
metrics = {
"span":
Loss(torch.nn.CrossEntropyLoss(ignore_index=-100),
output_transform=lambda x: (x[0][0], x[1][0])),
"response":
Loss(torch.nn.CrossEntropyLoss(ignore_index=-100),
output_transform=lambda x: (x[0][1], x[1][1]))
}
metrics.update({
"average_span":
MetricsLambda(average_distributed_scalar, metrics["span"], args),
"average_response":
MetricsLambda(average_distributed_scalar, metrics["response"], args)
})
metrics["average_response"] = MetricsLambda(math.exp,
metrics["average_response"])
for name, metric in metrics.items():
metric.attach(evaluator, name)
# On the main process: add progress bar, tensorboard, checkpoints and save model, configuration and tokenizer before we start to train
if args.local_rank in [-1, 0]:
pbar = ProgressBar(persist=True)
pbar.attach(trainer, metric_names=["loss"])
evaluator.add_event_handler(
Events.COMPLETED, lambda _: pbar.log_message(
"Validation: %s" % pformat(evaluator.state.metrics)))
if not os.path.exists(f"pretrained_model/{args.save_name}"):
os.makedirs(f"pretrained_model/{args.save_name}")
log_dir = f"pretrained_model/{args.save_name}"
tb_logger = TensorboardLogger(log_dir)
tb_logger.attach(trainer,
log_handler=OutputHandler(tag="training",
metric_names=["loss"]),
event_name=Events.ITERATION_COMPLETED)
tb_logger.attach(trainer,
log_handler=OptimizerParamsHandler(optimizer),
event_name=Events.ITERATION_STARTED)
tb_logger.attach(evaluator,
log_handler=OutputHandler(tag="validation",
metric_names=list(
metrics.keys()),
another_engine=trainer),
event_name=Events.EPOCH_COMPLETED)
checkpoint_handler = ModelCheckpoint(log_dir,
'checkpoint',
save_interval=1,
n_saved=3)
trainer.add_event_handler(
Events.EPOCH_COMPLETED, checkpoint_handler,
{'mymodel': getattr(model, 'module', model)
}) # "getattr" takes care of distributed encapsulation
torch.save(args, log_dir + '/model_training_args.bin')
getattr(model, 'module',
model).config.to_json_file(os.path.join(log_dir, CONFIG_NAME))
tokenizer.save_pretrained(log_dir)
# Run the training
trainer.run(train_loader, max_epochs=args.n_epochs)
# On the main process: close tensorboard logger and rename the last checkpoint (for easy re-loading with OpenAIGPTModel.from_pretrained method)
if args.local_rank in [-1, 0] and args.n_epochs > 0:
os.rename(
os.path.join(log_dir, checkpoint_handler._saved[-1][1]),
os.path.join(log_dir, WEIGHTS_NAME)
) # TODO: PR in ignite to have better access to saved file paths (cleaner)
tb_logger.close()
def train():
parser = ArgumentParser()
parser.add_argument("--dataset_path", type=str, default="", help="Path or url of the dataset. If empty download from S3.")
parser.add_argument("--dataset_cache", type=str, default='./dataset_cache', help="Path or url of the dataset cache")
parser.add_argument("--model_checkpoint", type=str, default="openai-gpt", help="Path, url or short name of the model")
parser.add_argument("--num_candidates", type=int, default=2, help="Number of candidates for training")
parser.add_argument("--max_history", type=int, default=2, help="Number of previous exchanges to keep in history")
parser.add_argument("--train_batch_size", type=int, default=4, help="Batch size for training")
parser.add_argument("--valid_batch_size", type=int, default=4, help="Batch size for validation")
parser.add_argument("--gradient_accumulation_steps", type=int, default=8, help="Accumulate gradients on several steps")
parser.add_argument("--lr", type=float, default=6.25e-5, help="Learning rate")
parser.add_argument("--lm_coef", type=float, default=1.0, help="LM loss coefficient")
parser.add_argument("--mc_coef", type=float, default=1.0, help="Multiple-choice loss coefficient")
parser.add_argument("--max_norm", type=float, default=1.0, help="Clipping gradient norm")
parser.add_argument("--n_epochs", type=int, default=3, help="Number of training epochs")
parser.add_argument("--personality_permutations", type=int, default=1, help="Number of permutations of personality sentences")
parser.add_argument("--eval_before_start", action='store_true', help="If true start with a first evaluation before training")
parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu", help="Device (cuda or cpu)")
parser.add_argument("--fp16", type=str, default="", help="Set to O0, O1, O2 or O3 for fp16 training (see apex documentation)")
parser.add_argument("--local_rank", type=int, default=-1, help="Local rank for distributed training (-1: not distributed)")
args = parser.parse_args()
# logging is set to INFO (resp. WARN) for main (resp. auxiliary) process. logger.info => log main process only, logger.warning => log all processes
logging.basicConfig(level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN)
logger.warning("Running process %d", args.local_rank) # This is a logger.warning: it will be printed by all distributed processes
logger.info("Arguments: %s", pformat(args))
# Initialize distributed training if needed
args.distributed = (args.local_rank != -1)
if args.distributed:
torch.cuda.set_device(args.local_rank)
args.device = torch.device("cuda", args.local_rank)
torch.distributed.init_process_group(backend='nccl', init_method='env://')
logger.info("Prepare tokenizer, pretrained model and optimizer - add special tokens for fine-tuning")
tokenizer_class = GPT2Tokenizer if "gpt2" in args.model_checkpoint else OpenAIGPTTokenizer
tokenizer = tokenizer_class.from_pretrained(args.model_checkpoint)
model_class = GPT2LMHeadModel if "gpt2" in args.model_checkpoint else OpenAIGPTLMHeadModel
model = model_class.from_pretrained(args.model_checkpoint)
tokenizer.set_special_tokens(SPECIAL_TOKENS)
model.set_num_special_tokens(len(SPECIAL_TOKENS))
model.to(args.device)
optimizer = OpenAIAdam(model.parameters(), lr=args.lr)
# Prepare model for FP16 and distributed training if needed (order is important, distributed should be the last)
if args.fp16:
from apex import amp # Apex is only required if we use fp16 training
model, optimizer = amp.initialize(model, optimizer, opt_level=args.fp16)
if args.distributed:
model = DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank)
logger.info("Prepare datasets")
train_loader, val_loader, train_sampler, valid_sampler = get_data_loaders(args, tokenizer)
# Training function and trainer
def update(engine, batch):
model.train()
batch = tuple(input_tensor.to(args.device) for input_tensor in batch)
lm_loss, mc_loss = model(*batch)
loss = (lm_loss * args.lm_coef + mc_loss * args.mc_coef) / 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_norm)
else:
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_norm)
if engine.state.iteration % args.gradient_accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
return loss.item()
trainer = Engine(update)
# Evaluation function and evaluator (evaluator output is the input of the metrics)
def inference(engine, batch):
model.eval()
with torch.no_grad():
batch = tuple(input_tensor.to(args.device) for input_tensor in batch)
input_ids, mc_token_ids, lm_labels, mc_labels, token_type_ids = batch
logger.info(tokenizer.decode(input_ids[0, -1, :].tolist()))
model_outputs = model(input_ids, mc_token_ids, token_type_ids=token_type_ids)
lm_logits, mc_logits = model_outputs[0], model_outputs[1] # So we can also use GPT2 outputs
lm_logits_flat_shifted = lm_logits[..., :-1, :].contiguous().view(-1, lm_logits.size(-1))
lm_labels_flat_shifted = lm_labels[..., 1:].contiguous().view(-1)
return (lm_logits_flat_shifted, mc_logits), (lm_labels_flat_shifted, mc_labels)
evaluator = Engine(inference)
# Attach evaluation to trainer: we evaluate when we start the training and at the end of each epoch
trainer.add_event_handler(Events.EPOCH_COMPLETED, lambda _: evaluator.run(val_loader))
if args.n_epochs < 1:
trainer.add_event_handler(Events.COMPLETED, lambda _: evaluator.run(val_loader))
if args.eval_before_start:
trainer.add_event_handler(Events.STARTED, lambda _: evaluator.run(val_loader))
# Make sure distributed data samplers split the dataset nicely between the distributed processes
if args.distributed:
trainer.add_event_handler(Events.EPOCH_STARTED, lambda engine: train_sampler.set_epoch(engine.state.epoch))
evaluator.add_event_handler(Events.EPOCH_STARTED, lambda engine: valid_sampler.set_epoch(engine.state.epoch))
# Linearly decrease the learning rate from lr to zero
scheduler = PiecewiseLinear(optimizer, "lr", [(0, args.lr), (args.n_epochs * len(train_loader), 0.0)])
trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
# Prepare metrics - note how we compute distributed metrics
RunningAverage(output_transform=lambda x: x).attach(trainer, "loss")
metrics = {"nll": Loss(torch.nn.CrossEntropyLoss(ignore_index=-1), output_transform=lambda x: (x[0][0], x[1][0])),
"accuracy": Accuracy(output_transform=lambda x: (x[0][1], x[1][1]))}
metrics.update({"average_nll": MetricsLambda(average_distributed_scalar, metrics["nll"], args),
"average_accuracy": MetricsLambda(average_distributed_scalar, metrics["accuracy"], args)})
metrics["average_ppl"] = MetricsLambda(math.exp, metrics["average_nll"])
for name, metric in metrics.items():
metric.attach(evaluator, name)
# On the main process: add progress bar, tensorboard, checkpoints and save model, configuration and tokenizer before we start to train
if args.local_rank in [-1, 0]:
pbar = ProgressBar(persist=True)
pbar.attach(trainer, metric_names=["loss"])
evaluator.add_event_handler(Events.COMPLETED, lambda _: pbar.log_message("Validation: %s" % pformat(evaluator.state.metrics)))
tb_logger = TensorboardLogger(log_dir=None)
tb_logger.attach(trainer, log_handler=OutputHandler(tag="training", metric_names=["loss"]), event_name=Events.ITERATION_COMPLETED)
tb_logger.attach(trainer, log_handler=OptimizerParamsHandler(optimizer), event_name=Events.ITERATION_STARTED)
tb_logger.attach(evaluator, log_handler=OutputHandler(tag="validation", metric_names=list(metrics.keys()), another_engine=trainer), event_name=Events.EPOCH_COMPLETED)
checkpoint_handler = ModelCheckpoint(tb_logger.writer.log_dir, 'checkpoint', save_interval=1, n_saved=3)
trainer.add_event_handler(Events.EPOCH_COMPLETED, checkpoint_handler, {'mymodel': getattr(model, 'module', model)}) # "getattr" take care of distributed encapsulation
torch.save(args, tb_logger.writer.log_dir + '/model_training_args.bin')
getattr(model, 'module', model).config.to_json_file(os.path.join(tb_logger.writer.log_dir, CONFIG_NAME))
tokenizer.save_vocabulary(tb_logger.writer.log_dir)
# Run the training
trainer.run(train_loader, max_epochs=args.n_epochs)
# On the main process: close tensorboard logger and rename the last checkpoint (for easy re-loading with OpenAIGPTModel.from_pretrained method)
if args.local_rank in [-1, 0] and args.n_epochs > 0:
os.rename(checkpoint_handler._saved[-1][1][-1], os.path.join(tb_logger.writer.log_dir, WEIGHTS_NAME)) # TODO: PR in ignite to have better access to saved file paths (cleaner)
tb_logger.close()
def train(
distributed=False, local_rank=-1, lr = 6.25e-5, dataset_path='../data/personachat_self_original.json',
dataset_cache=cached_path('../data/personachat_self_original.json'),
model_checkpoint='gpt2', num_candidates=2, max_history=5, train_batch_size=2, valid_batch_size=2,
gradient_accumulation_steps=8, lm_coef=1.0, mc_coef=1.0, max_norm=1.0, n_epochs=10,
personality_permutations=1, eval_before_start=False, device = 'cuda' if torch.cuda.is_available() else 'cpu',
fp16=''
):
'''
parser = ArgumentParser()
parser.add_argument("--dataset_path", type=str, default="", help="Path or url of the dataset. If empty download from S3.")
parser.add_argument("--dataset_cache", type=str, default='./dataset_cache', help="Path or url of the dataset cache")
parser.add_argument("--model_checkpoint", type=str, default="openai-gpt", help="Path, url or short name of the model")
parser.add_argument("--num_candidates", type=int, default=2, help="Number of candidates for training")
parser.add_argument("--max_history", type=int, default=2, help="Number of previous exchanges to keep in history")
parser.add_argument("--train_batch_size", type=int, default=4, help="Batch size for training")
parser.add_argument("--valid_batch_size", type=int, default=4, help="Batch size for validation")
parser.add_argument("--gradient_accumulation_steps", type=int, default=8, help="Accumulate gradients on several steps")
parser.add_argument("--lr", type=float, default=6.25e-5, help="Learning rate")
parser.add_argument("--lm_coef", type=float, default=1.0, help="LM loss coefficient")
parser.add_argument("--mc_coef", type=float, default=1.0, help="Multiple-choice loss coefficient")
parser.add_argument("--max_norm", type=float, default=1.0, help="Clipping gradient norm")
parser.add_argument("--n_epochs", type=int, default=3, help="Number of training epochs")
parser.add_argument("--personality_permutations", type=int, default=1, help="Number of permutations of personality sentences")
parser.add_argument("--eval_before_start", action='store_true', help="If true start with a first evaluation before training")
parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu", help="Device (cuda or cpu)")
parser.add_argument("--fp16", type=str, default="", help="Set to O0, O1, O2 or O3 for fp16 training (see apex documentation)")
parser.add_argument("--local_rank", type=int, default=-1, help="Local rank for distributed training (-1: not distributed)")
args = parser.parse_args()
# logging is set to INFO (resp. WARN) for main (resp. auxiliary) process. logger.info => log main process only, logger.warning => log all processes
logging.basicConfig(level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN)
logger.warning("Running process %d", args.local_rank) # This is a logger.warning: it will be printed by all distributed processes
logger.info("Arguments: %s", pformat(args))
'''
args = None
# Initialize distributed training if needed
distributed = (local_rank != -1)
if distributed:
torch.cuda.set_device(local_rank)
device = torch.device("cuda", local_rank)
torch.distributed.init_process_group(backend='nccl', init_method='env://')
#logger.info("Prepare tokenizer, pretrained model and optimizer - add special tokens for fine-tuning")
print(f'{datetime.now()}: Prepare tokenizer, pretrained model and optimizer - add special tokens for fine-tuning')
model = GPT2DoubleHeadsModel.from_pretrained('gpt2')
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
# We will use 5 special tokens:
# - <bos> to indicate the start of the sequence
# - <eos> to indicate the end of the sequence
# - <speaker1> to indicate the beginning and the tokens of an utterance from the user
# - <speaker2> to indicate the beginning and the tokens of an utterance from the bot
# - <pad> as a padding token to build batches of sequences
special_tokens = {
'bos_token': '<bos>',
'eos_token': '<eos>',
'additional_special_tokens': ['<speaker1>', '<speaker2>'],
'pad_token': '<pad>'
}
# We can add these special tokens to the vocabulary and the embeddings of the model:
tokenizer.add_special_tokens(special_tokens)
#model.config.num_special_tokens = len(special_tokens)
model.resize_token_embeddings(len(tokenizer))
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model.to(device)
optimizer = AdamW(model.parameters(), lr=lr)
# Prepare model for FP16 and distributed training if needed (order is important, distributed should be the last)
if fp16:
from apex import amp # Apex is only required if we use fp16 training
model, optimizer = amp.initialize(model, optimizer, opt_level=fp16)
if distributed:
model = DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank)
#logger.info("Prepare datasets")
print(f'{datetime.now()}: prepare datasets')
train_loader, val_loader, train_sampler, valid_sampler = get_data_loaders(tokenizer)
# Training function and trainer
def update(engine, batch):
model.train()
batch = tuple(input_tensor.to(device) for input_tensor in batch)
lm_loss, mc_loss, _, _, _ = model(*batch)
loss = (lm_loss * lm_coef + mc_loss * mc_coef) / gradient_accumulation_steps
if fp16:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), max_norm)
else:
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
if engine.state.iteration % gradient_accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
return loss.item()
trainer = Engine(update)
# Evaluation function and evaluator (evaluator output is the input of the metrics)
def inference(engine, batch):
model.eval()
with torch.no_grad():
batch = tuple(input_tensor.to(device) for input_tensor in batch)
input_ids, mc_token_ids, lm_labels, mc_labels, token_type_ids = batch
#logger.info(tokenizer.decode(input_ids[0, -1, :].tolist()))
print(f'{datetime.now()}: {tokenizer.decode(input_ids[0, -1, :].tolist())}')
model_outputs = model(input_ids, mc_token_ids, token_type_ids=token_type_ids)
lm_logits, mc_logits = model_outputs[0], model_outputs[1] # So we can also use GPT2 outputs
lm_logits_flat_shifted = lm_logits[..., :-1, :].contiguous().view(-1, lm_logits.size(-1))
lm_labels_flat_shifted = lm_labels[..., 1:].contiguous().view(-1)
return (lm_logits_flat_shifted, mc_logits), (lm_labels_flat_shifted, mc_labels)
evaluator = Engine(inference)
# Attach evaluation to trainer: we evaluate when we start the training and at the end of each epoch
trainer.add_event_handler(Events.EPOCH_COMPLETED, lambda _: evaluator.run(val_loader))
if n_epochs < 1:
trainer.add_event_handler(Events.COMPLETED, lambda _: evaluator.run(val_loader))
if eval_before_start:
trainer.add_event_handler(Events.STARTED, lambda _: evaluator.run(val_loader))
# Make sure distributed data samplers split the dataset nicely between the distributed processes
if distributed:
trainer.add_event_handler(Events.EPOCH_STARTED, lambda engine: train_sampler.set_epoch(engine.state.epoch))
evaluator.add_event_handler(Events.EPOCH_STARTED, lambda engine: valid_sampler.set_epoch(engine.state.epoch))
# Linearly decrease the learning rate from lr to zero
scheduler = PiecewiseLinear(optimizer, "lr", [(0, lr), (n_epochs * len(train_loader), 0.0)])
trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
# Prepare metrics - note how we compute distributed metrics
RunningAverage(output_transform=lambda x: x).attach(trainer, "loss")
metrics = {"nll": Loss(torch.nn.CrossEntropyLoss(ignore_index=-1), output_transform=lambda x: (x[0][0], x[1][0])),
"accuracy": Accuracy(output_transform=lambda x: (x[0][1], x[1][1]))}
#metrics.update({"average_nll": MetricsLambda(average_distributed_scalar, metrics["nll"], args),
# "average_accuracy": MetricsLambda(average_distributed_scalar, metrics["accuracy"], args)})
metrics["average_ppl"] = MetricsLambda(math.exp, metrics["nll"])
for name, metric in metrics.items():
metric.attach(evaluator, name)
# On the main process: add progress bar, tensorboard, checkpoints and save model, configuration and tokenizer before we start to train
if local_rank in [-1, 0]:
pbar = ProgressBar(persist=True)
pbar.attach(trainer, metric_names=["loss"])
evaluator.add_event_handler(Events.COMPLETED, lambda _: pbar.log_message("Validation: %s" % pformat(evaluator.state.metrics)))
tb_logger = TensorboardLogger(log_dir='../logs')
#tb_logger.attach(trainer, log_handler=OutputHandler(tag="training", metric_names=["loss"]), event_name=Events.ITERATION_COMPLETED)
#tb_logger.attach(trainer, log_handler=OptimizerParamsHandler(optimizer), event_name=Events.ITERATION_STARTED)
#tb_logger.attach(evaluator, log_handler=OutputHandler(tag="validation", metric_names=list(metrics.keys()), another_engine=trainer), event_name=Events.EPOCH_COMPLETED)
#checkpoint_handler = ModelCheckpoint(tb_logger.writer.log_dir, 'checkpoint', save_interval=1, n_saved=3)
#trainer.add_event_handler(Events.EPOCH_COMPLETED, checkpoint_handler,
# {'mymodel': getattr(model, 'module', model)})
# "getattr" take care of distributed encapsulation
#torch.save(args, tb_logger.writer.log_dir + '/model_training_args.bin')
#getattr(model, 'module', model).config.to_json_file(os.path.join(tb_logger.writer.log_dir, CONFIG_NAME))
#tokenizer.save_vocabulary(tb_logger.writer.log_dir)
# Run the training
trainer.run(train_loader, max_epochs=n_epochs)
