def test_regression_model(self): n_classes = 1 # Only one label for regression config = dummy_config(model_arch=ModelArchitecture.F_NET) with config.unlocked(): config.dataset_name = "glue/stsb" # regression task dataset frozen_config = ml_collections.FrozenConfigDict(config) model = models.SequenceClassificationModel(config=frozen_config, n_classes=n_classes) rng = jax.random.PRNGKey(0) init_batch = init_encoder_batch(config) params = init_model_params(rng, model, init_batch) self.assertEqual(params.keys(), {"encoder", "classification"}) # Logits for eval/prediction (no labels supplied). eval_inputs = dummy_inputs(rng, config) eval_inputs["deterministic"] = True logits = model.apply({"params": params}, **eval_inputs) expected_logits_shape = (config.train_batch_size, n_classes) self.assertEqual(jnp.shape(logits), expected_logits_shape) # Metrics for training (labels supplied). train_inputs = dummy_inputs(rng, config) _, label_key = jax.random.split(rng) train_inputs["labels"] = jax.random.uniform( label_key, (config.train_batch_size, ), minval=0., maxval=1.) metrics = model.apply({"params": params}, rngs={"dropout": rng}, **train_inputs) self.assertEqual(metrics.keys(), {"loss", "num_labels"})
def test_classification_model(self): n_classes = 2 config = dummy_config(model_arch=ModelArchitecture.BERT) with config.unlocked(): config.dataset_name = "dummy/classification_dataset" frozen_config = ml_collections.FrozenConfigDict(config) model = models.SequenceClassificationModel(config=frozen_config, n_classes=n_classes) rng = jax.random.PRNGKey(0) init_batch = init_encoder_batch(config) params = init_model_params(rng, model, init_batch) self.assertEqual(params.keys(), {"encoder", "classification"}) # Logits for eval/prediction (no labels supplied). eval_inputs = dummy_inputs(rng, config) eval_inputs["deterministic"] = True logits = model.apply({"params": params}, **eval_inputs) expected_logits_shape = (config.train_batch_size, n_classes) self.assertEqual(jnp.shape(logits), expected_logits_shape) # Metrics for training (labels supplied). train_inputs = dummy_inputs(rng, config) train_inputs["labels"] = jnp.ones(config.train_batch_size, jnp.int32) metrics = model.apply({"params": params}, rngs={"dropout": rng}, **train_inputs) self.assertEqual(metrics.keys(), {"loss", "correct_predictions", "num_labels"})
def train_and_evaluate(config, workdir, vocab_filepath): """Runs a training and evaluation loop. Args: config: Model and training configuration. workdir: Working directory for checkpoints and Tensorboard summaries. If this contains a checkpoint, training will be resumed from the latest checkpoint. vocab_filepath: Absolute path to SentencePiece vocab model. Raises: ValueError: If training or eval batch sizes won't fit number of processes and devices, or config is underspecified. """ n_processes = jax.process_count() # Number of processes n_devices = jax.local_device_count() # Number of local devices per process if config.train_batch_size % (n_processes * n_devices) > 0: raise ValueError( "Training batch size must be divisible by the total number of devices, " "but training batch size = %d, while total number of devices = %d " "(%d processes, each with %d devices)" % (config.train_batch_size, n_processes * n_devices, n_processes, n_devices)) if config.eval_batch_size % (n_processes * n_devices) > 0: raise ValueError( "Eval batch size must be divisible by the total number of devices, " "but eval batch size = %d, while total number of devices = %d " "(%d processes, each with %d devices)" % (config.eval_batch_size, n_processes * n_devices, n_processes, n_devices)) per_process_train_batch_size = config.train_batch_size // n_processes per_process_eval_batch_size = config.eval_batch_size // n_processes if jax.process_index() == 0: train_summary_writer = tensorboard.SummaryWriter( os.path.join(workdir, "train")) eval_summary_writer = tensorboard.SummaryWriter( os.path.join(workdir, "eval")) else: train_summary_writer = None eval_summary_writer = None rng = random.PRNGKey(config.seed) rng, init_rng = random.split(rng) ds_info = tfds.builder(config.dataset_name).info num_train_examples = ds_info.splits[tfds.Split.TRAIN].num_examples num_train_steps = int(num_train_examples * config.num_train_epochs // config.train_batch_size) num_warmup_steps = int(config.warmup_proportion * num_train_steps) # Round up evaluation frequency to power of 10. eval_frequency = int( math.ceil(config.eval_proportion * num_train_steps / 10)) * 10 is_regression_task = config.dataset_name == "glue/stsb" num_classes = (1 if is_regression_task else ds_info.features["label"].num_classes) tokenizer = spm.SentencePieceProcessor() tokenizer.Load(vocab_filepath) with config.unlocked(): config.vocab_size = tokenizer.GetPieceSize() frozen_config = ml_collections.FrozenConfigDict(config) model = models.SequenceClassificationModel(config=frozen_config, n_classes=num_classes) params = _init_params(model, init_rng, config) optimizer = _create_adam_optimizer(config.learning_rate, params) # In case current job restarts, ensure that we continue from where we left # off. optimizer = checkpoints.restore_checkpoint(workdir, optimizer) start_step = int(optimizer.state.step) # Otherwise, try to restore optimizer and model state from config checkpoint. if (start_step == 0 and "init_checkpoint_dir" in config and config.init_checkpoint_dir): optimizer = _restore_pretrained_model(optimizer, params, config) # We access model state only from optimizer via optimizer.target. del params optimizer = jax_utils.replicate(optimizer) if is_regression_task: compute_stats = functools.partial(_compute_regression_stats, model=model, pad_id=tokenizer.pad_id()) else: compute_stats = functools.partial(_compute_classification_stats, model=model, pad_id=tokenizer.pad_id()) learning_rate_fn = train_utils.create_learning_rate_scheduler( factors="constant * linear_warmup * linear_decay", base_learning_rate=config.learning_rate, warmup_steps=num_warmup_steps, decay_steps=num_train_steps - num_warmup_steps, ) glue_inputs = functools.partial(input_pipeline.glue_inputs, dataset_name=config.dataset_name, max_seq_length=config.max_seq_length, tokenizer=tokenizer) train_ds = glue_inputs(split=tfds.Split.TRAIN, batch_size=per_process_train_batch_size, training=True) train_iter = iter(train_ds) if config.dataset_name == "glue/mnli": # MNLI contains two validation and test datasets. split_suffixes = ["_matched", "_mismatched"] else: split_suffixes = [""] # We init the first set of dropout PRNG keys, but update it afterwards inside # the main pmap'd training update for performance. rngs = random.split(rng, n_devices) loss_and_metrics_fn = functools.partial(_compute_loss_and_metrics, model=model, pad_id=tokenizer.pad_id()) p_train_step = jax.pmap(functools.partial( train_utils.train_step, loss_and_metrics_fn=loss_and_metrics_fn, learning_rate_fn=learning_rate_fn), axis_name="batch") p_eval_step = jax.pmap(functools.partial(train_utils.eval_step, metric_fn=compute_stats), axis_name="batch") eval_metrics_fn = _create_eval_metrics_fn(config.dataset_name, is_regression_task) train_metrics = [] logging.info("Starting training loop.") logging.info("====================") for step in range(start_step, num_train_steps): with jax.profiler.StepTraceAnnotation("train", step_num=step): train_batch = next(train_iter) train_batch = common_utils.shard(train_batch) optimizer, train_step_metrics, rngs = p_train_step(optimizer, train_batch, rng=rngs) train_metrics.append(train_step_metrics) if ((step > 0 and config.save_checkpoints_steps and step % config.save_checkpoints_steps == 0) or step == num_train_steps - 1) and jax.process_index() == 0: # Save un-replicated optimizer and model state. checkpoints.save_checkpoint(workdir, jax_utils.unreplicate(optimizer), step, keep=2) # Periodic metric handling. if step % eval_frequency != 0 and step < num_train_steps - 1: continue logging.info("Gathering training metrics at step: %d", step) train_metrics = common_utils.get_metrics(train_metrics) train_summary = { "loss": jnp.sum(train_metrics["loss"]) / jnp.sum(train_metrics["num_labels"]), "learning_rate": learning_rate_fn(step) } if not is_regression_task: train_summary["accuracy"] = jnp.sum( train_metrics["correct_predictions"]) / jnp.sum( train_metrics["num_labels"]) if jax.process_index() == 0: assert train_summary_writer for key, val in train_summary.items(): train_summary_writer.scalar(key, val, step) train_summary_writer.flush() # Reset metric accumulation for next evaluation cycle. train_metrics = [] logging.info("Gathering validation metrics at step: %d", step) for split_suffix in split_suffixes: eval_ds = glue_inputs(split=tfds.Split.VALIDATION + split_suffix, batch_size=per_process_eval_batch_size, training=False) all_stats = [] for _, eval_batch in zip(range(config.max_num_eval_steps), eval_ds): all_stats.append( _evaluate(p_eval_step, optimizer.target, eval_batch, n_devices)) flat_stats = {} for k in all_stats[ 0]: # All batches of output stats are the same size flat_stats[k] = np.concatenate([stat[k] for stat in all_stats], axis=0) eval_summary = eval_metrics_fn(flat_stats) if jax.process_index() == 0: assert eval_summary_writer for key, val in eval_summary.items(): eval_summary_writer.scalar(f"{key}{split_suffix}", val, step) eval_summary_writer.flush()