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 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)
