def test_optimizer_params_handler_wrong_setup():
with pytest.raises(TypeError):
OptimizerParamsHandler(optimizer=None)
optimizer = MagicMock(spec=torch.optim.Optimizer)
handler = OptimizerParamsHandler(optimizer=optimizer)
mock_logger = MagicMock()
mock_engine = MagicMock()
with pytest.raises(RuntimeError, match="Handler OptimizerParamsHandler works only with TensorboardLogger"):
handler(mock_engine, mock_logger, Events.ITERATION_STARTED)
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 add_tensorboard(engine_train, optimizer, model, log_dir):
"""Creates an ignite logger object and adds training elements such as weight and gradient histograms
Args:
engine_train (:obj:`ignite.engine`): the train engine to attach to the logger
optimizer (:obj:`torch.optim`): the model's optimizer
model (:obj:`torch.nn.Module`): the model being trained
log_dir (string): path to where tensorboard data should be saved
"""
# Create a logger
tb_logger = TensorboardLogger(log_dir=log_dir)
# Attach the logger to the trainer to log training loss at each iteration
tb_logger.attach(engine_train,
log_handler=OutputHandler(
tag="training",
output_transform=lambda loss: {"loss": loss}),
event_name=Events.ITERATION_COMPLETED)
# Attach the logger to the trainer to log optimizer's parameters, e.g. learning rate at each iteration
tb_logger.attach(engine_train,
log_handler=OptimizerParamsHandler(optimizer),
event_name=Events.EPOCH_COMPLETED)
# Attach the logger to the trainer to log model's weights as a histogram after each epoch
tb_logger.attach(engine_train,
log_handler=WeightsHistHandler(model),
event_name=Events.EPOCH_COMPLETED)
# Attach the logger to the trainer to log model's gradients as a histogram after each epoch
tb_logger.attach(engine_train,
log_handler=GradsHistHandler(model),
event_name=Events.EPOCH_COMPLETED)
tb_logger.close()
def on_training_started(engine):
# construct an optimizer
logger.info('Started Training...')
params = [p for p in model.parameters() if p.requires_grad]
engine.state.optimizer = torch.optim.SGD(params,
lr=lr,
momentum=momentum,
weight_decay=weight_decay)
tb_logger.attach(
trainer,
log_handler=OptimizerParamsHandler(engine.state.optimizer),
event_name=Events.ITERATION_STARTED
)
engine.state.scheduler = torch.optim.lr_scheduler.StepLR(engine.state.optimizer, step_size=step_size,
gamma=gamma)
if input_checkpoint:
# Load traininer states
trainer.state.epoch = input_checkpoint['epoch']
if 'iteration' in input_checkpoint:
trainer.state.iteration = input_checkpoint['iteration']
else:
trainer.state.iteration = int(hparam_dict['training_set_size'] / batch_size * input_checkpoint['epoch'])
if load_optimizer:
print('loading optimizer')
logger.info('Loading optimizer and scheduler...')
engine.state.optimizer.load_state_dict(input_checkpoint['optimizer'])
engine.state.scheduler.load_state_dict(input_checkpoint['lr_scheduler'])
engine.state.scheduler.last_epoch = trainer.state.epoch
else:
print('not loading optimizer')
def custom_setup(self):
if self.tensorboard_logs:
tb_logger = TensorboardLogger(log_dir=self.tensorboard_logs)
tb_logger.attach(self.trainer,
log_handler=OutputHandler(
tag="training",
output_transform=lambda loss: {'loss': loss}),
event_name=Events.ITERATION_COMPLETED)
tb_logger.attach(self.evaluator,
log_handler=OutputHandler(
tag="validation",
metric_names=["LossMetric"],
another_engine=self.trainer),
event_name=Events.EPOCH_COMPLETED)
if self.optional_tensorboard_features:
tb_logger.attach(self.trainer,
log_handler=OptimizerParamsHandler(
self.optimizer),
event_name=Events.ITERATION_STARTED)
tb_logger.attach(self.trainer,
log_handler=WeightsScalarHandler(self.model),
event_name=Events.ITERATION_COMPLETED)
tb_logger.attach(self.trainer,
log_handler=WeightsHistHandler(self.model),
event_name=Events.EPOCH_COMPLETED)
tb_logger.attach(self.trainer,
log_handler=GradsScalarHandler(self.model),
event_name=Events.ITERATION_COMPLETED)
# This is important to close the tensorboard file logger
@self.trainer.on(Events.COMPLETED)
def end_tensorboard(trainer):
logger.info("Training completed")
tb_logger.close()
if self.embeddings_name:
@self.trainer.on(Events.COMPLETED)
def log_embeddings(trainer):
if hasattr(self.model, self.embeddings_name) and hasattr(
self.dataset_splits, "vectorizer") and TENSORBOARD:
logger.info(
f"Logging embeddings ({self.embeddings_name}) to Tensorboard!"
)
embeddings = getattr(self.model,
self.embeddings_name).weight.data
metadata = [
str(self.dataset_splits.vectorizer.data_vocab.
_id2token[token_index]).encode('utf-8')
for token_index in range(embeddings.shape[0])
]
self.writer.add_embedding(
mat=embeddings,
metadata=metadata,
global_step=self.trainer.state.epoch)
def add_optimizer_params_logging(optimizer: torch.optim.Optimizer,
tb_logger: TensorboardLogger,
engine: Engine) -> None:
for parameter_name in optimizer.defaults.keys():
tb_logger.attach(
engine,
log_handler=OptimizerParamsHandler(optimizer, parameter_name),
event_name=Events.ITERATION_STARTED,
)
def test_optimizer_params():
optimizer = torch.optim.SGD([torch.Tensor(0)], lr=0.01)
wrapper = OptimizerParamsHandler(optimizer=optimizer, param_name="lr")
mock_logger = MagicMock(spec=TensorboardLogger)
mock_logger.writer = MagicMock()
mock_engine = MagicMock()
mock_engine.state = State()
mock_engine.state.iteration = 123
wrapper(mock_engine, mock_logger, Events.ITERATION_STARTED)
mock_logger.writer.add_scalar.assert_called_once_with("lr/group_0", 0.01, 123)
wrapper = OptimizerParamsHandler(optimizer, param_name="lr", tag="generator")
mock_logger = MagicMock(spec=TensorboardLogger)
mock_logger.writer = MagicMock()
wrapper(mock_engine, mock_logger, Events.ITERATION_STARTED)
mock_logger.writer.add_scalar.assert_called_once_with("generator/lr/group_0", 0.01, 123)
def train():
device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
model = Bert_SQG()
optimizer = AdamW(model.parameters(), lr=3e-5)
ds = dataloader.BertSQG_DataClass()
dl = DataLoader(ds, num_workers=4, batch_size=4)
scheduler = PiecewiseLinear(optimizer, "lr",
[(0, 3e-5),
(EPOCHS * len(ds) // BATCH_SIZE, 0.0)])
metrics = {"nll": Loss(torch.nn.CrossEntropyLoss(ignore_index=-1))}
def update(engine, batch):
model.train()
for i in range(0, len(batch) - 1):
x = batch[i].to(device)
y = batch[i + 1].to(device)
y_prime = model(x)
loss = criterion(y_prime[-1], y[-1]) / ITERATION_STEP
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
if engine.state.iteration % ITERATION_STEP == 0:
optimizer.step()
optimizer.zero_grad()
return loss.item()
trainer = Engine(update)
trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
RunningAverage(output_transform=lambda x: x).attach(trainer, "loss")
pbar = ProgressBar(persist=True)
pbar.attach(trainer, metric_names=["loss"])
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)
checkpoint_handler = ModelCheckpoint('./checkpoint',
'_checkpoint',
save_interval=1,
n_saved=3)
trainer.add_event_handler(Events.EPOCH_COMPLETED, checkpoint_handler,
{'bert_sqg': getattr(model, 'module', model)})
trainer.run(dl, max_epochs=EPOCHS)
tb_loger.close()
def __init__(self, name, model, log_dir, lr, lr_decay_step, adam=False):
"""
Initialize to train the given model.
:param name: The name of the model to be trained.
:param model: The model to be trained.
:param log_dir: String. The log directory of the tensorboard.
:param lr: Float. The learning rate.
:param lr_decay_step: Integer. The amount of steps the learning rate decays.
:param adam: Bool. Whether to use adam optimizer or not.
"""
super(Trainer, self).__init__(self.update_model)
self.model = model
# tqdm
ProgressBar(persist=True).attach(self)
# Optimizer
params = [p for p in model.parameters() if p.requires_grad]
if adam:
self.optimizer = torch.optim.Adam(params, lr=lr)
else:
self.optimizer = torch.optim.SGD(params, lr=lr, momentum=0.9)
# Scheduler
if lr_decay_step > 0:
self.scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer, step_size=lr_decay_step, gamma=0.1)
self.add_event_handler(Events.EPOCH_COMPLETED, lambda e: e.scheduler.step())
else:
self.scheduler = None
# Terminate if nan values found
self.add_event_handler(Events.ITERATION_COMPLETED, TerminateOnNan())
# Tensorboard logging
self.tb_logger = TensorboardLogger(log_dir=os.path.join(log_dir, name))
self.add_event_handler(Events.COMPLETED, lambda x: self.tb_logger.close())
self.tb_logger.attach(self,
log_handler=OptimizerParamsHandler(self.optimizer),
event_name=Events.EPOCH_COMPLETED)
self.tb_logger.attach(self,
log_handler=OutputHandler(tag='training', output_transform=lambda x: {
'rpn_box_loss': round(self.state.output['loss_rpn_box_reg'].item(), 4),
'rpn_cls_loss': round(self.state.output['loss_objectness'].item(), 4),
'roi_box_loss': round(self.state.output['loss_box_reg'].item(), 4),
'roi_cls_loss': round(self.state.output['loss_classifier'].item(), 4)
}),
event_name=Events.EPOCH_COMPLETED)
# Run on GPU (cuda) if available
if torch.cuda.is_available():
torch.cuda.set_device(int(get_free_gpu()))
model.cuda(torch.cuda.current_device())
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 logging_board(model_name="densenet121"):
from ignite.contrib.handlers.tensorboard_logger import (
TensorboardLogger,
OutputHandler,
OptimizerParamsHandler,
GradsHistHandler,
)
tb_logger = TensorboardLogger("board/" + model_name)
tb_logger.attach(
trainer,
log_handler=OutputHandler(
tag="training", output_transform=lambda loss: {"loss": loss}),
event_name=Events.ITERATION_COMPLETED,
)
tb_logger.attach(
val_evaluator,
log_handler=OutputHandler(
tag="validation",
metric_names=["accuracy", "loss"],
another_engine=trainer,
),
event_name=Events.EPOCH_COMPLETED,
)
tb_logger.attach(
trainer,
log_handler=OptimizerParamsHandler(IGTrainer.optimizer),
event_name=Events.ITERATION_STARTED,
)
tb_logger.attach(
trainer,
log_handler=GradsHistHandler(IGTrainer.model),
event_name=Events.EPOCH_COMPLETED,
)
tb_logger.close()
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()
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)
# 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)
def main():
args = get_args()
if 'e-SNLI-VE' in args.data_path:
args.no_image = False
else:
args.no_image = True
if not args.no_image:
args.no_premise = True
args.with_expl = True
'''Setup'''
t = datetime.today()
output_dir = os.path.join(args.output_folder,
f"{t.month}_{t.day}_{t.hour}_{t.minute}_{t.second}")
if not os.path.exists(output_dir):
os.makedirs(output_dir, exist_ok=True)
# 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(filename=os.path.join(output_dir, 'app.log'),
filemode='a',
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(f"Running process {args.local_rank}")
logger.info(f"Arguments: {pformat(args)}")
logger.info(f'Image not used:{args.no_image}')
logger.info(f'Premise not used:{args.no_premise}')
logger.info(f'Explanations used:{args.with_expl}')
'''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 = GPT2Tokenizer.from_pretrained(args.model_checkpoint)
tokenizer.add_special_tokens(SPECIAL_TOKENS_DICT)
if args.no_image:
model = GPT2LMHeadModel.from_pretrained(args.model_checkpoint)
else:
import image_gpt2_291
model = image_gpt2_291.GPT2LMHeadModel.from_pretrained(
args.model_checkpoint)
model.resize_token_embeddings(len(tokenizer))
model.to(args.device)
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)
model = model.module
logger.info("Prepare datasets")
train_loader, val_loader = get_data_loaders(args, tokenizer)
'''Training function and trainer'''
def train(engine, batch):
model.train()
batch = tuple(input_tensor.to(args.device) for input_tensor in batch)
if args.no_image:
input_ids, lm_label, label, input_mask = batch
else:
image, input_ids, lm_label, label, input_mask = batch
if args.no_image:
output = model(input_ids=input_ids,
# attention_mask=input_mask,
labels=lm_label)
else:
output = model(input_ids=input_ids,
images=image,
# attention_mask=input_mask,
labels=lm_label)
loss, logits, _ = output
loss = loss / args.gradient_accumulation_steps
if args.fp16:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
torch.nn.utils.clip_grad_norm_(
amp.master_params(optimizer), args.max_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()
if not args.with_expl:
lbl_accuracy = torch.eq(label, logits.argmax(
dim=1)).float().sum() / len(label)
return {
'loss': loss.item(),
'lbl_accuracy': lbl_accuracy.item()
}
else:
if engine.state.iteration % (args.gradient_accumulation_steps * 500) == 0:
input_output = list(zip(input_ids, logits))
random_item = random.choice(input_output)
in_sent = tokenizer.decode(list(filter(
lambda x: x != tokenizer.eos_token_id,
random_item[0])))
out_expl = tokenizer.decode(random_item[1].argmax(dim=1),
skip_special_tokens=True)
logger.info(f'MODEL INPUT: {in_sent}')
logger.info(f'GEN. EXPL {out_expl}')
logger.info('--------------------------------')
return {
'loss': loss.item(),
}
'''Validation function and validator (validator output is the input of the metrics)'''
def validation(engine, batch):
model.eval()
with torch.no_grad():
batch = tuple(input_tensor.to(args.device)
for input_tensor in batch)
if args.no_image:
input_ids, lm_label, label, input_mask = batch
else:
image, input_ids, lm_label, label, input_mask = batch
if args.no_image:
output = model(input_ids=input_ids,
# attention_mask=input_mask
)
else:
output = model(input_ids=input_ids,
images=image,
# attention_mask=input_mask
)
logits, _ = output
logits_shifted = logits[..., :-1, :].contiguous().view(-1,
logits.size(-1))
labels_shifted = lm_label[..., 1:].contiguous().view(-1)
return logits_shifted, labels_shifted
'''Engines'''
trainer = Engine(train)
validator = Engine(validation)
# t_total = len(
# train_loader) // args.gradient_accumulation_steps * args.n_epochs
# scheduler = get_linear_schedule_with_warmup(
# optimizer, num_warmup_steps=args.warmup_steps, num_training_steps=t_total)
'''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)
'''
Attach validation to trainer: we evaluate when we start the training and at the end of each epoch
'''
trainer.add_event_handler(Events.EPOCH_COMPLETED,
lambda _: validator.run(val_loader))
if args.eval_before_start:
trainer.add_event_handler(Events.STARTED,
lambda _: validator.run(val_loader))
'''Prepare metrics - note how we compute distributed metrics'''
RunningAverage(output_transform=lambda x: x['loss']).attach(
trainer, "loss")
RunningAverage(output_transform=lambda x: math.exp(
average_distributed_scalar(x['loss'], args))).attach(trainer, "ppl")
if not args.with_expl:
RunningAverage(output_transform=lambda x: 100 * x['lbl_accuracy']).attach(
trainer, "lbl_accuracy")
metrics = {}
metrics["lbl_loss"] = Loss(torch.nn.CrossEntropyLoss(),
output_transform=lambda x: (x[0], x[1]))
metrics["loss"] = MetricsLambda(
lambda l, a: average_distributed_scalar(
l / a.gradient_accumulation_steps, a), metrics["lbl_loss"], args)
metrics["ppl"] = MetricsLambda(math.exp, metrics["loss"])
if not args.with_expl:
metrics["lbl_accuracy"] = 100 * \
Accuracy(output_transform=lambda x: (x[0], x[1]))
for name, metric in metrics.items():
metric.attach(validator, 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", 'ppl'] if args.with_expl else ["loss", 'lbl_accuracy', 'ppl'])
validator.add_event_handler(Events.COMPLETED,
lambda _: pbar.log_message(
"Validation: %s" % pformat(validator.state.metrics)))
tb_logger = TensorboardLogger(log_dir=output_dir)
tb_logger.attach(trainer,
log_handler=OptimizerParamsHandler(optimizer),
event_name=Events.ITERATION_STARTED)
tb_logger.attach(trainer,
log_handler=OutputHandler(
tag="training",
metric_names=["loss"]),
event_name=Events.ITERATION_COMPLETED)
tb_logger.attach(trainer,
log_handler=OutputHandler(
tag="training",
metric_names=["ppl"] if args.with_expl else ["lbl_accuracy", "ppl"]),
event_name=Events.EPOCH_COMPLETED)
tb_logger.attach(validator,
log_handler=OutputHandler(
tag="validation",
metric_names=[
'ppl', 'loss'] if args.with_expl else['ppl', 'loss', 'lbl_accuracy'],
global_step_transform=lambda *args, **kwargs: trainer.state.iteration),
event_name=Events.EPOCH_COMPLETED)
checkpoint_handler = ModelCheckpoint(output_dir,
'checkpoint',
n_saved=8,
require_empty=False)
trainer.add_event_handler(Events.EPOCH_COMPLETED(every=1),
checkpoint_handler,
{'mymodel': getattr(model, 'module', model)})
# "getattr" take care of distributed encapsulation
torch.save(args, os.path.join(output_dir, 'model_training_args.bin'))
getattr(model, 'module', model).config.to_json_file(
os.path.join(output_dir, CONFIG_NAME))
tokenizer.save_vocabulary(output_dir)
'''Run the training'''
trainer.run(train_loader, max_epochs=args.n_epochs)
def run(train_loader, val_loader, epochs, lr, momentum, weight_decay, lr_step,
k1, k2, es_patience, log_dir):
model = Vgg16()
device = 'cpu'
if torch.cuda.is_available():
device = 'cuda'
model.to(device)
optimizer = optim.SGD(model.parameters(),
lr=lr,
momentum=momentum,
weight_decay=weight_decay)
lr_scheduler = ExponentialLR(optimizer, gamma=0.975)
# criterion = VAELoss(k1=k1, k2=k2).to(device)
def update_fn(engine, batch):
x, y = _prepare_batch(batch, device=device, non_blocking=True)
model.train()
optimizer.zero_grad()
output = model(x)
# Compute loss
loss = F.nll_loss(output, y)
loss.backward()
optimizer.step()
return {
"batchloss": loss.item(),
}
trainer = Engine(update_fn)
try:
GpuInfo().attach(trainer)
except RuntimeError:
print(
"INFO: By default, in this example it is possible to log GPU information (used memory, utilization). "
"As there is no pynvml python package installed, GPU information won't be logged. Otherwise, please "
"install it : `pip install pynvml`")
trainer.add_event_handler(Events.ITERATION_COMPLETED(every=lr_step),
lambda engine: lr_scheduler.step())
metric_names = [
'batchloss',
]
def output_transform(x, name):
return x[name]
for n in metric_names:
# We compute running average values on the output (batch loss) across all devices
RunningAverage(output_transform=partial(output_transform, name=n),
epoch_bound=False,
device=device).attach(trainer, n)
exp_name = datetime.now().strftime("%Y%m%d-%H%M%S")
log_path = log_dir + "/vgg_vae/{}".format(exp_name)
tb_logger = TensorboardLogger(log_dir=log_path)
tb_logger.attach(trainer,
log_handler=OutputHandler(tag="training",
metric_names=metric_names),
event_name=Events.ITERATION_COMPLETED)
tb_logger.attach(trainer,
log_handler=OptimizerParamsHandler(optimizer, "lr"),
event_name=Events.ITERATION_STARTED)
ProgressBar(persist=True,
bar_format="").attach(trainer,
event_name=Events.EPOCH_STARTED,
closing_event_name=Events.COMPLETED)
ProgressBar(persist=False, bar_format="").attach(trainer,
metric_names=metric_names)
# val process definition
def loss_output_transform(output):
return output
def acc_output_transform(output):
return output
customed_loss = Loss(loss_fn=F.nll_loss,
output_transform=loss_output_transform,
device=device)
customed_accuracy = Accuracy(output_transform=acc_output_transform,
device=device)
metrics = {'Loss': customed_loss, 'Accuracy': customed_accuracy}
def val_update_fn(engine, batch):
model.eval()
with torch.no_grad():
x, y = _prepare_batch(batch, device=device, non_blocking=True)
output = model(x)
return output, y
val_evaluator = Engine(val_update_fn)
for name, metric in metrics.items():
metric.attach(val_evaluator, name)
def run_evaluation(engine):
val_evaluator.run(val_loader)
trainer.add_event_handler(Events.EPOCH_COMPLETED, run_evaluation)
trainer.add_event_handler(Events.COMPLETED, run_evaluation)
ProgressBar(persist=False, desc="Train evaluation").attach(val_evaluator)
# Log val metrics:
tb_logger.attach(val_evaluator,
log_handler=OutputHandler(tag="val",
metric_names=list(
metrics.keys()),
another_engine=trainer),
event_name=Events.EPOCH_COMPLETED)
# trainer.add_event_handler(Events.ITERATION_COMPLETED, TerminateOnNan())
# Store the best model
def default_score_fn(engine):
score = engine.state.metrics['Accuracy']
return score
best_model_handler = ModelCheckpoint(dirname=log_path,
filename_prefix="best",
n_saved=3,
score_name="val_acc",
score_function=default_score_fn)
val_evaluator.add_event_handler(Events.COMPLETED, best_model_handler, {
'model': model,
})
# Add early stopping
es_patience = es_patience
es_handler = EarlyStopping(patience=es_patience,
score_function=default_score_fn,
trainer=trainer)
val_evaluator.add_event_handler(Events.COMPLETED, es_handler)
setup_logger(es_handler._logger)
setup_logger(logging.getLogger("ignite.engine.engine.Engine"))
def empty_cuda_cache(engine):
torch.cuda.empty_cache()
import gc
gc.collect()
trainer.add_event_handler(Events.EPOCH_COMPLETED, empty_cuda_cache)
val_evaluator.add_event_handler(Events.COMPLETED, empty_cuda_cache)
trainer.run(train_loader, max_epochs=epochs)
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 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_fold(fold, args):
# loggers
logging_logger = args.logging_logger
if args.tb_log:
tb_logger = args.tb_logger
num_classes = utils.problem_class[args.problem_type]
# init model
model = eval(args.model)(in_channels=3, num_classes=num_classes, bn=False)
model = nn.DataParallel(model, device_ids=args.device_ids).cuda()
# transform for train/valid data
train_transform, valid_transform = get_transform(args.model)
# loss function
loss_func = LossMulti(num_classes, args.jaccard_weight)
if args.semi:
loss_func_semi = LossMultiSemi(num_classes, args.jaccard_weight, args.semi_loss_alpha, args.semi_method)
# train/valid filenames
train_filenames, valid_filenames = utils.trainval_split(args.train_dir, fold)
# DataLoader and Dataset args
train_shuffle = True
train_ds_kwargs = {
'filenames': train_filenames,
'problem_type': args.problem_type,
'transform': train_transform,
'model': args.model,
'mode': 'train',
'semi': args.semi,
}
valid_num_workers = args.num_workers
valid_batch_size = args.batch_size
if 'TAPNet' in args.model:
# for TAPNet, cancel default shuffle, use self-defined shuffle in torch.Dataset instead
train_shuffle = False
train_ds_kwargs['batch_size'] = args.batch_size
train_ds_kwargs['mf'] = args.mf
if args.semi == True:
train_ds_kwargs['semi_method'] = args.semi_method
train_ds_kwargs['semi_percentage'] = args.semi_percentage
# additional valid dataset kws
valid_ds_kwargs = {
'filenames': valid_filenames,
'problem_type': args.problem_type,
'transform': valid_transform,
'model': args.model,
'mode': 'valid',
}
if 'TAPNet' in args.model:
# in validation, num_workers should be set to 0 for sequences
valid_num_workers = 0
# in validation, batch_size should be set to 1 for sequences
valid_batch_size = 1
valid_ds_kwargs['mf'] = args.mf
# train dataloader
train_loader = DataLoader(
dataset=RobotSegDataset(**train_ds_kwargs),
shuffle=train_shuffle, # set to False to disable pytorch dataset shuffle
num_workers=args.num_workers,
batch_size=args.batch_size,
pin_memory=True
)
# valid dataloader
valid_loader = DataLoader(
dataset=RobotSegDataset(**valid_ds_kwargs),
shuffle=False, # in validation, no need to shuffle
num_workers=valid_num_workers,
batch_size=valid_batch_size, # in valid time. have to use one image by one
pin_memory=True
)
# optimizer
optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
# optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=0.9,
# weight_decay=args.weight_decay, nesterov=True)
# ignite trainer process function
def train_step(engine, batch):
# set model to train
model.train()
# clear gradients
optimizer.zero_grad()
# additional params to feed into model
add_params = {}
inputs = batch['input'].cuda(non_blocking=True)
with torch.no_grad():
targets = batch['target'].cuda(non_blocking=True)
# for TAPNet, add attention maps
if 'TAPNet' in args.model:
add_params['attmap'] = batch['attmap'].cuda(non_blocking=True)
outputs = model(inputs, **add_params)
loss_kwargs = {}
if args.semi:
loss_kwargs['labeled'] = batch['labeled']
if args.semi_method == 'rev_flow':
loss_kwargs['optflow'] = batch['optflow']
loss = loss_func_semi(outputs, targets, **loss_kwargs)
else:
loss = loss_func(outputs, targets, **loss_kwargs)
loss.backward()
optimizer.step()
return_dict = {
'output': outputs,
'target': targets,
'loss_kwargs': loss_kwargs,
'loss': loss.item(),
}
# for TAPNet, update attention maps after each iteration
if 'TAPNet' in args.model:
# output_classes and target_classes: <b, h, w>
output_softmax_np = torch.softmax(outputs, dim=1).detach().cpu().numpy()
# update attention maps
train_loader.dataset.update_attmaps(output_softmax_np, batch['abs_idx'].numpy())
return_dict['attmap'] = add_params['attmap']
return return_dict
# init trainer
trainer = engine.Engine(train_step)
# lr scheduler and handler
# cyc_scheduler = optim.lr_scheduler.CyclicLR(optimizer, args.lr / 100, args.lr)
# lr_scheduler = c_handlers.param_scheduler.LRScheduler(cyc_scheduler)
# trainer.add_event_handler(engine.Events.ITERATION_COMPLETED, lr_scheduler)
step_scheduler = optim.lr_scheduler.StepLR(optimizer,
step_size=args.lr_decay_epochs, gamma=args.lr_decay)
lr_scheduler = c_handlers.param_scheduler.LRScheduler(step_scheduler)
trainer.add_event_handler(engine.Events.EPOCH_STARTED, lr_scheduler)
@trainer.on(engine.Events.STARTED)
def trainer_start_callback(engine):
logging_logger.info('training fold {}, {} train / {} valid files'. \
format(fold, len(train_filenames), len(valid_filenames)))
# resume training
if args.resume:
# ckpt for current fold fold_<fold>_model_<epoch>.pth
ckpt_dir = Path(args.ckpt_dir)
ckpt_filename = ckpt_dir.glob('fold_%d_model_[0-9]*.pth' % fold)[0]
res = re.match(r'fold_%d_model_(\d+).pth' % fold, ckpt_filename)
# restore epoch
engine.state.epoch = int(res.groups()[0])
# load model state dict
model.load_state_dict(torch.load(str(ckpt_filename)))
logging_logger.info('restore model [{}] from epoch {}.'.format(args.model, engine.state.epoch))
else:
logging_logger.info('train model [{}] from scratch'.format(args.model))
# record metrics history every epoch
engine.state.metrics_records = {}
@trainer.on(engine.Events.EPOCH_STARTED)
def trainer_epoch_start_callback(engine):
# log learning rate on pbar
train_pbar.log_message('model: %s, problem type: %s, fold: %d, lr: %.5f, batch size: %d' % \
(args.model, args.problem_type, fold, lr_scheduler.get_param(), args.batch_size))
# for TAPNet, change dataset schedule to random after the first epoch
if 'TAPNet' in args.model and engine.state.epoch > 1:
train_loader.dataset.set_dataset_schedule("shuffle")
@trainer.on(engine.Events.ITERATION_COMPLETED)
def trainer_iter_comp_callback(engine):
# logging_logger.info(engine.state.metrics)
pass
# monitor loss
# running average loss
train_ra_loss = imetrics.RunningAverage(output_transform=
lambda x: x['loss'], alpha=0.98)
train_ra_loss.attach(trainer, 'train_ra_loss')
# monitor train loss over epoch
if args.semi:
train_loss = imetrics.Loss(loss_func_semi, output_transform=lambda x: (x['output'], x['target'], x['loss_kwargs']))
else:
train_loss = imetrics.Loss(loss_func, output_transform=lambda x: (x['output'], x['target']))
train_loss.attach(trainer, 'train_loss')
# progress bar
train_pbar = c_handlers.ProgressBar(persist=True, dynamic_ncols=True)
train_metric_names = ['train_ra_loss']
train_pbar.attach(trainer, metric_names=train_metric_names)
# tensorboardX: log train info
if args.tb_log:
tb_logger.attach(trainer, log_handler=OptimizerParamsHandler(optimizer, 'lr'),
event_name=engine.Events.EPOCH_STARTED)
tb_logger.attach(trainer, log_handler=OutputHandler('train_iter', train_metric_names),
event_name=engine.Events.ITERATION_COMPLETED)
tb_logger.attach(trainer, log_handler=OutputHandler('train_epoch', ['train_loss']),
event_name=engine.Events.EPOCH_COMPLETED)
tb_logger.attach(trainer,
log_handler=WeightsScalarHandler(model, reduction=torch.norm),
event_name=engine.Events.ITERATION_COMPLETED)
# tb_logger.attach(trainer, log_handler=tb_log_train_vars,
# event_name=engine.Events.ITERATION_COMPLETED)
# ignite validator process function
def valid_step(engine, batch):
with torch.no_grad():
model.eval()
inputs = batch['input'].cuda(non_blocking=True)
targets = batch['target'].cuda(non_blocking=True)
# additional arguments
add_params = {}
# for TAPNet, add attention maps
if 'TAPNet' in args.model:
add_params['attmap'] = batch['attmap'].cuda(non_blocking=True)
# output logits
outputs = model(inputs, **add_params)
# loss
loss = loss_func(outputs, targets)
output_softmaxs = torch.softmax(outputs, dim=1)
output_argmaxs = output_softmaxs.argmax(dim=1)
# output_classes and target_classes: <b, h, w>
output_classes = output_argmaxs.cpu().numpy()
target_classes = targets.cpu().numpy()
# record current batch metrics
iou_mRecords = MetricRecord()
dice_mRecords = MetricRecord()
cm_b = np.zeros((num_classes, num_classes), dtype=np.uint32)
for output_class, target_class in zip(output_classes, target_classes):
# calculate metrics for each frame
# calculate using confusion matrix or dirctly using definition
cm = calculate_confusion_matrix_from_arrays(output_class, target_class, num_classes)
iou_mRecords.update_record(calculate_iou(cm))
dice_mRecords.update_record(calculate_dice(cm))
cm_b += cm
######## calculate directly using definition ##########
# iou_mRecords.update_record(iou_multi_np(target_class, output_class))
# dice_mRecords.update_record(dice_multi_np(target_class, output_class))
# accumulate batch metrics to engine state
engine.state.epoch_metrics['confusion_matrix'] += cm_b
engine.state.epoch_metrics['iou'].merge(iou_mRecords)
engine.state.epoch_metrics['dice'].merge(dice_mRecords)
return_dict = {
'loss': loss.item(),
'output': outputs,
'output_argmax': output_argmaxs,
'target': targets,
# for monitoring
'iou': iou_mRecords,
'dice': dice_mRecords,
}
if 'TAPNet' in args.model:
# for TAPNet, update attention maps after each iteration
valid_loader.dataset.update_attmaps(output_softmaxs.cpu().numpy(), batch['abs_idx'].numpy())
# for TAPNet, return extra internal values
return_dict['attmap'] = add_params['attmap']
# TODO: for TAPNet, return internal self-learned attention maps
return return_dict
# validator engine
validator = engine.Engine(valid_step)
# monitor loss
valid_ra_loss = imetrics.RunningAverage(output_transform=
lambda x: x['loss'], alpha=0.98)
valid_ra_loss.attach(validator, 'valid_ra_loss')
# monitor validation loss over epoch
valid_loss = imetrics.Loss(loss_func, output_transform=lambda x: (x['output'], x['target']))
valid_loss.attach(validator, 'valid_loss')
# monitor <data> mean metrics
valid_data_miou = imetrics.RunningAverage(output_transform=
lambda x: x['iou'].data_mean()['mean'], alpha=0.98)
valid_data_miou.attach(validator, 'mIoU')
valid_data_mdice = imetrics.RunningAverage(output_transform=
lambda x: x['dice'].data_mean()['mean'], alpha=0.98)
valid_data_mdice.attach(validator, 'mDice')
# show metrics on progress bar (after every iteration)
valid_pbar = c_handlers.ProgressBar(persist=True, dynamic_ncols=True)
valid_metric_names = ['valid_ra_loss', 'mIoU', 'mDice']
valid_pbar.attach(validator, metric_names=valid_metric_names)
# ## monitor ignite IoU (the same as iou we are using) ###
# cm = imetrics.ConfusionMatrix(num_classes,
# output_transform=lambda x: (x['output'], x['target']))
# imetrics.IoU(cm,
# ignore_index=0
# ).attach(validator, 'iou')
# # monitor ignite mean iou (over all classes even not exist in gt)
# mean_iou = imetrics.mIoU(cm,
# ignore_index=0
# ).attach(validator, 'mean_iou')
@validator.on(engine.Events.STARTED)
def validator_start_callback(engine):
pass
@validator.on(engine.Events.EPOCH_STARTED)
def validator_epoch_start_callback(engine):
engine.state.epoch_metrics = {
# directly use definition to calculate
'iou': MetricRecord(),
'dice': MetricRecord(),
'confusion_matrix': np.zeros((num_classes, num_classes), dtype=np.uint32),
}
# evaluate after iter finish
@validator.on(engine.Events.ITERATION_COMPLETED)
def validator_iter_comp_callback(engine):
pass
# evaluate after epoch finish
@validator.on(engine.Events.EPOCH_COMPLETED)
def validator_epoch_comp_callback(engine):
# log ignite metrics
# logging_logger.info(engine.state.metrics)
# ious = engine.state.metrics['iou']
# msg = 'IoU: '
# for ins_id, iou in enumerate(ious):
# msg += '{:d}: {:.3f}, '.format(ins_id + 1, iou)
# logging_logger.info(msg)
# logging_logger.info('nonzero mean IoU for all data: {:.3f}'.format(ious[ious > 0].mean()))
# log monitored epoch metrics
epoch_metrics = engine.state.epoch_metrics
######### NOTICE: Two metrics are available but different ##########
### 1. mean metrics for all data calculated by confusion matrix ####
'''
compared with using confusion_matrix[1:, 1:] in original code,
we use the full confusion matrix and only present non-background result
'''
confusion_matrix = epoch_metrics['confusion_matrix']# [1:, 1:]
ious = calculate_iou(confusion_matrix)
dices = calculate_dice(confusion_matrix)
mean_ious = np.mean(list(ious.values()))
mean_dices = np.mean(list(dices.values()))
std_ious = np.std(list(ious.values()))
std_dices = np.std(list(dices.values()))
logging_logger.info('mean IoU: %.3f, std: %.3f, for each class: %s' %
(mean_ious, std_ious, ious))
logging_logger.info('mean Dice: %.3f, std: %.3f, for each class: %s' %
(mean_dices, std_dices, dices))
### 2. mean metrics for all data calculated by definition ###
iou_data_mean = epoch_metrics['iou'].data_mean()
dice_data_mean = epoch_metrics['dice'].data_mean()
logging_logger.info('data (%d) mean IoU: %.3f, std: %.3f' %
(len(iou_data_mean['items']), iou_data_mean['mean'], iou_data_mean['std']))
logging_logger.info('data (%d) mean Dice: %.3f, std: %.3f' %
(len(dice_data_mean['items']), dice_data_mean['mean'], dice_data_mean['std']))
# record metrics in trainer every epoch
# trainer.state.metrics_records[trainer.state.epoch] = \
# {'miou': mean_ious, 'std_miou': std_ious,
# 'mdice': mean_dices, 'std_mdice': std_dices}
trainer.state.metrics_records[trainer.state.epoch] = \
{'miou': iou_data_mean['mean'], 'std_miou': iou_data_mean['std'],
'mdice': dice_data_mean['mean'], 'std_mdice': dice_data_mean['std']}
# log interal variables(attention maps, outputs, etc.) on validation
def tb_log_valid_iter_vars(engine, logger, event_name):
log_tag = 'valid_iter'
output = engine.state.output
batch_size = output['output'].shape[0]
res_grid = tvutils.make_grid(torch.cat([
output['output_argmax'].unsqueeze(1),
output['target'].unsqueeze(1),
]), padding=2,
normalize=False, # show origin image
nrow=batch_size).cpu()
logger.writer.add_image(tag='%s (outputs, targets)' % (log_tag), img_tensor=res_grid)
if 'TAPNet' in args.model:
# log attention maps and other internal values
inter_vals_grid = tvutils.make_grid(torch.cat([
output['attmap'],
]), padding=2, normalize=True, nrow=batch_size).cpu()
logger.writer.add_image(tag='%s internal vals' % (log_tag), img_tensor=inter_vals_grid)
def tb_log_valid_epoch_vars(engine, logger, event_name):
log_tag = 'valid_iter'
# log monitored epoch metrics
epoch_metrics = engine.state.epoch_metrics
confusion_matrix = epoch_metrics['confusion_matrix']# [1:, 1:]
ious = calculate_iou(confusion_matrix)
dices = calculate_dice(confusion_matrix)
mean_ious = np.mean(list(ious.values()))
mean_dices = np.mean(list(dices.values()))
logger.writer.add_scalar('mIoU', mean_ious, engine.state.epoch)
logger.writer.add_scalar('mIoU', mean_dices, engine.state.epoch)
if args.tb_log:
# log internal values
tb_logger.attach(validator, log_handler=tb_log_valid_iter_vars,
event_name=engine.Events.ITERATION_COMPLETED)
tb_logger.attach(validator, log_handler=tb_log_valid_epoch_vars,
event_name=engine.Events.EPOCH_COMPLETED)
# tb_logger.attach(validator, log_handler=OutputHandler('valid_iter', valid_metric_names),
# event_name=engine.Events.ITERATION_COMPLETED)
tb_logger.attach(validator, log_handler=OutputHandler('valid_epoch', ['valid_loss']),
event_name=engine.Events.EPOCH_COMPLETED)
# score function for model saving
ckpt_score_function = lambda engine: \
np.mean(list(calculate_iou(engine.state.epoch_metrics['confusion_matrix']).values()))
# ckpt_score_function = lambda engine: engine.state.epoch_metrics['iou'].data_mean()['mean']
ckpt_filename_prefix = 'fold_%d' % fold
# model saving handler
model_ckpt_handler = handlers.ModelCheckpoint(
dirname=args.model_save_dir,
filename_prefix=ckpt_filename_prefix,
score_function=ckpt_score_function,
create_dir=True,
require_empty=False,
save_as_state_dict=True,
atomic=True)
validator.add_event_handler(event_name=engine.Events.EPOCH_COMPLETED,
handler=model_ckpt_handler,
to_save={
'model': model,
})
# early stop
# trainer=trainer, but should be handled by validator
early_stopping = handlers.EarlyStopping(patience=args.es_patience,
score_function=ckpt_score_function,
trainer=trainer
)
validator.add_event_handler(event_name=engine.Events.EPOCH_COMPLETED,
handler=early_stopping)
# evaluate after epoch finish
@trainer.on(engine.Events.EPOCH_COMPLETED)
def trainer_epoch_comp_callback(engine):
validator.run(valid_loader)
trainer.run(train_loader, max_epochs=args.max_epochs)
if args.tb_log:
# close tb_logger
tb_logger.close()
return trainer.state.metrics_records
def run(output_path, config):
distributed = dist.is_available() and dist.is_initialized()
rank = dist.get_rank() if distributed else 0
manual_seed(config["seed"] + rank)
# Setup dataflow, model, optimizer, criterion
train_loader, test_loader = utils.get_dataflow(config, distributed)
model, optimizer = utils.get_model_optimizer(config, distributed)
criterion = nn.CrossEntropyLoss().to(utils.device)
le = len(train_loader)
milestones_values = [
(0, 0.0),
(le * config["num_warmup_epochs"], config["learning_rate"]),
(le * config["num_epochs"], 0.0),
]
lr_scheduler = PiecewiseLinear(optimizer, param_name="lr", milestones_values=milestones_values)
# Setup Ignite trainer:
# - let's define training step
# - add other common handlers:
# - TerminateOnNan,
# - handler to setup learning rate scheduling,
# - ModelCheckpoint
# - RunningAverage` on `train_step` output
# - Two progress bars on epochs and optionally on iterations
def train_step(engine, batch):
x = convert_tensor(batch[0], device=utils.device, non_blocking=True)
y = convert_tensor(batch[1], device=utils.device, non_blocking=True)
model.train()
# Supervised part
y_pred = model(x)
loss = criterion(y_pred, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
return {
"batch loss": loss.item(),
}
if config["deterministic"] and rank == 0:
print("Setup deterministic trainer")
trainer = Engine(train_step) if not config["deterministic"] else DeterministicEngine(train_step)
train_sampler = train_loader.sampler if distributed else None
to_save = {"trainer": trainer, "model": model, "optimizer": optimizer, "lr_scheduler": lr_scheduler}
metric_names = [
"batch loss",
]
common.setup_common_training_handlers(
trainer,
train_sampler=train_sampler,
to_save=to_save,
save_every_iters=config["checkpoint_every"],
output_path=output_path,
lr_scheduler=lr_scheduler,
output_names=metric_names,
with_pbar_on_iters=config["display_iters"],
log_every_iters=10,
)
if rank == 0:
# Setup Tensorboard logger - wrapper on SummaryWriter
tb_logger = TensorboardLogger(log_dir=output_path)
# Attach logger to the trainer and log trainer's metrics (stored in trainer.state.metrics) every iteration
tb_logger.attach(
trainer,
log_handler=OutputHandler(tag="train", metric_names=metric_names),
event_name=Events.ITERATION_COMPLETED,
)
# log optimizer's parameters: "lr" every iteration
tb_logger.attach(
trainer, log_handler=OptimizerParamsHandler(optimizer, param_name="lr"), event_name=Events.ITERATION_STARTED
)
# Let's now setup evaluator engine to perform model's validation and compute metrics
metrics = {
"accuracy": Accuracy(device=utils.device if distributed else None),
"loss": Loss(criterion, device=utils.device if distributed else None),
}
# We define two evaluators as they wont have exactly similar roles:
# - `evaluator` will save the best model based on validation score
evaluator = create_supervised_evaluator(model, metrics=metrics, device=utils.device, non_blocking=True)
train_evaluator = create_supervised_evaluator(model, metrics=metrics, device=utils.device, non_blocking=True)
def run_validation(engine):
train_evaluator.run(train_loader)
evaluator.run(test_loader)
trainer.add_event_handler(Events.EPOCH_STARTED(every=config["validate_every"]), run_validation)
trainer.add_event_handler(Events.COMPLETED, run_validation)
if rank == 0:
# Setup progress bar on evaluation engines
if config["display_iters"]:
ProgressBar(persist=False, desc="Train evaluation").attach(train_evaluator)
ProgressBar(persist=False, desc="Test evaluation").attach(evaluator)
# Let's log metrics of `train_evaluator` stored in `train_evaluator.state.metrics` when validation run is done
tb_logger.attach(
train_evaluator,
log_handler=OutputHandler(
tag="train", metric_names="all", global_step_transform=global_step_from_engine(trainer)
),
event_name=Events.COMPLETED,
)
# Let's log metrics of `evaluator` stored in `evaluator.state.metrics` when validation run is done
tb_logger.attach(
evaluator,
log_handler=OutputHandler(
tag="test", metric_names="all", global_step_transform=global_step_from_engine(trainer)
),
event_name=Events.COMPLETED,
)
# Store 3 best models by validation accuracy:
common.save_best_model_by_val_score(
output_path, evaluator, model=model, metric_name="accuracy", n_saved=3, trainer=trainer, tag="test"
)
# Optionally log model gradients
if config["log_model_grads_every"] is not None:
tb_logger.attach(
trainer,
log_handler=GradsHistHandler(model, tag=model.__class__.__name__),
event_name=Events.ITERATION_COMPLETED(every=config["log_model_grads_every"]),
)
# In order to check training resuming we can emulate a crash
if config["crash_iteration"] is not None:
@trainer.on(Events.ITERATION_STARTED(once=config["crash_iteration"]))
def _(engine):
raise Exception("STOP at iteration: {}".format(engine.state.iteration))
resume_from = config["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())
print("Resume from a checkpoint: {}".format(checkpoint_fp.as_posix()))
checkpoint = torch.load(checkpoint_fp.as_posix())
Checkpoint.load_objects(to_load=to_save, checkpoint=checkpoint)
try:
trainer.run(train_loader, max_epochs=config["num_epochs"])
except Exception as e:
import traceback
print(traceback.format_exc())
if rank == 0:
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 trainer(
train_batch,
evaluate_batch,
evaluate_data_loaders,
metrics,
optimizers,
):
'''
Create standard trainer with evaluators.
Parameters
----------
train_batch : function
function that trains on given batch
evaluate_batch : function
function that evaluates a given batch
evaluate_data_loaders: list
data loaders that yield batches to evaluate on
metrics : dict
dict with one dict each for 'train' and evaluate data loader. Wrap a
metric with trainer.Progress to show in progress bar.
optimizers : dict
dict with optimizers for logging
Returns
-------
tuple
trainer engine
list of evaluator engines
tensorboard logger
'''
trainer = ignite.engine.Engine(train_batch)
for name, metric in metrics.get(PROGRESS_DESC, dict()).items():
metric.attach(trainer, name)
for name, metric in metrics.get(TRAIN_DESC, dict()).items():
metric.attach(trainer, name)
evaluators = {
evaluator_name: ignite.engine.Engine(evaluate_batch)
for evaluator_name in evaluate_data_loaders.keys()
}
for evaluator_name, evaluator in evaluators.items():
for metric_name, metric in metrics[evaluator_name].items():
metric.attach(evaluator, metric_name)
tensorboard_logger = TensorboardLogger(log_dir='tb')
EpochLogger().attach(trainer)
# Order of attaching progress bars is important for vscode / atom
ProgressBar(desc=TRAIN_DESC).attach(trainer,
metric_names=list(
metrics.get(PROGRESS_DESC,
dict()).keys()))
tensorboard_logger.attach(
trainer,
OutputHandler(
tag=PROGRESS_DESC,
metric_names=list(metrics.get(PROGRESS_DESC, dict()).keys()),
),
Events.ITERATION_COMPLETED,
)
MetricsLogger(TRAIN_DESC).attach(trainer,
metrics.get(TRAIN_DESC, dict()).keys())
tensorboard_logger.attach(
trainer,
OutputHandler(
tag=TRAIN_DESC,
metric_names=list(metrics.get(TRAIN_DESC, dict()).keys()),
),
Events.ITERATION_COMPLETED,
)
def run_evaluator(evaluator_desc):
return lambda engine: evaluators[evaluator_desc].run(
evaluate_data_loaders[evaluator_desc])
for evaluator_desc, evaluator in evaluators.items():
evaluator_metric_names = list(metrics[evaluator_desc].keys())
trainer.add_event_handler(
Events.EPOCH_COMPLETED,
run_evaluator(evaluator_desc),
)
ProgressBar(desc=evaluator_desc).attach(evaluator)
MetricsLogger(evaluator_desc).attach(evaluator, evaluator_metric_names)
tensorboard_logger.attach(
evaluator,
OutputHandler(
tag=evaluator_desc,
metric_names=evaluator_metric_names,
global_step_transform=global_step_from_engine(trainer),
),
Events.EPOCH_COMPLETED,
)
if type(optimizers) is not dict:
optimizers = dict(optimizer=optimizers)
for name, optimizer in optimizers.items():
tensorboard_logger.attach(
trainer,
log_handler=OptimizerParamsHandler(
tag=f'{TRAIN_DESC}/{name}',
param_name='lr',
optimizer=optimizer,
),
event_name=Events.ITERATION_COMPLETED,
)
return trainer, evaluators, tensorboard_logger
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():
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 train(epochs=500,
batch_size=32,
bptt_len=70,
lr=0.00025,
log_steps=200,
clip_grad=0.25,
log_dir="experiments"):
###################################################################
# Dataset
###################################################################
wt = wikitext103(batch_size=batch_size, bptt_len=bptt_len)
# wt = wikitext2(batch_size=batch_size, bptt_len=bptt_len)
###################################################################
# Configs
###################################################################
embedding_config = DropEmbedding.Hyperparams(len(wt.text_field.vocab) + 3,
ninp=512)
encoder_config = TransformerEncoder.Hyperparams(
att_num_units=[512, 512, 512, 512, 512, 512], max_ext=384)
###################################################################
# Models
###################################################################
base_embedding = DropEmbedding(embedding_config)
embedding = TransformerEmbedding(embedding=base_embedding,
max_length=bptt_len,
embedding_size=embedding_config.ninp,
use_positional_embedding=False)
encoder = TransformerEncoder(encoder_config)
model = TransformerLanguageModel(embedding, encoder)
model.init_weight()
###################################################################
# Loss
###################################################################
criterion = lm_criterion(in_features=encoder_config.att_num_units[-1],
vocab_size=len(wt.text_field.vocab))
###################################################################
# Parameters + Train ops
###################################################################
parameters = (list(model.parameters()) + list(criterion.parameters()))
tot_params = 0
for p in parameters:
tot_params += reduce(lambda x, y: x * y, p.size())
print("Total Parameters: ", tot_params)
opt = optim.Adam(parameters, lr=lr)
model.to(DEVICE)
criterion.to(DEVICE)
###################################################################
# Train + Evaluation
###################################################################
def train_step(engine, batch):
model.train()
opt.zero_grad()
text = batch.text.to(DEVICE).t().contiguous()
target = batch.target.to(DEVICE).t().contiguous()
out, out_past = model(text, engine.state.train_past)
engine.state.train_past = out_past
raw_loss = criterion(out.view(-1, out.size(2)), target.view(-1))
loss = raw_loss[1]
loss.backward()
nn.utils.clip_grad_norm_(parameters, clip_grad)
opt.step()
return {"train_loss": loss.item(), "train_ppl": loss.exp().item()}
def eval_step(engine, batch):
model.eval()
if not hasattr(engine.state, "eval_past"):
engine.state.eval_past = None
with torch.no_grad():
text = batch.text.to(DEVICE).t().contiguous()
target = batch.target.to(DEVICE).t().contiguous()
out, out_past = model(text, engine.state.eval_past)
engine.state.eval_past = out_past
raw_loss = criterion(out.view(-1, out.size(2)), target.view(-1))
loss = raw_loss[1]
return {"val_loss": loss.item()}
train_engine = Engine(train_step)
eval_engine = Engine(eval_step)
def reset_state(engine):
engine.state.train_past = None
def run_eval(_):
print("start running eval")
eval_engine.run(wt.valid_iter)
metrics = eval_engine.state.metrics
print("Validation loss: ", metrics["val_loss"], ", ppl: ",
np.exp(metrics["val_loss"]))
train_engine.add_event_handler(Events.EPOCH_STARTED, reset_state)
train_engine.add_event_handler(Events.EPOCH_COMPLETED, run_eval)
###################################################################
# LR Scheduler
###################################################################
cosine_scheduler = CosineAnnealingScheduler(opt.param_groups[0],
"lr",
0.0,
2.5e-4,
cycle_size=len(wt.train_iter))
warmup_scheduler = create_lr_scheduler_with_warmup(cosine_scheduler, 0.0,
2.5e-4, 200)
train_engine.add_event_handler(Events.ITERATION_STARTED, warmup_scheduler)
###################################################################
# Metrics
###################################################################
RunningAverage(output_transform=lambda x: x["train_ppl"]).attach(
train_engine, "train_ppl")
RunningAverage(output_transform=lambda x: x["train_loss"]).attach(
train_engine, "train_loss")
RunningAverage(output_transform=lambda x: x["val_loss"]).attach(
eval_engine, "val_loss")
progress_bar = ProgressBar(persist=True)
progress_bar.attach(train_engine, ["train_ppl", "train_loss"])
progress_bar_val = ProgressBar(persist=True)
progress_bar_val.attach(eval_engine, ["val_loss"])
###################################################################
# Tensorboard
###################################################################
tb_logger = TensorboardLogger(log_dir=log_dir)
def stepn_logger(num_steps, handler):
def logger_runner(engine, log_handler, event_name):
if engine.state.iteration % num_steps == 0:
handler(engine, log_handler, event_name)
return logger_runner
tb_logger.attach(train_engine,
log_handler=stepn_logger(
log_steps,
OutputHandler(tag="training",
output_transform=lambda loss: loss)),
event_name=Events.ITERATION_COMPLETED)
tb_logger.attach(eval_engine,
log_handler=OutputHandler(
tag="validation",
output_transform=lambda loss: loss,
another_engine=train_engine),
event_name=Events.EPOCH_COMPLETED)
tb_logger.attach(train_engine,
log_handler=stepn_logger(log_steps,
OptimizerParamsHandler(opt)),
event_name=Events.ITERATION_STARTED)
tb_logger.attach(train_engine,
log_handler=stepn_logger(log_steps,
WeightsScalarHandler(model)),
event_name=Events.ITERATION_COMPLETED)
tb_logger.attach(train_engine,
log_handler=stepn_logger(log_steps,
GradsScalarHandler(model)),
event_name=Events.ITERATION_COMPLETED)
tb_logger.attach(train_engine,
log_handler=stepn_logger(500, WeightsHistHandler(model)),
event_name=Events.ITERATION_COMPLETED)
tb_logger.attach(train_engine,
log_handler=stepn_logger(500, GradsHistHandler(model)),
event_name=Events.ITERATION_COMPLETED)
try:
train_engine.run(wt.train_iter, max_epochs=epochs)
except Exception:
pass
finally:
tb_logger.close()
def attach_handlers(run, model, optimizer, learning_rule, trainer, evaluator, train_loader, val_loader, params):
# Metrics
UnitConvergence(model[0], learning_rule.norm).attach(trainer.engine, 'unit_conv')
# Tqdm logger
pbar = ProgressBar(persist=True, bar_format=config.IGNITE_BAR_FORMAT)
pbar.attach(trainer.engine, metric_names='all')
tqdm_logger = TqdmLogger(pbar=pbar)
# noinspection PyTypeChecker
tqdm_logger.attach_output_handler(
evaluator.engine,
event_name=Events.COMPLETED,
tag="validation",
global_step_transform=global_step_from_engine(trainer.engine),
)
# Evaluator
evaluator.attach(trainer.engine, Events.EPOCH_COMPLETED(every=100), train_loader, val_loader)
# Learning rate scheduling
lr_scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer=optimizer,
lr_lambda=lambda epoch: 1 - epoch / params['epochs'])
lr_scheduler = LRScheduler(lr_scheduler)
trainer.engine.add_event_handler(Events.EPOCH_COMPLETED, lr_scheduler)
# Early stopping
mc_handler = ModelCheckpoint(config.MODELS_DIR, run.replace('/', '-'), n_saved=1, create_dir=True,
require_empty=False,
global_step_transform=global_step_from_engine(trainer.engine))
trainer.engine.add_event_handler(Events.EPOCH_COMPLETED, mc_handler, {'m': model})
# Create a TensorBoard logger
tb_logger = TensorboardLogger(log_dir=os.path.join(config.TENSORBOARD_DIR, run))
images, labels = next(iter(train_loader))
tb_logger.writer.add_graph(copy.deepcopy(model).cpu(), images)
tb_logger.writer.add_hparams(params, {})
# noinspection PyTypeChecker
tb_logger.attach_output_handler(
evaluator.engine,
event_name=Events.COMPLETED,
tag="validation",
metric_names="all",
global_step_transform=global_step_from_engine(trainer.engine),
)
# noinspection PyTypeChecker
tb_logger.attach_output_handler(
trainer.engine,
event_name=Events.EPOCH_COMPLETED,
tag="train",
metric_names=["unit_conv"]
)
input_shape = tuple(next(iter(train_loader))[0].shape[1:])
tb_logger.attach(trainer.engine,
log_handler=WeightsImageHandler(model, input_shape),
event_name=Events.EPOCH_COMPLETED)
tb_logger.attach(trainer.engine, log_handler=OptimizerParamsHandler(optimizer), event_name=Events.EPOCH_STARTED)
# tb_logger.attach(trainer.engine,
# log_handler=WeightsScalarHandler(model, layer_names=['linear1', 'linear2']),
# event_name=Events.EPOCH_COMPLETED)
# tb_logger.attach(trainer.engine,
# log_handler=WeightsHistHandler(model, layer_names=['linear1', 'linear2']),
# event_name=Events.EPOCH_COMPLETED)
# tb_logger.attach(trainer.engine,
# log_handler=ActivationsHistHandler(model, layer_names=['batch_norm', 'repu']),
# event_name=Events.ITERATION_COMPLETED)
# tb_logger.attach(trainer.engine,
# log_handler=NumActivationsScalarHandler(model, layer_names=['repu']),
# event_name=Events.ITERATION_COMPLETED)
# tb_logger.attach(trainer.engine,
# log_handler=ActivationsScalarHandler(model, reduction=torch.mean,
# layer_names=['batch_norm', 'repu']),
# event_name=Events.ITERATION_COMPLETED)
# tb_logger.attach(trainer.engine,
# log_handler=ActivationsScalarHandler(model, reduction=torch.std,
# layer_names=['batch_norm', 'repu']),
# event_name=Events.ITERATION_COMPLETED)
return tb_logger
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()
output_transform=lambda output: {'loss': output['loss']},
metric_names=[f"gpu:{args.gpu} mem(%)"])
# FIRE
tb_logger = TensorboardLogger(log_dir=TENSORBOARD_RUN_LOG_DIR_PATH)
tb_logger.attach(
trainer,
log_handler=OutputHandler(
tag='training',
output_transform=lambda output: {'loss': output['loss']}),
event_name=Events.ITERATION_COMPLETED(
every=LOG_TRAINING_PROGRESS_EVERY_N))
tb_logger.attach(
evaluator,
log_handler=OutputHandler(
tag='validation',
metric_names='all',
global_step_transform=global_step_from_engine(trainer)),
event_name=Events.EPOCH_COMPLETED)
tb_logger.attach(trainer,
log_handler=OptimizerParamsHandler(opt),
event_name=Events.ITERATION_STARTED)
tb_logger.attach(trainer,
log_handler=WeightsHistHandler(mude),
event_name=Events.EPOCH_COMPLETED)
trainer.run(train_ld, max_epochs=EPOCHS)
tb_logger.close()
torch.save(mude.state_dict(),
CHECKPOINTS_RUN_DIR_PATH.joinpath(f"{RUN_NAME}-last.pth"))
def setup(self, training_metrics):
def metric_name(n) -> str:
if n.endswith('Accuracy'):
n = 'acc'
else:
n = n[:-6] if n.endswith('Metric') else n
return n
def print_metrics(metrics) -> str:
rv = ''
metric_keys = sorted(k for k in metrics)
for k in metric_keys:
if k == 'Accuracy':
rv += f'{metric_name(k)}: {metrics[k]:.3}'
else:
rv += f'{metric_name(k)}: {metrics[k]:.6}'
return rv
if self.seed:
set_seed_everywhere(self.seed, self.cuda)
pbar = ProgressBar()
names = []
for k, v in training_metrics.items():
name = f'r{k}'
names.append(name)
RunningAverage(v).attach(self.trainer, name)
RunningAverage(None,
output_transform=lambda x: x[-1] * self.
loss_accumulation_steps).attach(self.trainer, 'rloss')
names.append('rloss')
pbar.attach(self.trainer, names)
pbar = ProgressBar()
pbar.attach(self.evaluator)
# A few events handler. To add / modify the events handler, you need to extend the __init__ method of RunnerABC
# Ignite provides the necessary abstractions and a furnished repository of useful tools
@self.trainer.on(Events.EPOCH_COMPLETED)
def log_validation_results(trainer):
self.evaluator.run(self.dataset_splits.val_data_loader())
metrics = self.evaluator.state.metrics
logger.info(
f"Validation Results - Epoch: {trainer.state.epoch} {print_metrics(metrics)}"
)
if self.scheduler:
self.scheduler.step(
metrics[self.loss_metric.__class__.__name__])
@self.trainer.on(Events.COMPLETED)
def log_test_results(trainer):
self.evaluator.run(self.dataset_splits.test_data_loader())
metrics = self.evaluator.state.metrics
logger.info(
f"Test Results - Epoch: {trainer.state.epoch} {print_metrics(metrics)}"
)
if self.tensorboard_logs:
tb_logger = TensorboardLogger(log_dir=self.tensorboard_logs)
tb_logger.attach(self.trainer,
log_handler=OutputHandler(
tag="training",
output_transform=lambda loss: {'loss': loss}),
event_name=Events.ITERATION_COMPLETED)
tb_logger.attach(self.evaluator,
log_handler=OutputHandler(
tag="validation",
metric_names=["LossMetric"],
another_engine=self.trainer),
event_name=Events.EPOCH_COMPLETED)
tb_logger.attach(self.trainer,
log_handler=OptimizerParamsHandler(
self.optimizer),
event_name=Events.ITERATION_STARTED)
tb_logger.attach(self.trainer,
log_handler=WeightsScalarHandler(self.model),
event_name=Events.ITERATION_COMPLETED)
tb_logger.attach(self.trainer,
log_handler=WeightsHistHandler(self.model),
event_name=Events.EPOCH_COMPLETED)
tb_logger.attach(self.trainer,
log_handler=GradsScalarHandler(self.model),
event_name=Events.ITERATION_COMPLETED)
# This is important to close the tensorboard file logger
@self.trainer.on(Events.COMPLETED)
def end_tensorboard(trainer):
logger.info("Training completed")
tb_logger.close()
if self.embeddings_name:
@self.trainer.on(Events.COMPLETED)
def log_embeddings(trainer):
if hasattr(self.model, self.embeddings_name) and hasattr(
self.dataset_splits, "vectorizer"):
logger.info(
f"Logging embeddings ({self.embeddings_name}) to Tensorboard!"
)
embeddings = getattr(self.model,
self.embeddings_name).weight.data
metadata = [
str(self.dataset_splits.vectorizer.data_vocab.
_id2token[token_index]).encode('utf-8')
for token_index in range(embeddings.shape[0])
]
self.writer.add_embedding(
mat=embeddings,
metadata=metadata,
global_step=self.trainer.state.epoch)
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():
parser = ArgumentParser()
parser.add_argument("--local_rank", type=int, default=-1)
args = parser.parse_args()
device = torch.device("cuda" if torch.cuda.device_count() > 1 else "cpu")
model = GPT2DoubleHeadsModel.from_pretrained('gpt2')
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
DISTRIBUTED = args.local_rank != -1
if DISTRIBUTED and torch.distributed.is_available():
print("Distributed")
torch.cuda.set_device(args.local_rank)
device = torch.device("cuda", args.local_rank)
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
#BATCH_SIZE *= 2
def average_distributed_scalar(scalar):
if (not DISTRIBUTED):
return scalar
scalar_t = torch.tensor(
scalar, dtype=torch.float,
device=device) / torch.distributed.get_world_size()
torch.distributed.all_reduce(scalar_t,
op=torch.distributed.ReduceOp.SUM)
return scalar_t.item()
optimizer = AdamW(model.parameters(), lr=6.25e-5)
ds = dataloader.Conv_GPT2_DataClass(tokenizer)
v_ds = dataloader.Conv_GPT2_DataClass(tokenizer, dev=True)
orig_added_tokens = len(tokenizer.encoder)
num_added_tokens = tokenizer.add_special_tokens(
dataloader.ATTR_SPECIAL_TOKENS)
if (num_added_tokens > 0):
model.resize_token_embeddings(new_num_tokens=orig_added_tokens +
num_added_tokens)
model = model.to(device)
train_sampler = torch.utils.data.distributed.DistributedSampler(
ds) if DISTRIBUTED else None
valid_sampler = torch.utils.data.distributed.DistributedSampler(
v_ds) if DISTRIBUTED else None
dl = DataLoader(ds,
sampler=train_sampler,
batch_size=BATCH_SIZE,
shuffle=not DISTRIBUTED)
v_dl = DataLoader(v_ds, sampler=valid_sampler, shuffle=False)
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"]),
})
metrics["average_ppl"] = MetricsLambda(math.exp, metrics["average_nll"])
def update(engine, batch):
model.train()
batch = tuple(t.to(device) for t in batch)
lm_loss, *_ = model(batch[0],
token_type_ids=batch[1],
lm_labels=batch[2])
loss = lm_loss / ITERATION_STEP
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
if engine.state.iteration % ITERATION_STEP == 0:
optimizer.step()
optimizer.zero_grad()
return loss.item()
def inference(engine, batch):
model.eval()
with torch.no_grad():
batch = tuple(t.to(device) for t in batch)
input_ids, token_type_ids, lm_labels = batch
model_outputs = model(input_ids, token_type_ids=token_type_ids)
lm_logits = model_outputs[0]
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)
scheduler = PiecewiseLinear(optimizer, "lr",
[(0, 6.25e-5),
(EPOCHS * len(ds) // BATCH_SIZE, 0.0)])
trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
# trainer.add_event_handler(Events.STARTED, lambda _: evaluator.run(v_dl))
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))
RunningAverage(output_transform=lambda x: x).attach(trainer, "loss")
for name, metric in metrics.items():
metric.attach(evaluator, name)
if (args.local_rank in [0, -1]):
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('./checkpoint',
'_checkpoint',
n_saved=3)
trainer.add_event_handler(Events.EPOCH_COMPLETED, checkpoint_handler,
{'gpt2_qg': getattr(model, 'module', model)})
getattr(model, 'module', model).config.to_json_file(
os.path.join('./checkpoint', 'config'))
tokenizer.save_pretrained('./checkpoint')
trainer.run(dl, max_epochs=EPOCHS)
if (args.local_rank in [0, -1]):
tb_logger.close()