def test_TFBertForQuestionAnswering(self):
from transformers import BertConfig, TFBertForQuestionAnswering
keras.backend.clear_session()
# pretrained_weights = 'bert-base-uncased'
tokenizer_file = 'bert_bert-base-uncased.pickle'
tokenizer = self._get_tokenzier(tokenizer_file)
text, inputs, inputs_onnx = self._prepare_inputs(tokenizer)
config = BertConfig()
model = TFBertForQuestionAnswering(config)
predictions = model.predict(inputs)
onnx_model = keras2onnx.convert_keras(model, model.name)
self.assertTrue(run_onnx_runtime(onnx_model.graph.name, onnx_model, inputs_onnx, predictions, self.model_files))
def fine_tune_squad(**config):
# Get checkpoint from Hugging Face's models repo or from our own checkpoint.
bert_model_name = config.get("bert_model_name", None)
pretrained_ckpt_path = config.get("pretrained_ckpt_path", None) # This can be passed from both cli and config file.
assert bert_model_name is not None or pretrained_ckpt_path is not None, \
"SQuAD requires a pretrained model, either `bert_model_name` (via config file) or `pretrained_ckpt_path` " \
"(via config file or `--pretrained-ckpt-path` command line argument) but none provided."
assert (bert_model_name is not None and pretrained_ckpt_path is None) \
or (bert_model_name is None and pretrained_ckpt_path is not None), \
f"Only one checkpoint is accepted, but two provided: `bert_model_name`={bert_model_name}, " \
f"and `pretrained_ckpt_path`={pretrained_ckpt_path}."
if pretrained_ckpt_path is not None:
bert_config_params = config["bert_config"]
# Get required options
micro_batch_size = config["micro_batch_size"]
num_epochs = config["num_epochs"]
optimizer_opts = config["optimizer_opts"]
learning_rate = config['learning_rate']
replicas = config["replicas"]
grad_acc_steps_per_replica = config["grad_acc_steps_per_replica"]
wandb_opts = config["wandb_opts"]
use_outlining = config["use_outlining"]
replace_layers = config["replace_layers"]
enable_recomputation = config["enable_recomputation"]
embedding_serialization_factor = config["embedding_serialization_factor"]
optimizer_state_offchip = config["optimizer_state_offchip"]
matmul_available_memory_proportion_per_pipeline_stage = config[
"matmul_available_memory_proportion_per_pipeline_stage"]
matmul_partials_type = config["matmul_partials_type"]
pipeline_stages = config["pipeline_stages"]
device_mapping = config["device_mapping"]
global_batches_per_log = config["global_batches_per_log"]
seed = config["seed"]
cache_dir = config["cache_dir"]
output_dir = config["output_dir"]
# Get optional options
save_ckpt_path = config.get("save_ckpt_path", Path(__file__).parent.joinpath("checkpoints").absolute())
ckpt_every_n_steps_per_execution = config.get("ckpt_every_n_steps_per_execution", 2000)
universal_run_name = config.get("name", f"{Path(config['config']).stem}-{wandb_opts['init']['name']}")
universal_run_name += f"-{datetime.now().strftime('%Y%m%d_%H%M%S')}"
print(f"Universal name for run: {universal_run_name}")
set_random_seeds(seed)
num_pipeline_stages = len(device_mapping)
num_ipus_per_replicas = max(device_mapping) + 1
num_ipus = replicas * num_ipus_per_replicas
# Load training and validation data
# =================================
train_dataset, eval_dataset, num_train_samples, num_eval_samples, raw_datasets = get_squad_data(
micro_batch_size,
cache_dir
)
train_batch_config = BatchConfig(micro_batch_size=micro_batch_size,
total_num_train_samples=num_epochs * num_train_samples.numpy(),
num_replicas=replicas,
gradient_accumulation_count=grad_acc_steps_per_replica,
dataset_size=num_train_samples.numpy(),
global_batches_per_log=global_batches_per_log,
task=Task.OTHER)
# Create model
# ============
policy = tf.keras.mixed_precision.Policy("float16")
tf.keras.mixed_precision.set_global_policy(policy)
strategy = create_ipu_strategy(num_ipus, enable_recomputation=enable_recomputation)
with strategy.scope():
# Instantiate the pretrained model given in the config.
if bert_model_name is not None:
model = TFBertForQuestionAnswering.from_pretrained(bert_model_name)
else:
bert_config = BertConfig(**bert_config_params, hidden_act=ipu.nn_ops.gelu)
model = TFBertForQuestionAnswering(config=bert_config)
# Convert subclass model to functional, expand main layers to enable pipelining, and replace some layers to
# optimise performance.
model = convert_tf_bert_model(
model,
train_dataset,
post_process_bert_input_layer,
replace_layers=replace_layers,
use_outlining=use_outlining,
embedding_serialization_factor=embedding_serialization_factor,
rename_outputs={'tf.compat.v1.squeeze': 'start_positions', 'tf.compat.v1.squeeze_1': 'end_positions'}
)
# Load from pretrained checkpoint if requested.
if pretrained_ckpt_path is not None:
print(f"Attempting to load pretrained checkpoint from path {pretrained_ckpt_path}. "
f"This will overwrite the current weights")
load_checkpoint_into_model(model, pretrained_ckpt_path)
# Configure pipeline stages
# =========================
if num_pipeline_stages > 1:
pipeline_assigner = PipelineStagesAssigner(PIPELINE_ALLOCATE_PREVIOUS, PIPELINE_NAMES)
assignments = model.get_pipeline_stage_assignment()
assignments = pipeline_assigner.assign_pipeline_stages(assignments, pipeline_stages)
model.set_pipeline_stage_assignment(assignments)
model.print_pipeline_stage_assignment_summary()
poplar_options_per_pipeline_stage = get_poplar_options_per_pipeline_stage(
num_ipus_per_replicas,
device_mapping,
matmul_available_memory_proportion_per_pipeline_stage,
matmul_partials_type
)
model.set_pipelining_options(
gradient_accumulation_steps_per_replica=grad_acc_steps_per_replica,
pipeline_schedule=ipu.ops.pipelining_ops.PipelineSchedule.Grouped,
device_mapping=device_mapping,
offload_weight_update_variables=optimizer_state_offchip,
forward_propagation_stages_poplar_options=poplar_options_per_pipeline_stage,
backward_propagation_stages_poplar_options=poplar_options_per_pipeline_stage,
recomputation_mode=ipu.pipelining_ops.RecomputationMode.RecomputeAndBackpropagateInterleaved,
)
# Compile the model for training
# ==============================
# Wrap loss in an out-feed queue.
loss_outfeed_queue = ipu.ipu_outfeed_queue.IPUOutfeedQueue()
qa_loss = wrap_loss_in_enqueuer(QuestionAnsweringLossFunction,
loss_outfeed_queue,
["end_positions_loss", "start_positions_loss"])()
# Define optimiser with polynomial decay learning rate.
learning_rate['lr_schedule_params']['total_steps'] = train_batch_config.num_train_steps
lr_outfeed_queue = ipu.ipu_outfeed_queue.IPUOutfeedQueue(outfeed_mode=ipu.ipu_outfeed_queue.IPUOutfeedMode.LAST)
lr_scheduler = get_lr_scheduler(scheduler_name=learning_rate["lr_schedule"],
schedule_params=learning_rate["lr_schedule_params"],
queue=lr_outfeed_queue)
# Prepare optimizer.
outline_optimizer_apply_gradients = use_outlining
optimizer = get_optimizer(
optimizer_opts["name"],
grad_acc_steps_per_replica,
replicas,
lr_scheduler,
outline_optimizer_apply_gradients,
weight_decay_rate=optimizer_opts["params"]["weight_decay_rate"],
)
# Compile the model.
model.compile(
optimizer=optimizer,
loss={"end_positions": qa_loss, "start_positions": qa_loss},
metrics='accuracy',
steps_per_execution=train_batch_config.steps_per_execution
)
# Train the model
# ===============
# Set up callbacks
callbacks = CallbackFactory.get_callbacks(
universal_run_name=universal_run_name,
batch_config=train_batch_config,
model=model,
checkpoint_path=save_ckpt_path,
ckpt_every_n_steps_per_execution=ckpt_every_n_steps_per_execution,
outfeed_queues=[lr_outfeed_queue, loss_outfeed_queue],
config=config,
)
# Print configs to be logged in wandb's terminal.
print(config)
print(f"Training batch config:\n{train_batch_config}")
# Train the model
history = model.fit(
train_dataset,
steps_per_epoch=train_batch_config.num_micro_batches_per_epoch,
epochs=num_epochs,
callbacks=callbacks
)
# Evaluate the model on the validation set
# ========================================
# Prepare the dataset to be evaluated in the IPU.
eval_batch_config = BatchConfig(micro_batch_size=micro_batch_size,
total_num_train_samples=num_eval_samples.numpy(),
num_replicas=replicas,
gradient_accumulation_count=grad_acc_steps_per_replica,
dataset_size=num_eval_samples.numpy(),
task=Task.OTHER)
max_eval_samples = eval_batch_config.micro_batch_size * eval_batch_config.num_micro_batches_per_epoch
eval_pred_dataset = get_prediction_dataset(eval_dataset, max_eval_samples)
with strategy.scope():
# Re-compile the model for prediction if needed.
if train_batch_config.steps_per_execution != eval_batch_config.steps_per_execution:
model.compile(steps_per_execution=eval_batch_config.steps_per_execution)
# Get predictions for the validation data.
print(f"Running inference:\nGenerating predictions on the validation data...")
predictions = model.predict(
eval_pred_dataset,
batch_size=eval_batch_config.micro_batch_size
)
# The predictions for the end position goes first in the model outputs tuple (note the output of model.summary()).
end_predictions, start_predictions = predictions
# Match the predictions to answers in the original context.
# This will also write out the predictions to a json file in the directory given by `output_dir`.
final_predictions = postprocess_qa_predictions(
list(raw_datasets["validation"].as_numpy_iterator()),
list(eval_dataset.unbatch().as_numpy_iterator()),
(start_predictions, end_predictions),
output_dir=output_dir
)
# Format the predictions and the actual labels as expected by the metric.
formatted_predictions = [{"id": k, "prediction_text": v} for k, v in final_predictions.items()]
formatted_labels = format_raw_data_for_metric(raw_datasets["validation"])
metric = load_metric("squad")
metrics = metric.compute(predictions=formatted_predictions, references=formatted_labels)
print("Evaluation metrics:")
for key, value in metrics.items():
print(f"{key}: {value:.3f}")
return history