def create_optimizer_and_scheduler(self, num_training_steps: int):
"""
Setup the optimizer and the learning rate scheduler.
We provide a reasonable default that works well. If you want to use something else, you can pass a tuple in the
TFTrainer's init through :obj:`optimizers`, or subclass and override this method.
"""
if not self.optimizer and not self.lr_scheduler:
self.optimizer, self.lr_scheduler = create_optimizer(
self.args.learning_rate,
num_training_steps,
self.args.warmup_steps,
adam_beta1=self.args.adam_beta1,
adam_beta2=self.args.adam_beta2,
adam_epsilon=self.args.adam_epsilon,
weight_decay_rate=self.args.weight_decay,
)
def get_optimizer(num_train_examples, options):
steps_per_epoch = ceil(num_train_examples / options.batch_size)
num_train_steps = steps_per_epoch * options.epochs
num_warmup_steps = floor(num_train_steps * options.warmup_proportion)
# Mostly defaults from transformers.optimization_tf
optimizer, lr_scheduler = create_optimizer(
options.lr,
num_train_steps=num_train_steps,
num_warmup_steps=num_warmup_steps,
min_lr_ratio=0.0,
adam_beta1=0.9,
adam_beta2=0.999,
adam_epsilon=1e-8,
weight_decay_rate=0.01,
power=1.0,
)
return optimizer
def get_optimizer(lr, epochs, batch_size, warmup_proportion,
num_train_examples):
from transformers.optimization_tf import create_optimizer
steps_per_epoch = math.ceil(num_train_examples / batch_size)
num_train_steps = steps_per_epoch * epochs
num_warmup_steps = math.floor(num_train_steps * warmup_proportion)
# Mostly defaults from transformers.optimization_tf
optimizer, lr_schedule = create_optimizer(
lr,
num_train_steps=num_train_steps,
num_warmup_steps=num_warmup_steps,
min_lr_ratio=0.0,
adam_beta1=0.9,
adam_beta2=0.999,
adam_epsilon=1e-8,
weight_decay_rate=0.01,
power=1.0,
)
return optimizer, lr_schedule
def get_optimizers(
self,
) -> Tuple[tf.keras.optimizers.Optimizer,
tf.keras.optimizers.schedules.LearningRateSchedule]:
"""
Setup the optimizer and the learning rate scheduler.
We provide a reasonable default that works well.
If you want to use something else, you can pass a tuple in the Trainer's init,
or override this method in a subclass.
"""
if self.optimizers is not None:
return self.optimizers
optimizer, scheduler = create_optimizer(
self.args.learning_rate,
self.train_steps * self.args.num_train_epochs,
self.args.warmup_steps,
adam_epsilon=self.args.adam_epsilon,
weight_decay_rate=self.args.weight_decay,
)
return optimizer, scheduler
def __init__(self, directory, settings, **kwargs):
self.dir = directory
self.model = None
self.tokenizer = None
self.settings = {
"model_file": "model.tf",
"vocab_file": "vocab",
"model": 'jplu/tf-xlm-roberta-base',
"batch_size": 16,
"maxlen": 150,
"n_classes": 2,
"epochs": 10,
"steps_per_epoch": 40000,
"patience": 3,
"dropout": 0.1,
"n_hidden": 2048,
"activation": 'relu',
"loss": "binary_crossentropy",
"lr": 2e-6,
"decay_rate": 0.1,
"warmup_steps": 1000,
"clipnorm": 1.0,
**settings,
}
scheduler = InverseTimeDecay(self.settings["lr"],
decay_steps=32.0,
decay_rate=0.1)
self.settings["scheduler"] = scheduler
optimizer, scheduler = create_optimizer(
self.settings["lr"],
self.settings["steps_per_epoch"] * self.settings["epochs"],
self.settings["warmup_steps"],
weight_decay_rate=self.settings["decay_rate"])
self.settings["scheduler"] = scheduler
self.settings["optimizer"] = optimizer
def main():
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
try:
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
print(gpu)
except RuntimeError as e:
print(e)
parser = argparse.ArgumentParser()
parser.add_argument('--casing',
type=str,
default="bert-base-uncased",
help='BERT model')
parser.add_argument('--bottleneck_size',
type=int,
default=64,
help='Bottleneck size of adapters')
parser.add_argument('--non_linearity',
type=str,
default='gelu_new',
help='non_linearity function in adapters')
parser.add_argument('--task', type=str, default='mrpc', help='GLUE task')
parser.add_argument('--batch_size',
type=int,
default=32,
help='batch size')
parser.add_argument('--epochs',
type=int,
default=10,
help='The number of training epochs')
parser.add_argument('--max_seq_length',
type=int,
default=128,
help='max sequence length')
parser.add_argument('--learning_rate',
type=float,
default=3e-4,
help='learning rate')
parser.add_argument('--warmup_ratio',
type=float,
default=0.1,
help='warmup ratio')
parser.add_argument('--saved_models_dir',
type=str,
default='saved_models',
help='save directory')
args = parser.parse_args()
if not os.path.isdir(args.saved_models_dir):
os.mkdir(args.saved_models_dir)
if not os.path.isdir(os.path.join(args.saved_models_dir, args.task)):
os.mkdir(os.path.join(args.saved_models_dir, args.task))
if args.task == "sst-2":
TFDS_TASK = "sst2"
postfix = ""
elif args.task == "sts-b":
TFDS_TASK = "stsb"
postfix = ""
elif args.task == "mnli":
TFDS_TASK = "mnli"
postfix = "_matched"
elif args.task == "mnli-mm":
TFDS_TASK = "mnli"
postfix = "_mismatched"
else:
TFDS_TASK = args.task
postfix = ""
TFDS_TASK = "glue/" + TFDS_TASK
num_labels = len(glue_processors[args.task]().get_labels())
# Load Model, Tokenizer & Datasets
config = BertConfig.from_pretrained(args.casing, num_labels=num_labels)
config.bottleneck_size = args.bottleneck_size
config.non_linearity = args.non_linearity
tokenizer = BertTokenizer.from_pretrained(args.casing)
bert_model = TFBertModel(config).from_pretrained(args.casing)
model = modeling_tf_adapter_bert.AdapterBertModel(bert_model, num_labels)
data, info = tensorflow_datasets.load(TFDS_TASK, with_info=True)
train_examples = info.splits["train"].num_examples
eval_examples = info.splits["validation" + postfix].num_examples
train_dataset = glue_convert_examples_to_features(
data["train"],
tokenizer,
max_length=args.max_seq_length,
task=args.task)
eval_dataset = glue_convert_examples_to_features(
data["validation" + postfix],
tokenizer,
max_length=args.max_seq_length,
task=args.task)
train_dataset = train_dataset.repeat().shuffle(buffer_size=100).batch(
args.batch_size)
eval_dataset = eval_dataset.batch(args.batch_size)
train_steps = int(np.ceil(train_examples / args.batch_size))
eval_steps = int(np.ceil(eval_examples / args.batch_size))
# Add Adapters
for i in range(config.num_hidden_layers):
# instantiate
model.bert.bert.encoder.layer[
i].attention.dense_output = modeling_tf_adapter_bert.TFBertSelfOutput(
model.bert.bert.encoder.layer[i].attention.dense_output.dense,
model.bert.bert.encoder.layer[i].attention.dense_output.
LayerNorm, config)
model.bert.bert.encoder.layer[
i].bert_output = modeling_tf_adapter_bert.TFBertOutput(
model.bert.bert.encoder.layer[i].bert_output.dense,
model.bert.bert.encoder.layer[i].bert_output.LayerNorm, config)
# Freeze BERT
model.bert.bert.embeddings.trainable = False
model.bert.bert.pooler.trainable = False
for i in range(config.num_hidden_layers):
model.bert.bert.encoder.layer[
i].attention.self_attention.trainable = False
model.bert.bert.encoder.layer[
i].attention.dense_output.dense.trainable = False
#model.bert.bert.encoder.layer[i].attention.dense_output.LayerNorm.trainable = False
model.bert.bert.encoder.layer[i].intermediate.trainable = False
model.bert.bert.encoder.layer[i].bert_output.dense.trainable = False
#model.bert.bert.encoder.layer[i].bert_output.LayerNorm.trainable = False
# Loss & Optimizer
if num_labels == 1:
loss = tf.keras.losses.MeanSquaredError()
monitor = 'val_spearmanr'
else:
loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
if args.task == 'cola':
monitor = 'val_matthews_corrcoef'
elif args.task in ['mrpc', 'qqp']:
monitor = 'val_f1'
else:
monitor = 'val_accuracy'
opt, scheduler = create_optimizer(
init_lr=args.learning_rate,
num_train_steps=train_steps * args.epochs,
num_warmup_steps=int(train_steps * args.epochs * args.warmup_ratio),
adam_epsilon=1e-6,
weight_decay_rate=0)
model.compile(optimizer=opt, loss=loss)
# Callback to save the best model
checkpoint = ModelCheckpoint(
eval_dataset, args.batch_size, eval_steps, monitor,
os.path.join(args.saved_models_dir, args.task))
# Fine-tuning
history = model.fit(train_dataset,
epochs=args.epochs,
steps_per_epoch=train_steps,
callbacks=[checkpoint])
# train_dataset = train_dataset.shuffle(128).batch(BATCH_SIZE).repeat(-1)
# valid_dataset = valid_dataset.batch(EVAL_BATCH_SIZE)
# test_dataset = test_dataset.batch(EVAL_BATCH_SIZE)
# train_steps = train_examples // BATCH_SIZE
# valid_steps = valid_examples // EVAL_BATCH_SIZE
# test_steps = test_examples // EVAL_BATCH_SIZE
len_train_features = len(train_features)
len_dev_features = len(dev_features)
total_train_steps = int(len_train_features * EPOCHS / BATCH_SIZE) + 1
train_dataset = get_dataset_from_features(train_features, BATCH_SIZE)
dev_dataset = get_dataset_from_features(dev_features, EVAL_BATCH_SIZE)
# Prepare training: Compile tf.keras model with optimizer, loss and learning rate schedule
opt = create_optimizer(init_lr=3e-5,
num_train_steps=total_train_steps,
num_warmup_steps=int(0.1 * total_train_steps))
# opt = tf.keras.optimizers.Adam(learning_rate=3e-5, epsilon=1e-08)
if USE_AMP:
# loss scaling is currently required when using mixed precision
opt = tf.keras.mixed_precision.experimental.LossScaleOptimizer(
opt, "dynamic")
if num_labels == 1:
loss = tf.keras.losses.MeanSquaredError()
else:
loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
metric = tf.keras.metrics.SparseCategoricalAccuracy("accuracy")
model.compile(optimizer=opt, loss=loss, metrics=[metric])
train_loss_results, train_accuracy_results = [], []
def train(self,
finetune=False,
previous_epoch=10,
do_eval_train=False) -> None:
"""
Train method to train the model.
"""
train_ds = self.get_train_tfdataset()
if self.args.debug:
tf.summary.trace_on(graph=True, profiler=True)
self.gradient_accumulator.reset()
with self.args.strategy.scope():
optimizer, lr_scheduler = self.get_optimizers()
iterations = optimizer.iterations
folder = os.path.join(self.args.output_dir, PREFIX_CHECKPOINT_DIR)
ckpt = tf.train.Checkpoint(optimizer=optimizer, model=self.model)
self.model.ckpt_manager = tf.train.CheckpointManager(
ckpt, folder, max_to_keep=self.args.save_total_limit)
if self.model.ckpt_manager.latest_checkpoint:
logger.info(
"Checkpoint file %s found and restoring from checkpoint",
self.model.ckpt_manager.latest_checkpoint)
ckpt.restore(self.model.ckpt_manager.latest_checkpoint
).expect_partial()
if finetune:
logger.info("Reset Optimizer in finetune mode")
optimizer, lr_scheduler = create_optimizer(
self.args.learning_rate,
self.train_steps * self.args.num_train_epochs,
self.args.warmup_steps,
adam_epsilon=self.args.adam_epsilon,
weight_decay_rate=self.args.weight_decay,
)
iterations = optimizer.iterations
if iterations.numpy() > 0:
logger.info("Start the training from the last checkpoint")
start_epoch = (iterations.numpy() // self.train_steps) + 1
else:
if finetune:
logger.info("Start the finetune from the last checkpoint")
start_epoch = 1 + previous_epoch
else:
start_epoch = 1
tf.summary.experimental.set_step(iterations)
epochs = 1 if self.args.max_steps > 0 else self.args.num_train_epochs
if finetune and self.args.max_steps <= 0:
epochs += previous_epoch
if self.args.fp16:
policy = tf.keras.mixed_precision.experimental.Policy(
"mixed_float16")
tf.keras.mixed_precision.experimental.set_policy(policy)
with self.tb_writer.as_default():
tf.summary.text("args", self.args.to_json_string())
self.tb_writer.flush()
logger.info("***** Running training *****")
logger.info(" Num examples = %d", self.num_train_examples)
logger.info(" Num Epochs = %d", epochs)
logger.info(" Total optimization steps = %d", self.train_steps)
for epoch_iter in range(start_epoch, int(epochs + 1)):
time_start = time.time()
for step, training_loss in enumerate(
self._training_steps(train_ds, optimizer)):
self.global_step = iterations.numpy()
self.epoch_logging = epoch_iter - 1 + (
self.global_step % self.train_steps + 1) / self.train_steps
if self.args.debug:
logs = {}
logs["loss"] = training_loss.numpy()
logs["epoch"] = self.epoch_logging
self._log(logs)
if self.global_step == 1 and self.args.debug:
with self.tb_writer.as_default():
tf.summary.trace_export(
name="training",
step=self.global_step,
profiler_outdir=self.args.logging_dir)
if self.args.evaluate_during_training and self.global_step % self.args.eval_steps == 0:
if do_eval_train:
self.evaluate_train()
self.evaluate()
if self.global_step % self.args.logging_steps == 0:
logs = {}
logs["loss"] = training_loss.numpy()
logs["learning_rate"] = lr_scheduler(
self.global_step).numpy()
logs["epoch"] = self.epoch_logging
self._log(logs)
if self.global_step % self.args.save_steps == 0:
ckpt_save_path = self.model.ckpt_manager.save()
logger.info("Saving checkpoint for step {} at {}".format(
self.global_step, ckpt_save_path))
if self.global_step % 1000 == 0:
time_elapse = time.time() - time_start
time_est = time_elapse / (step + 1) * self.train_steps
logger.info(
'Epoch: %d; Step: %d; Current Epoch Elapse/Estimate: %0.2fs/%0.2fs'
% (epoch_iter, step + 1, time_elapse, time_est))
if self.global_step % self.train_steps == 0:
break