def write_outputs(self, tasks, trial, split):
"""Write model prediction to disk."""
utils.log("Writing out predictions for", tasks, split)
distill_input_fn, _, _ = self._preprocessor.prepare_predict(
tasks, split)
results = self._estimator.predict(input_fn=distill_input_fn,
yield_single_examples=True)
# task name -> eid -> model-logits
logits = collections.defaultdict(dict)
for r in results:
if r["task_id"] != len(self._tasks):
r = utils.nest_dict(r, self._config.task_names)
task_name = self._config.task_names[r["task_id"]]
logits[task_name][r[task_name]["eid"]] = (
r[task_name]["logits"] if "logits" in r[task_name] else
r[task_name]["predictions"])
for task_name in logits:
utils.log("Pickling predictions for {:} {:} examples ({:})".format(
len(logits[task_name]), task_name, split))
if split == "train":
if trial <= self._config.n_writes_distill:
utils.write_pickle(
logits[task_name],
self._config.distill_outputs(task_name, trial))
else:
if trial <= self._config.n_writes_test:
utils.write_pickle(
logits[task_name],
self._config.test_outputs(task_name, split, trial))
def serialize_examples(self, examples, is_training, output_file):
"""Convert a set of `InputExample`s to a TFRecord file."""
with tf.python_io.TFRecordWriter(output_file) as writer:
for (ex_index, example) in enumerate(examples):
if ex_index % 50000 == 0:
utils.log("Writing example {:} of {:}".format(
ex_index, len(examples)))
tf_example = self._example_to_tf_example(example, is_training)
writer.write(tf_example.SerializeToString())
def write_results(config, results):
"""Write out evaluate metrics to disk."""
utils.log("Writing results to", config.results_txt)
utils.mkdir(config.results_txt.rsplit("/", 1)[0])
utils.write_pickle(results, config.results_pkl)
with tf.gfile.GFile(config.results_txt, "w") as f:
results_str = ""
for trial_results in results:
for task_name, task_results in trial_results.items():
results_str += task_name + ": " + " - ".join(
["{:}: {:.2f}".format(k, v)
for k, v in task_results.items()]) + "\n"
f.write(results_str)
def log(self, run_values=None):
msg = '{:}/{:} = {:.1f}%'.format(
self._iter_count, self._n_steps,
100.0 * self._iter_count / self._n_steps)
time_elapsed = time.time() - self._start_time
time_per_step = time_elapsed / self._iter_count
msg += ', GPS: {:.1f}'.format(1 / time_per_step)
msg += ', ELAP: ' + secs_to_str(time_elapsed)
msg += ', ETA: ' + secs_to_str(
(self._n_steps - self._iter_count) * time_per_step)
if run_values is not None:
for tag, value in run_values.results.items():
msg += ' - ' + str(tag) + (': {:.4f}'.format(value))
utils.log(msg)
def model_fn(features, labels, mode, params): # pylint: disable=unused-argument
"""The `model_fn` for TPUEstimator."""
utils.log("Building model")
is_training = (mode == tf.estimator.ModeKeys.TRAIN)
model = MultitaskModel(config, tasks, task_weights, is_training,
features, num_train_steps)
# Load pre-trained weights from checkpoint
tvars = tf.trainable_variables()
scaffold_fn = None
if not config.debug:
assignment_map, _ = modeling.get_assignment_map_from_checkpoint(
tvars, config.init_checkpoint)
if config.use_tpu:
def tpu_scaffold():
tf.train.init_from_checkpoint(config.init_checkpoint,
assignment_map)
return tf.train.Scaffold()
scaffold_fn = tpu_scaffold
else:
tf.train.init_from_checkpoint(config.init_checkpoint,
assignment_map)
# Run training or prediction
if mode == tf.estimator.ModeKeys.TRAIN:
train_op = optimization.create_optimizer(config, model.loss,
num_train_steps)
output_spec = tf.contrib.tpu.TPUEstimatorSpec(
mode=mode,
loss=model.loss,
train_op=train_op,
scaffold_fn=scaffold_fn,
training_hooks=[
training_utils.ETAHook(
config,
{} if config.use_tpu else dict(loss=model.loss),
num_train_steps)
])
else:
assert mode == tf.estimator.ModeKeys.PREDICT
output_spec = tf.contrib.tpu.TPUEstimatorSpec(
mode=mode,
predictions=utils.flatten_dict(model.outputs),
scaffold_fn=scaffold_fn)
utils.log("Building complete")
return output_spec
def main(_):
topdir, model_name, hparams = sys.argv[-3:] # pylint: disable=unbalanced-tuple-unpacking
config = configure.Config(topdir, model_name, **json.loads(hparams))
# Setup for training
tasks = task_builder.get_tasks(config)
results = []
trial = 1
utils.rmkdir(config.checkpoints_dir)
heading_info = "model={:}, trial {:}/{:}".format(config.model_name, trial,
config.num_trials)
heading = lambda msg: utils.heading(msg + ": " + heading_info)
# Train and evaluate num_trials models with different random seeds
while trial <= config.num_trials:
heading("Start training")
model_runner = ModelRunner(config, tasks)
model_runner.train()
utils.log()
heading("Run evaluation")
results.append(model_runner.evaluate())
write_results(config, results)
if ((config.write_distill_outputs and trial <= config.n_writes_distill)
or
(config.write_test_outputs and trial <= config.n_writes_test)):
heading("Write outputs")
for task in tasks:
if config.write_distill_outputs:
model_runner.write_outputs([task], trial, "train")
if config.write_test_outputs:
for split in task.get_test_splits():
model_runner.write_outputs([task], trial, split)
utils.rmkdir(config.checkpoints_dir)
trial += 1
def _serialize_dataset(self, tasks, is_training, split):
"""Writes out tfrecord examples for the specified tasks."""
dataset_name = "_".join(sorted([task.name for task in tasks]))
dataset_name += "_" + split
if self._config.distill:
dataset_name += "_distill"
dataset_prefix = os.path.join(self._config.preprocessed_data_dir,
dataset_name)
tfrecords_path = dataset_prefix + ".tfrecord"
metadata_path = dataset_prefix + ".metadata"
batch_size = (self._config.train_batch_size
if is_training else self._config.eval_batch_size)
utils.log("Loading dataset", dataset_name)
n_examples = None
sizes = {}
if (self._config.use_tfrecords_if_existing
and tf.gfile.Exists(metadata_path)):
utils.log("Using already-written tfrecords")
metadata = utils.load_json(metadata_path)
n_examples = metadata["n_examples"]
sizes = metadata["sizes"]
if n_examples is None:
utils.log("Existing tfrecords not found so creating")
examples = []
for task in tasks:
task_examples = task.get_examples(split)
sizes[task.name] = len(task_examples)
examples += task_examples
last_index = len(examples) - (len(examples) % batch_size)
#while len(examples) % batch_size != 0:
# examples.append(PaddingInputExample())
examples = examples[:
last_index] #drop the instances if batch size does not perfectly split the set
if is_training:
random.shuffle(examples)
n_examples = len(examples)
assert n_examples % batch_size == 0
utils.mkdir(tfrecords_path.rsplit("/", 1)[0])
self.serialize_examples(examples, is_training, tfrecords_path)
utils.write_json({
"n_examples": n_examples,
"sizes": sizes
}, metadata_path)
input_fn = self._input_fn_builder(tfrecords_path, is_training)
if is_training:
steps = int(n_examples // batch_size *
self._config.num_train_epochs)
else:
steps = n_examples // batch_size
return input_fn, steps, sizes
def _evaluate_task(self, task):
"""Evaluate the current model on the dev set."""
utils.log("Evaluating", task.name)
eval_input_fn, _, _ = self._preprocessor.prepare_eval(task)
results = self._estimator.predict(input_fn=eval_input_fn,
yield_single_examples=True)
scorer = task.get_scorer()
for r in results:
if r["task_id"] != len(self._tasks):
r = utils.nest_dict(r, self._config.task_names)
scorer.update(r[task.name])
utils.log(task.name + ": " + scorer.results_str())
utils.log()
return dict(scorer.get_results())
def train(self):
utils.log("Training for {:} steps".format(self.train_steps))
self._estimator.train(input_fn=self._train_input_fn,
max_steps=self.train_steps)
