def get_ensemble_stats(model_dir_ensemble, prediction_path):
"""Average results in all parallel ensemble directories."""
split_path = split_prediction_path(prediction_path)
filname = split_path[-1]
ensemble_dirs = gfile.glob(model_dir_ensemble)[:_ENSEMBLE_SIZE]
probs = 0.
count = 0
for d in ensemble_dirs:
# TODO(emilyaf) support replicas
glob_path = os.path.join(d, 'r0/predictions/*/*', filname)
paths = gfile.glob(glob_path)
if not paths:
continue
count += 1
# Use most recent prediction in ensemble
filestats = gfile.stat(glob_path)
idx = np.argmax([s.mtime_nsecs for s in filestats])
path = paths[idx]
stats = array_utils.load_stats_from_tfrecords(
path, max_records=_MAX_PREDICTIONS)
probs += stats['probs']
probs /= count
return probs.astype(np.float32), stats['labels'].astype(np.int32)
def run(prediction_path):
"""Run temperature scaling."""
stats = array_utils.load_stats_from_tfrecords(prediction_path)
probs = stats['probs'].astype(np.float32)
labels = stats['labels'].astype(np.int32)
if len(labels.shape) > 1:
labels = np.squeeze(labels, -1)
if probs.shape[0] > NUM_EXAMPLES:
probs = probs[:NUM_EXAMPLES, :]
labels = labels[:NUM_EXAMPLES]
probs = metrics_lib.soften_probabilities(probs=probs)
logits = uq_utils.np_inverse_softmax(probs)
temp = calibration_lib.find_scaling_temperature(labels, logits)
with gfile.GFile(
os.path.join(os.path.dirname(prediction_path),
'temperature_hparam.json'), 'w') as fh:
fh.write(json.dumps({'temperature': temp}))
def run(prediction_path, validation_path, model_dir_ensemble,
use_temp_scaling=False):
"""Runs predictions on the given dataset using the specified model."""
logging.info('Loading predictions...')
out_file_prefix = 'metrics_'
if model_dir_ensemble:
probs, labels = get_ensemble_stats(model_dir_ensemble, prediction_path)
out_file_prefix = 'metrics_ensemble_'
else:
stats = array_utils.load_stats_from_tfrecords(
prediction_path, max_records=_MAX_PREDICTIONS)
probs = stats['probs'].astype(np.float32)
labels = stats['labels'].astype(np.int32)
if len(labels.shape) > 1:
labels = np.squeeze(labels, -1)
probs = metrics_lib.soften_probabilities(probs=probs)
logits = inverse_softmax(probs)
if use_temp_scaling:
predictions_base_dir = os.path.dirname(
os.path.dirname(os.path.dirname(prediction_path)))
json_dir = os.path.join(predictions_base_dir, '*/*/temperature_hparam.json')
paths = gfile.glob(json_dir)
filestats = gfile.stat(json_dir)
idx = np.argmax([s.mtime_nsecs for s in filestats])
temperature_hparam_path = paths[idx]
with gfile.GFile(temperature_hparam_path) as fh:
temp = json.loads(fh.read())['temperature']
logits /= temp
probs = tf.nn.softmax(logits).numpy()
probs = metrics_lib.soften_probabilities(probs=probs)
out_file_prefix = 'metrics_temp_scaled_'
# confidence vs accuracy
thresholds = np.linspace(0, 1, 10, endpoint=False)
accuracies, counts = metrics_lib.compute_accuracies_at_confidences(
labels, probs, thresholds)
overall_accuracy = (probs.argmax(-1) == labels).mean()
accuracy_top5 = metrics_lib.accuracy_top_k(probs, labels, 5)
accuracy_top10 = metrics_lib.accuracy_top_k(probs, labels, 10)
probs_correct_class = probs[np.arange(probs.shape[0]), labels]
negative_log_likelihood = -np.log(probs_correct_class).mean()
entropy_per_example = tfd.Categorical(probs=probs).entropy().numpy()
uncertainty, resolution, reliability = metrics_lib.brier_decomposition(
labels=labels, logits=logits)
brier = tf.reduce_mean(metrics_lib.brier_scores(labels=labels, logits=logits))
ece = metrics_lib.expected_calibration_error_multiclass(
probs, labels, _ECE_BINS)
ece_top5 = metrics_lib.expected_calibration_error_multiclass(
probs, labels, _ECE_BINS, top_k=5)
ece_top10 = metrics_lib.expected_calibration_error_multiclass(
probs, labels, _ECE_BINS, top_k=10)
validation_stats = array_utils.load_stats_from_tfrecords(
validation_path, max_records=20000)
validation_probs = validation_stats['probs'].astype(np.float32)
bins = metrics_lib.get_quantile_bins(_ECE_BINS, validation_probs)
q_ece = metrics_lib.expected_calibration_error_multiclass(
probs, labels, bins)
bins = metrics_lib.get_quantile_bins(_ECE_BINS, validation_probs, top_k=5)
q_ece_top5 = metrics_lib.expected_calibration_error_multiclass(
probs, labels, bins, top_k=5)
bins = metrics_lib.get_quantile_bins(_ECE_BINS, validation_probs, top_k=10)
q_ece_top10 = metrics_lib.expected_calibration_error_multiclass(
probs, labels, bins, top_k=10)
metrics = {
'accuracy': overall_accuracy,
'accuracy_top5': accuracy_top5,
'accuracy_top10': accuracy_top10,
'accuracy_at_confidence': accuracies,
'confidence_thresholds': thresholds,
'confidence_counts': counts,
'ece': ece,
'ece_top5': ece_top5,
'ece_top10': ece_top10,
'q_ece': q_ece,
'q_ece_top5': q_ece_top5,
'q_ece_top10': q_ece_top10,
'ece_nbins': _ECE_BINS,
'entropy_per_example': entropy_per_example,
'brier_uncertainty': uncertainty.numpy(),
'brier_resolution': resolution.numpy(),
'brier_reliability': reliability.numpy(),
'brier_score': brier.numpy(),
'true_labels': labels,
'pred_labels': probs.argmax(-1),
'prob_true_label': probs[np.arange(len(labels)), labels],
'prob_pred_label': probs.max(-1),
'negative_log_likelihood': negative_log_likelihood,
}
save_dir = os.path.dirname(prediction_path)
split_path = split_prediction_path(prediction_path)
prediction_file = split_path[-1]
dataset_name = '-'.join(prediction_file.split('_')[2:])
out_file = out_file_prefix + dataset_name + '.npz'
array_utils.write_npz(save_dir, out_file, metrics)