def main(argv):
    del argv
    if (not FLAGS.use_gpu):
        raise ValueError('Only GPU is currently supported.')
    if (FLAGS.num_cores > 1):
        raise ValueError('Only a single accelerator is currently supported.')
    tf.enable_v2_behavior()
    tf.random.set_seed(FLAGS.seed)
    tf.io.gfile.makedirs(FLAGS.output_dir)
    batch_size = (FLAGS.per_core_batch_size * FLAGS.num_cores)
    steps_per_eval = (IMAGENET_VALIDATION_IMAGES // batch_size)
    dataset_test = utils.ImageNetInput(is_training=False, data_dir=FLAGS.data_dir, batch_size=FLAGS.per_core_batch_size, use_bfloat16=False).input_fn()
    test_datasets = {'clean': dataset_test}
    (corruption_types, max_intensity) = utils.load_corrupted_test_info()
    for name in corruption_types:
        for intensity in range(1, (max_intensity + 1)):
            dataset_name = '{0}_{1}'.format(name, intensity)
            test_datasets[dataset_name] = utils.load_corrupted_test_dataset(name=name, intensity=intensity, batch_size=FLAGS.per_core_batch_size, drop_remainder=True, use_bfloat16=False)
    model = deterministic_model.resnet50(input_shape=(224, 224, 3), num_classes=NUM_CLASSES)
    logging.info('Model input shape: %s', model.input_shape)
    logging.info('Model output shape: %s', model.output_shape)
    logging.info('Model number of weights: %s', model.count_params())
    ensemble_filenames = tf.io.gfile.glob(os.path.join(FLAGS.checkpoint_dir, '**/*.index'))
    ensemble_filenames = [filename[:(- 6)] for filename in ensemble_filenames]
    ensemble_size = len(ensemble_filenames)
    logging.info('Ensemble size: %s', ensemble_size)
    logging.info('Ensemble number of weights: %s', (ensemble_size * model.count_params()))
    logging.info('Ensemble filenames: %s', str(ensemble_filenames))
    checkpoint = tf.train.Checkpoint(model=model)
    num_datasets = len(test_datasets)
    for (m, ensemble_filename) in enumerate(ensemble_filenames):
        checkpoint.restore(ensemble_filename)
        for (n, (name, test_dataset)) in enumerate(test_datasets.items()):
            filename = '{dataset}_{member}.npy'.format(dataset=name, member=m)
            filename = os.path.join(FLAGS.output_dir, filename)
            if (not tf.io.gfile.exists(filename)):
                logits = []
                test_iterator = iter(test_dataset)
                for _ in range(steps_per_eval):
                    (features, _) = next(test_iterator)
                    logits.append(model(features, training=False))
                logits = tf.concat(logits, axis=0)
                with tf.io.gfile.GFile(filename, 'w') as f:
                    np.save(f, logits.numpy())
            percent = (((m * num_datasets) + (n + 1)) / (ensemble_size * num_datasets))
            message = '{:.1%} completion for prediction: ensemble member {:d}/{:d}. Dataset {:d}/{:d}'.format(percent, (m + 1), ensemble_size, (n + 1), num_datasets)
            logging.info(message)
    metrics = {'test/negative_log_likelihood': tf.keras.metrics.Mean(), 'test/gibbs_cross_entropy': tf.keras.metrics.Mean(), 'test/accuracy': tf.keras.metrics.SparseCategoricalAccuracy(), 'test/ece': ed.metrics.ExpectedCalibrationError(num_bins=FLAGS.num_bins)}
    corrupt_metrics = {}
    for name in test_datasets:
        corrupt_metrics['test/nll_{}'.format(name)] = tf.keras.metrics.Mean()
        corrupt_metrics['test/accuracy_{}'.format(name)] = tf.keras.metrics.SparseCategoricalAccuracy()
        corrupt_metrics['test/ece_{}'.format(name)] = ed.metrics.ExpectedCalibrationError(num_bins=FLAGS.num_bins)
    for (n, (name, test_dataset)) in enumerate(test_datasets.items()):
        logits_dataset = []
        for m in range(ensemble_size):
            filename = '{dataset}_{member}.npy'.format(dataset=name, member=m)
            filename = os.path.join(FLAGS.output_dir, filename)
            with tf.io.gfile.GFile(filename, 'rb') as f:
                logits_dataset.append(np.load(f))
        logits_dataset = tf.convert_to_tensor(logits_dataset)
        test_iterator = iter(test_dataset)
        for step in range(steps_per_eval):
            (_, labels) = next(test_iterator)
            logits = logits_dataset[:, (step * batch_size):((step + 1) * batch_size)]
            labels = tf.cast(tf.reshape(labels, [(- 1)]), tf.int32)
            negative_log_likelihood = tf.reduce_mean(ensemble_negative_log_likelihood(labels, logits))
            per_probs = tf.nn.softmax(logits)
            probs = tf.reduce_mean(per_probs, axis=0)
            if (name == 'clean'):
                gibbs_ce = tf.reduce_mean(gibbs_cross_entropy(labels, logits))
                metrics['test/negative_log_likelihood'].update_state(negative_log_likelihood)
                metrics['test/gibbs_cross_entropy'].update_state(gibbs_ce)
                metrics['test/accuracy'].update_state(labels, probs)
                metrics['test/ece'].update_state(labels, probs)
            else:
                corrupt_metrics['test/nll_{}'.format(name)].update_state(negative_log_likelihood)
                corrupt_metrics['test/accuracy_{}'.format(name)].update_state(labels, probs)
                corrupt_metrics['test/ece_{}'.format(name)].update_state(labels, probs)
        message = '{:.1%} completion for evaluation: dataset {:d}/{:d}'.format(((n + 1) / num_datasets), (n + 1), num_datasets)
        logging.info(message)
    corrupt_results = utils.aggregate_corrupt_metrics(corrupt_metrics, corruption_types, max_intensity, FLAGS.alexnet_errors_path)
    total_results = {name: metric.result() for (name, metric) in metrics.items()}
    total_results.update(corrupt_results)
    logging.info('Metrics: %s', total_results)
def main(argv):
    del argv
    if (not FLAGS.use_gpu):
        raise ValueError('Only GPU is currently supported.')
    if (FLAGS.num_cores > 1):
        raise ValueError('Only a single accelerator is currently supported.')
    tf.enable_v2_behavior()
    tf.random.set_seed(FLAGS.seed)
    tf.io.gfile.makedirs(FLAGS.output_dir)
    batch_size = (FLAGS.per_core_batch_size * FLAGS.num_cores)
    steps_per_eval = (IMAGENET_VALIDATION_IMAGES // batch_size)
    dataset_test = utils.ImageNetInput(is_training=False, data_dir=FLAGS.data_dir, batch_size=FLAGS.per_core_batch_size, use_bfloat16=False).input_fn()
    test_datasets = {'clean': dataset_test}
    (corruption_types, max_intensity) = utils.load_corrupted_test_info()
    for name in corruption_types:
        for intensity in range(1, (max_intensity + 1)):
            dataset_name = '{0}_{1}'.format(name, intensity)
            test_datasets[dataset_name] = utils.load_corrupted_test_dataset(name=name, intensity=intensity, batch_size=FLAGS.per_core_batch_size, drop_remainder=True, use_bfloat16=False)
    model = deterministic_model.resnet50(input_shape=(224, 224, 3), num_classes=NUM_CLASSES)
    logging.info('Model input shape: %s', model.input_shape)
    logging.info('Model output shape: %s', model.output_shape)
    logging.info('Model number of weights: %s', model.count_params())
    ensemble_filenames = tf.io.gfile.glob(os.path.join(FLAGS.checkpoint_dir, '**/*.index'))
    ensemble_filenames = [filename[:(- 6)] for filename in ensemble_filenames]
    ensemble_size = len(ensemble_filenames)
    logging.info('Ensemble size: %s', ensemble_size)
    logging.info('Ensemble number of weights: %s', (ensemble_size * model.count_params()))
    logging.info('Ensemble filenames: %s', str(ensemble_filenames))
    checkpoint = tf.train.Checkpoint(model=model)
    num_datasets = len(test_datasets)
    for (m, ensemble_filename) in enumerate(ensemble_filenames):
        checkpoint.restore(ensemble_filename)
        for (n, (name, test_dataset)) in enumerate(test_datasets.items()):
            filename = '{dataset}_{member}.npy'.format(dataset=name, member=m)
            filename = os.path.join(FLAGS.output_dir, filename)
            if (not tf.io.gfile.exists(filename)):
                logits = []
                test_iterator = iter(test_dataset)
                for _ in range(steps_per_eval):
                    (features, _) = next(test_iterator)
                    logits.append(model(features, training=False))
                logits = tf.concat(logits, axis=0)
                with tf.io.gfile.GFile(filename, 'w') as f:
                    np.save(f, logits.numpy())
            percent = (((m * num_datasets) + (n + 1)) / (ensemble_size * num_datasets))
            message = '{:.1%} completion for prediction: ensemble member {:d}/{:d}. Dataset {:d}/{:d}'.format(percent, (m + 1), ensemble_size, (n + 1), num_datasets)
            logging.info(message)
    metrics = {'test/negative_log_likelihood': tf.keras.metrics.Mean(), 'test/gibbs_cross_entropy': tf.keras.metrics.Mean(), 'test/accuracy': tf.keras.metrics.SparseCategoricalAccuracy(), 'test/ece': ed.metrics.ExpectedCalibrationError(num_bins=FLAGS.num_bins)}
    corrupt_metrics = {}
    for name in test_datasets:
        corrupt_metrics['test/nll_{}'.format(name)] = tf.keras.metrics.Mean()
        corrupt_metrics['test/accuracy_{}'.format(name)] = tf.keras.metrics.SparseCategoricalAccuracy()
        corrupt_metrics['test/ece_{}'.format(name)] = ed.metrics.ExpectedCalibrationError(num_bins=FLAGS.num_bins)
    for (n, (name, test_dataset)) in enumerate(test_datasets.items()):
        logits_dataset = []
        for m in range(ensemble_size):
            filename = '{dataset}_{member}.npy'.format(dataset=name, member=m)
            filename = os.path.join(FLAGS.output_dir, filename)
            with tf.io.gfile.GFile(filename, 'rb') as f:
                logits_dataset.append(np.load(f))
        logits_dataset = tf.convert_to_tensor(logits_dataset)
        test_iterator = iter(test_dataset)
        for step in range(steps_per_eval):
            (_, labels) = next(test_iterator)
            logits = logits_dataset[:, (step * batch_size):((step + 1) * batch_size)]
            labels = tf.cast(tf.reshape(labels, [(- 1)]), tf.int32)
            negative_log_likelihood = tf.reduce_mean(ensemble_negative_log_likelihood(labels, logits))
            per_probs = tf.nn.softmax(logits)
            probs = tf.reduce_mean(per_probs, axis=0)
            if (name == 'clean'):
                gibbs_ce = tf.reduce_mean(gibbs_cross_entropy(labels, logits))
                metrics['test/negative_log_likelihood'].update_state(negative_log_likelihood)
                metrics['test/gibbs_cross_entropy'].update_state(gibbs_ce)
                metrics['test/accuracy'].update_state(labels, probs)
                metrics['test/ece'].update_state(labels, probs)
            else:
                corrupt_metrics['test/nll_{}'.format(name)].update_state(negative_log_likelihood)
                corrupt_metrics['test/accuracy_{}'.format(name)].update_state(labels, probs)
                corrupt_metrics['test/ece_{}'.format(name)].update_state(labels, probs)
        message = '{:.1%} completion for evaluation: dataset {:d}/{:d}'.format(((n + 1) / num_datasets), (n + 1), num_datasets)
        logging.info(message)
    corrupt_results = utils.aggregate_corrupt_metrics(corrupt_metrics, corruption_types, max_intensity, FLAGS.alexnet_errors_path)
    total_results = {name: metric.result() for (name, metric) in metrics.items()}
    total_results.update(corrupt_results)
    logging.info('Metrics: %s', total_results)