def main(argv):
fmt = '[%(filename)s:%(lineno)s] %(message)s'
formatter = logging.PythonFormatter(fmt)
logging.get_absl_handler().setFormatter(formatter)
del argv # unused arg
tf.io.gfile.makedirs(FLAGS.output_dir)
logging.info('Saving checkpoints at %s', FLAGS.output_dir)
tf.random.set_seed(FLAGS.seed)
data_dir = FLAGS.data_dir
if FLAGS.use_gpu:
logging.info('Use GPU')
strategy = tf.distribute.MirroredStrategy()
else:
logging.info('Use TPU at %s',
FLAGS.tpu if FLAGS.tpu is not None else 'local')
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu=FLAGS.tpu)
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.TPUStrategy(resolver)
ds_info = tfds.builder(FLAGS.dataset).info
batch_size = FLAGS.per_core_batch_size * FLAGS.num_cores
train_dataset_size = (
ds_info.splits['train'].num_examples * FLAGS.train_proportion)
steps_per_epoch = int(train_dataset_size / batch_size)
logging.info('Steps per epoch %s', steps_per_epoch)
logging.info('Size of the dataset %s', ds_info.splits['train'].num_examples)
logging.info('Train proportion %s', FLAGS.train_proportion)
steps_per_eval = ds_info.splits['test'].num_examples // batch_size
num_classes = ds_info.features['label'].num_classes
aug_params = {
'augmix': FLAGS.augmix,
'aug_count': FLAGS.aug_count,
'augmix_depth': FLAGS.augmix_depth,
'augmix_prob_coeff': FLAGS.augmix_prob_coeff,
'augmix_width': FLAGS.augmix_width,
}
# Note that stateless_{fold_in,split} may incur a performance cost, but a
# quick side-by-side test seemed to imply this was minimal.
seeds = tf.random.experimental.stateless_split(
[FLAGS.seed, FLAGS.seed + 1], 2)[:, 0]
train_builder = ub.datasets.get(
FLAGS.dataset,
data_dir=data_dir,
download_data=FLAGS.download_data,
split=tfds.Split.TRAIN,
seed=seeds[0],
aug_params=aug_params,
shuffle_buffer_size=FLAGS.shuffle_buffer_size,
validation_percent=1. - FLAGS.train_proportion,
)
train_dataset = train_builder.load(batch_size=batch_size)
validation_dataset = None
steps_per_validation = 0
if FLAGS.train_proportion < 1.0:
validation_builder = ub.datasets.get(
FLAGS.dataset,
split=tfds.Split.VALIDATION,
validation_percent=1. - FLAGS.train_proportion,
data_dir=data_dir,
drop_remainder=FLAGS.drop_remainder_for_eval)
validation_dataset = validation_builder.load(batch_size=batch_size)
validation_dataset = strategy.experimental_distribute_dataset(
validation_dataset)
steps_per_validation = validation_builder.num_examples // batch_size
clean_test_builder = ub.datasets.get(
FLAGS.dataset,
split=tfds.Split.TEST,
data_dir=data_dir,
drop_remainder=FLAGS.drop_remainder_for_eval)
clean_test_dataset = clean_test_builder.load(batch_size=batch_size)
test_datasets = {
'clean': strategy.experimental_distribute_dataset(clean_test_dataset),
}
train_dataset = strategy.experimental_distribute_dataset(train_dataset)
steps_per_epoch = train_builder.num_examples // batch_size
steps_per_eval = clean_test_builder.num_examples // batch_size
num_classes = 100 if FLAGS.dataset == 'cifar100' else 10
if FLAGS.eval_on_ood:
ood_dataset_names = FLAGS.ood_dataset
ood_ds, steps_per_ood = ood_utils.load_ood_datasets(
ood_dataset_names,
clean_test_builder,
1. - FLAGS.train_proportion,
batch_size,
drop_remainder=FLAGS.drop_remainder_for_eval)
ood_datasets = {
name: strategy.experimental_distribute_dataset(ds)
for name, ds in ood_ds.items()
}
if FLAGS.corruptions_interval > 0:
if FLAGS.dataset == 'cifar100':
data_dir = FLAGS.cifar100_c_path
corruption_types, _ = utils.load_corrupted_test_info(FLAGS.dataset)
for corruption_type in corruption_types:
for severity in range(1, 6):
dataset = ub.datasets.get(
f'{FLAGS.dataset}_corrupted',
corruption_type=corruption_type,
severity=severity,
split=tfds.Split.TEST,
data_dir=data_dir).load(batch_size=batch_size)
test_datasets[f'{corruption_type}_{severity}'] = (
strategy.experimental_distribute_dataset(dataset))
summary_writer = tf.summary.create_file_writer(
os.path.join(FLAGS.output_dir, 'summaries'))
with strategy.scope():
logging.info('Building ResNet model')
model = ub.models.wide_resnet(
input_shape=(32, 32, 3),
depth=28,
width_multiplier=10,
num_classes=num_classes,
l2=FLAGS.l2,
hps=_extract_hyperparameter_dictionary(),
seed=seeds[1])
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())
# Linearly scale learning rate and the decay epochs by vanilla settings.
base_lr = FLAGS.base_learning_rate * batch_size / 128
lr_decay_epochs = [(int(start_epoch_str) * FLAGS.train_epochs) // 200
for start_epoch_str in FLAGS.lr_decay_epochs]
lr_schedule = ub.schedules.WarmUpPiecewiseConstantSchedule(
steps_per_epoch,
base_lr,
decay_ratio=FLAGS.lr_decay_ratio,
decay_epochs=lr_decay_epochs,
warmup_epochs=FLAGS.lr_warmup_epochs)
optimizer = tf.keras.optimizers.SGD(lr_schedule,
momentum=1.0 - FLAGS.one_minus_momentum,
nesterov=True)
metrics = {
'train/negative_log_likelihood':
tf.keras.metrics.Mean(),
'train/accuracy':
tf.keras.metrics.SparseCategoricalAccuracy(),
'train/loss':
tf.keras.metrics.Mean(),
'train/ece':
rm.metrics.ExpectedCalibrationError(num_bins=FLAGS.num_bins),
'test/negative_log_likelihood':
tf.keras.metrics.Mean(),
'test/accuracy':
tf.keras.metrics.SparseCategoricalAccuracy(),
'test/ece':
rm.metrics.ExpectedCalibrationError(num_bins=FLAGS.num_bins),
}
if validation_dataset:
metrics.update({
'validation/negative_log_likelihood': tf.keras.metrics.Mean(),
'validation/accuracy': tf.keras.metrics.SparseCategoricalAccuracy(),
'validation/ece': rm.metrics.ExpectedCalibrationError(
num_bins=FLAGS.num_bins),
})
if FLAGS.eval_on_ood:
ood_metrics = ood_utils.create_ood_metrics(ood_dataset_names)
metrics.update(ood_metrics)
if FLAGS.corruptions_interval > 0:
corrupt_metrics = {}
for intensity in range(1, 6):
for corruption in corruption_types:
dataset_name = '{0}_{1}'.format(corruption, intensity)
corrupt_metrics['test/nll_{}'.format(dataset_name)] = (
tf.keras.metrics.Mean())
corrupt_metrics['test/accuracy_{}'.format(dataset_name)] = (
tf.keras.metrics.SparseCategoricalAccuracy())
corrupt_metrics['test/ece_{}'.format(dataset_name)] = (
rm.metrics.ExpectedCalibrationError(num_bins=FLAGS.num_bins))
checkpoint = tf.train.Checkpoint(model=model, optimizer=optimizer)
latest_checkpoint = tf.train.latest_checkpoint(FLAGS.output_dir)
initial_epoch = 0
if latest_checkpoint:
# checkpoint.restore must be within a strategy.scope() so that optimizer
# slot variables are mirrored.
checkpoint.restore(latest_checkpoint)
logging.info('Loaded checkpoint %s', latest_checkpoint)
initial_epoch = optimizer.iterations.numpy() // steps_per_epoch
if FLAGS.saved_model_dir:
logging.info('Saved model dir : %s', FLAGS.saved_model_dir)
latest_checkpoint = tf.train.latest_checkpoint(FLAGS.saved_model_dir)
checkpoint.restore(latest_checkpoint)
logging.info('Loaded checkpoint %s', latest_checkpoint)
if FLAGS.eval_only:
initial_epoch = FLAGS.train_epochs - 1 # Run just one epoch of eval
@tf.function
def train_step(iterator):
"""Training StepFn."""
def step_fn(inputs):
"""Per-Replica StepFn."""
images = inputs['features']
labels = inputs['labels']
if FLAGS.augmix and FLAGS.aug_count >= 1:
# Index 0 at augmix processing is the unperturbed image.
# We take just 1 augmented image from the returned augmented images.
images = images[:, 1, ...]
with tf.GradientTape() as tape:
logits = model(images, training=True)
if FLAGS.label_smoothing == 0.:
negative_log_likelihood = tf.reduce_mean(
tf.keras.losses.sparse_categorical_crossentropy(labels,
logits,
from_logits=True))
else:
one_hot_labels = tf.one_hot(tf.cast(labels, tf.int32), num_classes)
negative_log_likelihood = tf.reduce_mean(
tf.keras.losses.categorical_crossentropy(
one_hot_labels,
logits,
from_logits=True,
label_smoothing=FLAGS.label_smoothing))
l2_loss = sum(model.losses)
loss = negative_log_likelihood + l2_loss
# Scale the loss given the TPUStrategy will reduce sum all gradients.
scaled_loss = loss / strategy.num_replicas_in_sync
grads = tape.gradient(scaled_loss, model.trainable_variables)
optimizer.apply_gradients(zip(grads, model.trainable_variables))
probs = tf.nn.softmax(logits)
metrics['train/ece'].add_batch(probs, label=labels)
metrics['train/loss'].update_state(loss)
metrics['train/negative_log_likelihood'].update_state(
negative_log_likelihood)
metrics['train/accuracy'].update_state(labels, logits)
for _ in tf.range(tf.cast(steps_per_epoch, tf.int32)):
strategy.run(step_fn, args=(next(iterator),))
@tf.function
def test_step(iterator, dataset_split, dataset_name, num_steps):
"""Evaluation StepFn."""
def step_fn(inputs):
"""Per-Replica StepFn."""
images = inputs['features']
labels = inputs['labels']
logits = model(images, training=False)
probs = tf.nn.softmax(logits)
negative_log_likelihood = tf.reduce_mean(
tf.keras.losses.sparse_categorical_crossentropy(labels, probs))
if dataset_name == 'clean':
metrics[f'{dataset_split}/negative_log_likelihood'].update_state(
negative_log_likelihood)
metrics[f'{dataset_split}/accuracy'].update_state(labels, probs)
metrics[f'{dataset_split}/ece'].add_batch(probs, label=labels)
elif dataset_name.startswith('ood/'):
ood_labels = 1 - inputs['is_in_distribution']
if FLAGS.dempster_shafer_ood:
ood_scores = ood_utils.DempsterShaferUncertainty(logits)
else:
ood_scores = 1 - tf.reduce_max(probs, axis=-1)
# Edgecase for if dataset_name contains underscores
for name, metric in metrics.items():
if dataset_name in name:
metric.update_state(ood_labels, ood_scores)
else:
corrupt_metrics['test/nll_{}'.format(dataset_name)].update_state(
negative_log_likelihood)
corrupt_metrics['test/accuracy_{}'.format(dataset_name)].update_state(
labels, probs)
corrupt_metrics['test/ece_{}'.format(dataset_name)].add_batch(
probs, label=labels)
for _ in tf.range(tf.cast(num_steps, tf.int32)):
strategy.run(step_fn, args=(next(iterator),))
@tf.function
def cifar10h_test_step(iterator, num_steps):
"""Evaluation StepFn."""
def step_fn(inputs):
"""Per-Replica StepFn."""
images = inputs['features']
labels = inputs['labels']
logits = model(images, training=False)
negative_log_likelihood = tf.keras.losses.CategoricalCrossentropy(
from_logits=True,
reduction=tf.keras.losses.Reduction.NONE)(labels, logits)
negative_log_likelihood = tf.reduce_mean(negative_log_likelihood)
metrics['cifar10h/nll'].update_state(negative_log_likelihood)
label_diversity, sample_diversity, ged = _generalized_energy_distance(
labels, tf.nn.softmax(logits), 10)
metrics['cifar10h/ged'].update_state(ged)
metrics['cifar10h/ged_label_diversity'].update_state(
tf.reduce_mean(label_diversity))
metrics['cifar10h/ged_sample_diversity'].update_state(
tf.reduce_mean(sample_diversity))
for _ in tf.range(tf.cast(num_steps, tf.int32)):
strategy.run(step_fn, args=(next(iterator),))
metrics.update({'test/ms_per_example': tf.keras.metrics.Mean()})
metrics.update({'train/ms_per_example': tf.keras.metrics.Mean()})
train_iterator = iter(train_dataset)
start_time = time.time()
tb_callback = None
if FLAGS.collect_profile:
tb_callback = tf.keras.callbacks.TensorBoard(
profile_batch=(100, 102),
log_dir=os.path.join(FLAGS.output_dir, 'logs'))
tb_callback.set_model(model)
for epoch in range(initial_epoch, FLAGS.train_epochs):
logging.info('Starting to run epoch: %s', epoch)
if tb_callback:
tb_callback.on_epoch_begin(epoch)
if not FLAGS.eval_only:
train_start_time = time.time()
train_step(train_iterator)
ms_per_example = (time.time() - train_start_time) * 1e6 / batch_size
metrics['train/ms_per_example'].update_state(ms_per_example)
current_step = (epoch + 1) * steps_per_epoch
max_steps = steps_per_epoch * FLAGS.train_epochs
time_elapsed = time.time() - start_time
steps_per_sec = float(current_step) / time_elapsed
eta_seconds = (max_steps - current_step) / steps_per_sec
message = ('{:.1%} completion: epoch {:d}/{:d}. {:.1f} steps/s. '
'ETA: {:.0f} min. Time elapsed: {:.0f} min'.format(
current_step / max_steps, epoch + 1, FLAGS.train_epochs,
steps_per_sec, eta_seconds / 60, time_elapsed / 60))
logging.info(message)
if tb_callback:
tb_callback.on_epoch_end(epoch)
if validation_dataset:
validation_iterator = iter(validation_dataset)
test_step(
validation_iterator, 'validation', 'clean', steps_per_validation)
datasets_to_evaluate = {'clean': test_datasets['clean']}
if (FLAGS.corruptions_interval > 0 and
(epoch + 1) % FLAGS.corruptions_interval == 0):
datasets_to_evaluate = test_datasets
for dataset_name, test_dataset in datasets_to_evaluate.items():
test_iterator = iter(test_dataset)
logging.info('Testing on dataset %s', dataset_name)
logging.info('Starting to run eval at epoch: %s', epoch)
test_start_time = time.time()
test_step(test_iterator, 'test', dataset_name, steps_per_eval)
ms_per_example = (time.time() - test_start_time) * 1e6 / batch_size
metrics['test/ms_per_example'].update_state(ms_per_example)
logging.info('Done with testing on %s', dataset_name)
if FLAGS.eval_on_ood:
for ood_dataset_name, ood_dataset in ood_datasets.items():
ood_iterator = iter(ood_dataset)
logging.info('Calculating OOD on dataset %s', ood_dataset_name)
logging.info('Running OOD eval at epoch: %s', epoch)
test_step(ood_iterator, 'test', ood_dataset_name,
steps_per_ood[ood_dataset_name])
logging.info('Done with OOD eval on %s', ood_dataset_name)
corrupt_results = {}
if (FLAGS.corruptions_interval > 0 and
(epoch + 1) % FLAGS.corruptions_interval == 0):
corrupt_results = utils.aggregate_corrupt_metrics(corrupt_metrics,
corruption_types)
logging.info('Train Loss: %.4f, Accuracy: %.2f%%',
metrics['train/loss'].result(),
metrics['train/accuracy'].result() * 100)
logging.info('Test NLL: %.4f, Accuracy: %.2f%%',
metrics['test/negative_log_likelihood'].result(),
metrics['test/accuracy'].result() * 100)
total_results = {name: metric.result() for name, metric in metrics.items()}
total_results.update(corrupt_results)
# Metrics from Robustness Metrics (like ECE) will return a dict with a
# single key/value, instead of a scalar.
total_results = {
k: (list(v.values())[0] if isinstance(v, dict) else v)
for k, v in total_results.items()
}
with summary_writer.as_default():
for name, result in total_results.items():
tf.summary.scalar(name, result, step=epoch + 1)
for metric in metrics.values():
metric.reset_states()
if FLAGS.corruptions_interval > 0:
for metric in corrupt_metrics.values():
metric.reset_states()
if (FLAGS.checkpoint_interval > 0 and
(epoch + 1) % FLAGS.checkpoint_interval == 0):
checkpoint_name = checkpoint.save(
os.path.join(FLAGS.output_dir, 'checkpoint'))
logging.info('Saved checkpoint to %s', checkpoint_name)
final_checkpoint_name = checkpoint.save(
os.path.join(FLAGS.output_dir, 'checkpoint'))
logging.info('Saved last checkpoint to %s', final_checkpoint_name)
with summary_writer.as_default():
hp.hparams({
'base_learning_rate': FLAGS.base_learning_rate,
'one_minus_momentum': FLAGS.one_minus_momentum,
'l2': FLAGS.l2,
})
def main(argv):
del argv # unused arg
tf.io.gfile.makedirs(FLAGS.output_dir)
logging.info('Saving checkpoints at %s', FLAGS.output_dir)
tf.random.set_seed(FLAGS.seed)
# Split the seed into a 2-tuple, for passing into dataset builder.
dataset_seed = (FLAGS.seed, FLAGS.seed + 1)
data_dir = FLAGS.data_dir
if FLAGS.use_gpu:
logging.info('Use GPU')
strategy = tf.distribute.MirroredStrategy()
else:
logging.info('Use TPU at %s',
FLAGS.tpu if FLAGS.tpu is not None else 'local')
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu=FLAGS.tpu)
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.TPUStrategy(resolver)
batch_size = FLAGS.total_batch_size // FLAGS.num_dropout_samples_training
test_batch_size = FLAGS.total_batch_size
num_classes = 10 if FLAGS.dataset == 'cifar10' else 100
aug_params = {
'augmix': FLAGS.augmix,
'aug_count': FLAGS.aug_count,
'augmix_depth': FLAGS.augmix_depth,
'augmix_prob_coeff': FLAGS.augmix_prob_coeff,
'augmix_width': FLAGS.augmix_width,
'ensemble_size': 1,
'mixup_alpha': FLAGS.mixup_alpha,
}
validation_proportion = 1. - FLAGS.train_proportion
use_validation_set = validation_proportion > 0.
if FLAGS.dataset == 'cifar10':
dataset_builder_class = ub.datasets.Cifar10Dataset
else:
dataset_builder_class = ub.datasets.Cifar100Dataset
train_dataset_builder = dataset_builder_class(
data_dir=data_dir,
download_data=FLAGS.download_data,
split=tfds.Split.TRAIN,
use_bfloat16=FLAGS.use_bfloat16,
aug_params=aug_params,
validation_percent=validation_proportion,
shuffle_buffer_size=FLAGS.shuffle_buffer_size,
seed=dataset_seed)
train_dataset = train_dataset_builder.load(batch_size=batch_size)
if validation_proportion > 0.:
validation_dataset_builder = dataset_builder_class(
data_dir=data_dir,
download_data=FLAGS.download_data,
split=tfds.Split.VALIDATION,
use_bfloat16=FLAGS.use_bfloat16,
validation_percent=validation_proportion,
drop_remainder=FLAGS.drop_remainder_for_eval)
validation_dataset = validation_dataset_builder.load(
batch_size=test_batch_size)
validation_dataset = strategy.experimental_distribute_dataset(
validation_dataset)
val_sample_size = validation_dataset_builder.num_examples
steps_per_val = steps_per_epoch = int(val_sample_size / test_batch_size)
clean_test_dataset_builder = dataset_builder_class(
data_dir=data_dir,
download_data=FLAGS.download_data,
split=tfds.Split.TEST,
use_bfloat16=FLAGS.use_bfloat16,
drop_remainder=FLAGS.drop_remainder_for_eval)
clean_test_dataset = clean_test_dataset_builder.load(
batch_size=test_batch_size)
steps_per_epoch = train_dataset_builder.num_examples // batch_size
steps_per_eval = clean_test_dataset_builder.num_examples // test_batch_size
train_dataset = strategy.experimental_distribute_dataset(train_dataset)
test_datasets = {
'clean': strategy.experimental_distribute_dataset(clean_test_dataset),
}
if FLAGS.eval_on_ood:
ood_dataset_names = FLAGS.ood_dataset
ood_ds, steps_per_ood = ood_utils.load_ood_datasets(
ood_dataset_names,
clean_test_dataset_builder,
validation_proportion,
test_batch_size,
drop_remainder=FLAGS.drop_remainder_for_eval)
ood_datasets = {
name: strategy.experimental_distribute_dataset(ds)
for name, ds in ood_ds.items()
}
if FLAGS.corruptions_interval > 0:
if FLAGS.dataset == 'cifar100':
data_dir = FLAGS.cifar100_c_path
corruption_types, _ = utils.load_corrupted_test_info(FLAGS.dataset)
for corruption_type in corruption_types:
for severity in range(1, 6):
dataset = ub.datasets.get(
f'{FLAGS.dataset}_corrupted',
corruption_type=corruption_type,
severity=severity,
split=tfds.Split.TEST,
data_dir=data_dir,
drop_remainder=FLAGS.drop_remainder_for_eval).load(
batch_size=batch_size)
test_datasets[f'{corruption_type}_{severity}'] = (
strategy.experimental_distribute_dataset(dataset))
if FLAGS.use_bfloat16:
tf.keras.mixed_precision.set_global_policy('mixed_bfloat16')
summary_writer = tf.summary.create_file_writer(
os.path.join(FLAGS.output_dir, 'summaries'))
with strategy.scope():
logging.info('Building ResNet model')
if FLAGS.use_spec_norm:
logging.info('Use Spectral Normalization with norm bound %.2f',
FLAGS.spec_norm_bound)
if FLAGS.use_gp_layer:
logging.info('Use GP layer with hidden units %d', FLAGS.gp_hidden_dim)
model = ub.models.wide_resnet_sngp(
input_shape=(32, 32, 3),
batch_size=batch_size,
depth=28,
width_multiplier=10,
num_classes=num_classes,
l2=FLAGS.l2,
use_mc_dropout=FLAGS.use_mc_dropout,
use_filterwise_dropout=FLAGS.use_filterwise_dropout,
dropout_rate=FLAGS.dropout_rate,
use_gp_layer=FLAGS.use_gp_layer,
gp_input_dim=FLAGS.gp_input_dim,
gp_hidden_dim=FLAGS.gp_hidden_dim,
gp_scale=FLAGS.gp_scale,
gp_bias=FLAGS.gp_bias,
gp_input_normalization=FLAGS.gp_input_normalization,
gp_random_feature_type=FLAGS.gp_random_feature_type,
gp_cov_discount_factor=FLAGS.gp_cov_discount_factor,
gp_cov_ridge_penalty=FLAGS.gp_cov_ridge_penalty,
use_spec_norm=FLAGS.use_spec_norm,
spec_norm_iteration=FLAGS.spec_norm_iteration,
spec_norm_bound=FLAGS.spec_norm_bound)
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())
# Linearly scale learning rate and the decay epochs by vanilla settings.
base_lr = FLAGS.base_learning_rate * batch_size / 128
lr_decay_epochs = [(int(start_epoch_str) * FLAGS.train_epochs) // 200
for start_epoch_str in FLAGS.lr_decay_epochs]
lr_schedule = ub.schedules.WarmUpPiecewiseConstantSchedule(
steps_per_epoch,
base_lr,
decay_ratio=FLAGS.lr_decay_ratio,
decay_epochs=lr_decay_epochs,
warmup_epochs=FLAGS.lr_warmup_epochs)
optimizer = tf.keras.optimizers.SGD(lr_schedule,
momentum=1.0 - FLAGS.one_minus_momentum,
nesterov=True)
metrics = {
'train/negative_log_likelihood': tf.keras.metrics.Mean(),
'train/accuracy': tf.keras.metrics.SparseCategoricalAccuracy(),
'train/loss': tf.keras.metrics.Mean(),
'train/ece': rm.metrics.ExpectedCalibrationError(
num_bins=FLAGS.num_bins),
'test/negative_log_likelihood': tf.keras.metrics.Mean(),
'test/accuracy': tf.keras.metrics.SparseCategoricalAccuracy(),
'test/ece': rm.metrics.ExpectedCalibrationError(
num_bins=FLAGS.num_bins),
'test/stddev': tf.keras.metrics.Mean(),
}
if use_validation_set:
metrics.update({
'val/negative_log_likelihood': tf.keras.metrics.Mean(),
'val/accuracy': tf.keras.metrics.SparseCategoricalAccuracy(),
'val/ece': rm.metrics.ExpectedCalibrationError(
num_bins=FLAGS.num_bins),
'val/stddev': tf.keras.metrics.Mean(),
})
if FLAGS.eval_on_ood:
ood_metrics = ood_utils.create_ood_metrics(
ood_dataset_names, tpr_list=FLAGS.ood_tpr_threshold)
metrics.update(ood_metrics)
if FLAGS.corruptions_interval > 0:
corrupt_metrics = {}
for intensity in range(1, 6):
for corruption in corruption_types:
dataset_name = '{0}_{1}'.format(corruption, intensity)
corrupt_metrics['test/nll_{}'.format(dataset_name)] = (
tf.keras.metrics.Mean())
corrupt_metrics['test/accuracy_{}'.format(dataset_name)] = (
tf.keras.metrics.SparseCategoricalAccuracy())
corrupt_metrics['test/ece_{}'.format(dataset_name)] = (
rm.metrics.ExpectedCalibrationError(num_bins=FLAGS.num_bins))
corrupt_metrics['test/stddev_{}'.format(dataset_name)] = (
tf.keras.metrics.Mean())
checkpoint = tf.train.Checkpoint(model=model, optimizer=optimizer)
latest_checkpoint = tf.train.latest_checkpoint(FLAGS.output_dir)
initial_epoch = 0
logging.info('Output dir : %s', FLAGS.output_dir)
if latest_checkpoint:
# checkpoint.restore must be within a strategy.scope() so that optimizer
# slot variables are mirrored.
checkpoint.restore(latest_checkpoint)
logging.info('Loaded checkpoint %s', latest_checkpoint)
initial_epoch = optimizer.iterations.numpy() // steps_per_epoch
if FLAGS.saved_model_dir:
logging.info('Saved model dir : %s', FLAGS.saved_model_dir)
latest_checkpoint = tf.train.latest_checkpoint(FLAGS.saved_model_dir)
checkpoint.restore(latest_checkpoint)
logging.info('Loaded checkpoint %s', latest_checkpoint)
if FLAGS.eval_only:
initial_epoch = FLAGS.train_epochs - 1 # Run just one epoch of eval
@tf.function
def train_step(iterator, step):
"""Training StepFn."""
def step_fn(inputs, step):
"""Per-Replica StepFn."""
images = inputs['features']
labels = inputs['labels']
if tf.equal(step, 0) and FLAGS.gp_cov_discount_factor < 0:
# Resetting covaraince estimator at the begining of a new epoch.
if FLAGS.use_gp_layer:
model.layers[-1].reset_covariance_matrix()
if FLAGS.augmix and FLAGS.aug_count >= 1:
# Index 0 at augmix preprocessing is the unperturbed image.
images = images[:, 1, ...]
# This is for the case of combining AugMix and Mixup.
if FLAGS.mixup_alpha > 0:
labels = tf.split(labels, FLAGS.aug_count + 1, axis=0)[1]
images = tf.tile(images, [FLAGS.num_dropout_samples_training, 1, 1, 1])
if FLAGS.mixup_alpha > 0:
labels = tf.tile(labels, [FLAGS.num_dropout_samples_training, 1])
else:
labels = tf.tile(labels, [FLAGS.num_dropout_samples_training])
with tf.GradientTape() as tape:
logits = model(images, training=True)
if isinstance(logits, (list, tuple)):
# If model returns a tuple of (logits, covmat), extract logits
logits, _ = logits
if FLAGS.use_bfloat16:
logits = tf.cast(logits, tf.float32)
if FLAGS.mixup_alpha > 0:
negative_log_likelihood = tf.reduce_mean(
tf.keras.losses.categorical_crossentropy(labels,
logits,
from_logits=True))
else:
negative_log_likelihood = tf.reduce_mean(
tf.keras.losses.sparse_categorical_crossentropy(labels,
logits,
from_logits=True))
l2_loss = sum(model.losses)
loss = negative_log_likelihood + l2_loss
# Scale the loss given the TPUStrategy will reduce sum all gradients.
scaled_loss = loss / strategy.num_replicas_in_sync
grads = tape.gradient(scaled_loss, model.trainable_variables)
optimizer.apply_gradients(zip(grads, model.trainable_variables))
probs = tf.nn.softmax(logits)
if FLAGS.mixup_alpha > 0:
labels = tf.argmax(labels, axis=-1)
metrics['train/ece'].add_batch(probs, label=labels)
metrics['train/loss'].update_state(loss)
metrics['train/negative_log_likelihood'].update_state(
negative_log_likelihood)
metrics['train/accuracy'].update_state(labels, logits)
strategy.run(step_fn, args=(next(iterator), step))
@tf.function
def test_step(iterator, dataset_name, num_steps):
"""Evaluation StepFn."""
def step_fn(inputs):
"""Per-Replica StepFn."""
images = inputs['features']
labels = inputs['labels']
logits_list = []
stddev_list = []
for _ in range(FLAGS.num_dropout_samples):
logits = model(images, training=False)
if isinstance(logits, (list, tuple)):
# If model returns a tuple of (logits, covmat), extract both
logits, covmat = logits
if FLAGS.use_bfloat16:
logits = tf.cast(logits, tf.float32)
logits = ed.layers.utils.mean_field_logits(
logits, covmat, mean_field_factor=FLAGS.gp_mean_field_factor)
else:
covmat = tf.eye(logits.shape[0])
if FLAGS.use_bfloat16:
logits = tf.cast(logits, tf.float32)
stddev = tf.sqrt(tf.linalg.diag_part(covmat))
stddev_list.append(stddev)
logits_list.append(logits)
# Logits dimension is (num_samples, batch_size, num_classes).
logits_list = tf.stack(logits_list, axis=0)
stddev_list = tf.stack(stddev_list, axis=0)
stddev = tf.reduce_mean(stddev_list, axis=0)
probs_list = tf.nn.softmax(logits_list)
probs = tf.reduce_mean(probs_list, axis=0)
logits = tf.reduce_mean(logits_list, axis=0)
labels_broadcasted = tf.broadcast_to(
labels, [FLAGS.num_dropout_samples,
tf.shape(labels)[0]])
log_likelihoods = -tf.keras.losses.sparse_categorical_crossentropy(
labels_broadcasted, logits_list, from_logits=True)
negative_log_likelihood = tf.reduce_mean(
-tf.reduce_logsumexp(log_likelihoods, axis=[0]) +
tf.math.log(float(FLAGS.num_dropout_samples)))
logging.info('Dataset name : %s', dataset_name)
if dataset_name == 'clean':
metrics['test/negative_log_likelihood'].update_state(
negative_log_likelihood)
metrics['test/accuracy'].update_state(labels, probs)
metrics['test/ece'].add_batch(probs, label=labels)
metrics['test/stddev'].update_state(stddev)
elif dataset_name == 'val':
metrics['val/negative_log_likelihood'].update_state(
negative_log_likelihood)
metrics['val/accuracy'].update_state(labels, probs)
metrics['val/ece'].add_batch(probs, label=labels)
metrics['val/stddev'].update_state(stddev)
elif dataset_name.startswith('ood/'):
ood_labels = 1 - inputs['is_in_distribution']
if FLAGS.dempster_shafer_ood:
ood_scores = ood_utils.DempsterShaferUncertainty(logits)
else:
ood_scores = 1 - tf.reduce_max(probs, axis=-1)
# Edgecase for if dataset_name contains underscores
for name, metric in metrics.items():
if dataset_name in name:
metric.update_state(ood_labels, ood_scores)
elif FLAGS.corruptions_interval > 0:
corrupt_metrics['test/nll_{}'.format(dataset_name)].update_state(
negative_log_likelihood)
corrupt_metrics['test/accuracy_{}'.format(dataset_name)].update_state(
labels, probs)
corrupt_metrics['test/ece_{}'.format(dataset_name)].add_batch(
probs, label=labels)
corrupt_metrics['test/stddev_{}'.format(dataset_name)].update_state(
stddev)
for _ in tf.range(tf.cast(num_steps, tf.int32)):
strategy.run(step_fn, args=(next(iterator),))
metrics.update({'test/ms_per_example': tf.keras.metrics.Mean()})
step_variable = tf.Variable(0, dtype=tf.int32)
train_iterator = iter(train_dataset)
start_time = time.time()
for epoch in range(initial_epoch, FLAGS.train_epochs):
logging.info('Starting to run epoch: %s', epoch)
if not FLAGS.eval_only:
for step in range(steps_per_epoch):
step_variable.assign(step)
# Pass `step` as a tf.Variable to train_step to prevent the tf.function
# train_step() re-compiling itself at each function call.
train_step(train_iterator, step_variable)
current_step = epoch * steps_per_epoch + (step + 1)
max_steps = steps_per_epoch * FLAGS.train_epochs
time_elapsed = time.time() - start_time
steps_per_sec = float(current_step) / time_elapsed
eta_seconds = (max_steps - current_step) / steps_per_sec
message = ('{:.1%} completion: epoch {:d}/{:d}. {:.1f} steps/s. '
'ETA: {:.0f} min. Time elapsed: {:.0f} min'.format(
current_step / max_steps,
epoch + 1,
FLAGS.train_epochs,
steps_per_sec,
eta_seconds / 60,
time_elapsed / 60))
if step % 20 == 0:
logging.info(message)
datasets_to_evaluate = {'clean': test_datasets['clean']}
if use_validation_set:
datasets_to_evaluate['val'] = validation_dataset
if (FLAGS.corruptions_interval > 0 and
(epoch + 1) % FLAGS.corruptions_interval == 0):
datasets_to_evaluate = test_datasets
for dataset_name, test_dataset in datasets_to_evaluate.items():
test_iterator = iter(test_dataset)
logging.info('Testing on dataset %s', dataset_name)
steps_per_eval = steps_per_val if dataset_name == 'val' else steps_per_eval
logging.info('Starting to run eval at epoch: %s', epoch)
test_start_time = time.time()
test_step(test_iterator, dataset_name, steps_per_eval)
ms_per_example = (time.time() - test_start_time) * 1e6 / batch_size
metrics['test/ms_per_example'].update_state(ms_per_example)
logging.info('Done with testing on %s', dataset_name)
if FLAGS.eval_on_ood:
for ood_dataset_name, ood_dataset in ood_datasets.items():
ood_iterator = iter(ood_dataset)
logging.info('Calculating OOD on dataset %s', ood_dataset_name)
logging.info('Running OOD eval at epoch: %s', epoch)
test_step(ood_iterator, ood_dataset_name,
steps_per_ood[ood_dataset_name])
logging.info('Done with OOD eval on %s', dataset_name)
corrupt_results = {}
if (FLAGS.corruptions_interval > 0 and
(epoch + 1) % FLAGS.corruptions_interval == 0):
corrupt_results = utils.aggregate_corrupt_metrics(corrupt_metrics,
corruption_types)
logging.info('Train Loss: %.4f, Accuracy: %.2f%%',
metrics['train/loss'].result(),
metrics['train/accuracy'].result() * 100)
if use_validation_set:
logging.info('Val NLL: %.4f, Accuracy: %.2f%%',
metrics['val/negative_log_likelihood'].result(),
metrics['val/accuracy'].result() * 100)
logging.info('Test NLL: %.4f, Accuracy: %.2f%%',
metrics['test/negative_log_likelihood'].result(),
metrics['test/accuracy'].result() * 100)
total_results = {name: metric.result() for name, metric in metrics.items()}
total_results.update(corrupt_results)
# Metrics from Robustness Metrics (like ECE) will return a dict with a
# single key/value, instead of a scalar.
total_results = {
k: (list(v.values())[0] if isinstance(v, dict) else v)
for k, v in total_results.items()
}
with summary_writer.as_default():
for name, result in total_results.items():
tf.summary.scalar(name, result, step=epoch + 1)
for metric in metrics.values():
metric.reset_states()
if FLAGS.corruptions_interval > 0:
for metric in corrupt_metrics.values():
metric.reset_states()
if (FLAGS.checkpoint_interval > 0 and
(epoch + 1) % FLAGS.checkpoint_interval == 0):
checkpoint_name = checkpoint.save(
os.path.join(FLAGS.output_dir, 'checkpoint'))
logging.info('Saved checkpoint to %s', checkpoint_name)
final_checkpoint_name = checkpoint.save(
os.path.join(FLAGS.output_dir, 'checkpoint'))
logging.info('Saved last checkpoint to %s', final_checkpoint_name)
final_save_name = os.path.join(FLAGS.output_dir, 'model')
model.save(final_save_name)
logging.info('Saved model to %s', final_save_name)
with summary_writer.as_default():
hp.hparams({
'base_learning_rate': FLAGS.base_learning_rate,
'one_minus_momentum': FLAGS.one_minus_momentum,
'l2': FLAGS.l2,
'gp_mean_field_factor': FLAGS.gp_mean_field_factor,
})
def main(argv):
del argv # unused arg
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.random.set_seed(FLAGS.seed)
tf.io.gfile.makedirs(FLAGS.output_dir)
logging.info('output_dir=%s', FLAGS.output_dir)
ds_info = tfds.builder(FLAGS.dataset).info
batch_size = FLAGS.per_core_batch_size * FLAGS.num_cores
steps_per_eval = ds_info.splits['test'].num_examples // batch_size
num_classes = ds_info.features['label'].num_classes
data_dir = FLAGS.data_dir
dataset_builder = ub.datasets.get(
FLAGS.dataset,
download_data=FLAGS.download_data,
data_dir=data_dir,
split=tfds.Split.TEST,
drop_remainder=FLAGS.drop_remainder_for_eval)
dataset = dataset_builder.load(batch_size=batch_size)
test_datasets = {'clean': dataset}
if FLAGS.eval_on_ood:
ood_dataset_names = FLAGS.ood_dataset
ood_datasets, steps_per_ood = ood_utils.load_ood_datasets(
ood_dataset_names,
dataset_builder,
1. - FLAGS.train_proportion,
batch_size,
drop_remainder=FLAGS.drop_remainder_for_eval)
test_datasets.update(ood_datasets)
if FLAGS.dataset == 'cifar100':
data_dir = FLAGS.cifar100_c_path
corruption_types, _ = utils.load_corrupted_test_info(FLAGS.dataset)
for corruption_type in corruption_types:
for severity in range(1, 6):
dataset = ub.datasets.get(
f'{FLAGS.dataset}_corrupted',
corruption_type=corruption_type,
download_data=FLAGS.download_data,
data_dir=data_dir,
severity=severity,
split=tfds.Split.TEST,
drop_remainder=FLAGS.drop_remainder_for_eval).load(
batch_size=batch_size)
test_datasets[f'{corruption_type}_{severity}'] = dataset
model = ub.models.wide_resnet(input_shape=ds_info.features['image'].shape,
depth=28,
width_multiplier=10,
num_classes=num_classes,
l2=0.,
version=2)
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())
# Search for checkpoints from their index file; then remove the index suffix.
ensemble_filenames = parse_checkpoint_dir(FLAGS.checkpoint_dir)
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)
# Write model predictions to files.
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.replace('/', '_'),
member=m) # ood dataset name has '/'
filename = os.path.join(FLAGS.output_dir, filename)
if not tf.io.gfile.exists(filename):
logits = []
test_iterator = iter(test_dataset)
steps = steps_per_eval if 'ood/' not in name else steps_per_ood[
name]
for _ in range(steps):
features = next(test_iterator)['features'] # pytype: disable=unsupported-operands
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':
rm.metrics.ExpectedCalibrationError(num_bins=FLAGS.num_bins),
'test/diversity': rm.metrics.AveragePairwiseDiversity(),
}
if FLAGS.eval_on_ood:
ood_metrics = ood_utils.create_ood_metrics(ood_dataset_names)
metrics.update(ood_metrics)
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)] = (
rm.metrics.ExpectedCalibrationError(num_bins=FLAGS.num_bins))
for i in range(ensemble_size):
metrics['test/nll_member_{}'.format(i)] = tf.keras.metrics.Mean()
metrics['test/accuracy_member_{}'.format(i)] = (
tf.keras.metrics.SparseCategoricalAccuracy())
# Evaluate model predictions.
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.replace('/', '_'),
member=m) # ood dataset name has '/'
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)
steps = steps_per_eval if 'ood/' not in name else steps_per_ood[name]
for step in range(steps):
inputs = next(test_iterator)
labels = inputs['labels'] # pytype: disable=unsupported-operands
logits = logits_dataset[:, (step * batch_size):((step + 1) *
batch_size)]
labels = tf.cast(labels, tf.int32)
negative_log_likelihood_metric = rm.metrics.EnsembleCrossEntropy()
negative_log_likelihood_metric.add_batch(logits, labels=labels)
negative_log_likelihood = list(
negative_log_likelihood_metric.result().values())[0]
per_probs = tf.nn.softmax(logits)
probs = tf.reduce_mean(per_probs, axis=0)
logits_mean = tf.reduce_mean(logits, axis=0)
if name == 'clean':
gibbs_ce_metric = rm.metrics.GibbsCrossEntropy()
gibbs_ce_metric.add_batch(logits, labels=labels)
gibbs_ce = list(gibbs_ce_metric.result().values())[0]
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'].add_batch(probs, label=labels)
for i in range(ensemble_size):
member_probs = per_probs[i]
member_loss = tf.keras.losses.sparse_categorical_crossentropy(
labels, member_probs)
metrics['test/nll_member_{}'.format(i)].update_state(
member_loss)
metrics['test/accuracy_member_{}'.format(i)].update_state(
labels, member_probs)
metrics['test/diversity'].add_batch(per_probs)
elif name.startswith('ood/'):
ood_labels = 1 - inputs['is_in_distribution'] # pytype: disable=unsupported-operands
if FLAGS.dempster_shafer_ood:
ood_scores = ood_utils.DempsterShaferUncertainty(
logits_mean)
else:
ood_scores = 1 - tf.reduce_max(probs, axis=-1)
for metric_name, metric in metrics.items():
if name in metric_name:
metric.update_state(ood_labels, ood_scores)
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)].add_batch(
probs, label=labels)
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)
total_results = {name: metric.result() for name, metric in metrics.items()}
total_results.update(corrupt_results)
# Results from Robustness Metrics themselves return a dict, so flatten them.
total_results = utils.flatten_dictionary(total_results)
logging.info('Metrics: %s', total_results)
def main(argv):
del argv # unused arg
tf.io.gfile.makedirs(FLAGS.output_dir)
logging.info('Saving checkpoints at %s', FLAGS.output_dir)
tf.random.set_seed(FLAGS.seed)
data_dir = FLAGS.data_dir
if FLAGS.use_gpu:
logging.info('Use GPU')
strategy = tf.distribute.MirroredStrategy()
else:
logging.info('Use TPU at %s',
FLAGS.tpu if FLAGS.tpu is not None else 'local')
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
tpu=FLAGS.tpu)
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.TPUStrategy(resolver)
ds_info = tfds.builder(FLAGS.dataset).info
batch_size = (FLAGS.per_core_batch_size * FLAGS.num_cores //
FLAGS.num_dropout_samples_training)
test_batch_size = FLAGS.per_core_batch_size * FLAGS.num_cores
steps_per_epoch = ds_info.splits['train'].num_examples // batch_size
steps_per_eval = ds_info.splits['test'].num_examples // test_batch_size
num_classes = ds_info.features['label'].num_classes
train_builder = ub.datasets.get(FLAGS.dataset,
data_dir=data_dir,
download_data=FLAGS.download_data,
split=tfds.Split.TRAIN,
validation_percent=1. -
FLAGS.train_proportion)
train_dataset = train_builder.load(batch_size=batch_size)
validation_dataset = None
steps_per_validation = 0
if FLAGS.train_proportion < 1.0:
validation_builder = ub.datasets.get(
FLAGS.dataset,
data_dir=data_dir,
split=tfds.Split.VALIDATION,
validation_percent=1. - FLAGS.train_proportion,
drop_remainder=FLAGS.drop_remainder_for_eval)
validation_dataset = validation_builder.load(
batch_size=test_batch_size)
validation_dataset = strategy.experimental_distribute_dataset(
validation_dataset)
steps_per_validation = validation_builder.num_examples // test_batch_size
clean_test_builder = ub.datasets.get(
FLAGS.dataset,
data_dir=data_dir,
split=tfds.Split.TEST,
drop_remainder=FLAGS.drop_remainder_for_eval)
clean_test_dataset = clean_test_builder.load(batch_size=test_batch_size)
train_dataset = strategy.experimental_distribute_dataset(train_dataset)
test_datasets = {
'clean': strategy.experimental_distribute_dataset(clean_test_dataset),
}
steps_per_epoch = train_builder.num_examples // batch_size
steps_per_eval = clean_test_builder.num_examples // batch_size
num_classes = 100 if FLAGS.dataset == 'cifar100' else 10
if FLAGS.eval_on_ood:
ood_dataset_names = FLAGS.ood_dataset
ood_ds, steps_per_ood = ood_utils.load_ood_datasets(
ood_dataset_names,
clean_test_builder,
1. - FLAGS.train_proportion,
batch_size,
drop_remainder=FLAGS.drop_remainder_for_eval)
ood_datasets = {
name: strategy.experimental_distribute_dataset(ds)
for name, ds in ood_ds.items()
}
if FLAGS.corruptions_interval > 0:
if FLAGS.dataset == 'cifar100':
data_dir = FLAGS.cifar100_c_path
corruption_types, _ = utils.load_corrupted_test_info(FLAGS.dataset)
for corruption_type in corruption_types:
for severity in range(1, 6):
dataset = ub.datasets.get(
f'{FLAGS.dataset}_corrupted',
corruption_type=corruption_type,
data_dir=data_dir,
severity=severity,
split=tfds.Split.TEST,
drop_remainder=FLAGS.drop_remainder_for_eval).load(
batch_size=test_batch_size)
test_datasets[f'{corruption_type}_{severity}'] = (
strategy.experimental_distribute_dataset(dataset))
summary_writer = tf.summary.create_file_writer(
os.path.join(FLAGS.output_dir, 'summaries'))
with strategy.scope():
logging.info('Building ResNet model')
model = ub.models.wide_resnet_dropout(
input_shape=(32, 32, 3),
depth=28,
width_multiplier=10,
num_classes=num_classes,
l2=FLAGS.l2,
dropout_rate=FLAGS.dropout_rate,
residual_dropout=FLAGS.residual_dropout,
filterwise_dropout=FLAGS.filterwise_dropout)
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())
# Linearly scale learning rate and the decay epochs by vanilla settings.
base_lr = FLAGS.base_learning_rate * batch_size / 128
lr_decay_epochs = [(int(start_epoch_str) * FLAGS.train_epochs) // 200
for start_epoch_str in FLAGS.lr_decay_epochs]
lr_schedule = ub.schedules.WarmUpPiecewiseConstantSchedule(
steps_per_epoch,
base_lr,
decay_ratio=FLAGS.lr_decay_ratio,
decay_epochs=lr_decay_epochs,
warmup_epochs=FLAGS.lr_warmup_epochs)
optimizer = tf.keras.optimizers.SGD(lr_schedule,
momentum=1.0 -
FLAGS.one_minus_momentum,
nesterov=True)
metrics = {
'train/negative_log_likelihood':
tf.keras.metrics.Mean(),
'train/accuracy':
tf.keras.metrics.SparseCategoricalAccuracy(),
'train/loss':
tf.keras.metrics.Mean(),
'train/ece':
rm.metrics.ExpectedCalibrationError(num_bins=FLAGS.num_bins),
'test/negative_log_likelihood':
tf.keras.metrics.Mean(),
'test/accuracy':
tf.keras.metrics.SparseCategoricalAccuracy(),
'test/ece':
rm.metrics.ExpectedCalibrationError(num_bins=FLAGS.num_bins),
}
if validation_dataset:
metrics.update({
'validation/negative_log_likelihood':
tf.keras.metrics.Mean(),
'validation/accuracy':
tf.keras.metrics.SparseCategoricalAccuracy(),
'validation/ece':
rm.metrics.ExpectedCalibrationError(num_bins=FLAGS.num_bins),
})
if FLAGS.eval_on_ood:
ood_metrics = ood_utils.create_ood_metrics(ood_dataset_names)
metrics.update(ood_metrics)
if FLAGS.corruptions_interval > 0:
corrupt_metrics = {}
for intensity in range(1, 6):
for corruption in corruption_types:
dataset_name = '{0}_{1}'.format(corruption, intensity)
corrupt_metrics['test/nll_{}'.format(dataset_name)] = (
tf.keras.metrics.Mean())
corrupt_metrics['test/accuracy_{}'.format(
dataset_name)] = (
tf.keras.metrics.SparseCategoricalAccuracy())
corrupt_metrics['test/ece_{}'.format(dataset_name)] = (
rm.metrics.ExpectedCalibrationError(
num_bins=FLAGS.num_bins))
checkpoint = tf.train.Checkpoint(model=model, optimizer=optimizer)
latest_checkpoint = tf.train.latest_checkpoint(FLAGS.output_dir)
initial_epoch = 0
if latest_checkpoint:
# checkpoint.restore must be within a strategy.scope() so that optimizer
# slot variables are mirrored.
checkpoint.restore(latest_checkpoint)
logging.info('Loaded checkpoint %s', latest_checkpoint)
initial_epoch = optimizer.iterations.numpy() // steps_per_epoch
@tf.function
def train_step(iterator):
"""Training StepFn."""
def step_fn(inputs):
"""Per-Replica StepFn."""
images = inputs['features']
labels = inputs['labels']
images = tf.tile(images,
[FLAGS.num_dropout_samples_training, 1, 1, 1])
labels = tf.tile(labels, [FLAGS.num_dropout_samples_training])
with tf.GradientTape() as tape:
logits = model(images, training=True)
negative_log_likelihood = tf.reduce_mean(
tf.keras.losses.sparse_categorical_crossentropy(
labels, logits, from_logits=True))
l2_loss = sum(model.losses)
loss = negative_log_likelihood + l2_loss
# Scale the loss given the TPUStrategy will reduce sum all gradients.
scaled_loss = loss / strategy.num_replicas_in_sync
grads = tape.gradient(scaled_loss, model.trainable_variables)
optimizer.apply_gradients(zip(grads, model.trainable_variables))
probs = tf.nn.softmax(logits)
metrics['train/ece'].add_batch(probs, label=labels)
metrics['train/loss'].update_state(loss)
metrics['train/negative_log_likelihood'].update_state(
negative_log_likelihood)
metrics['train/accuracy'].update_state(labels, logits)
for _ in tf.range(tf.cast(steps_per_epoch, tf.int32)):
strategy.run(step_fn, args=(next(iterator), ))
@tf.function
def test_step(iterator, dataset_split, dataset_name, num_steps):
"""Evaluation StepFn."""
def step_fn(inputs):
"""Per-Replica StepFn."""
images = inputs['features']
labels = inputs['labels']
logits_list = []
for _ in range(FLAGS.num_dropout_samples):
logits = model(images, training=False)
logits_list.append(logits)
# Logits dimension is (num_samples, batch_size, num_classes).
logits_list = tf.stack(logits_list, axis=0)
probs_list = tf.nn.softmax(logits_list)
probs = tf.reduce_mean(probs_list, axis=0)
labels_broadcasted = tf.broadcast_to(
labels, [FLAGS.num_dropout_samples,
tf.shape(labels)[0]])
log_likelihoods = -tf.keras.losses.sparse_categorical_crossentropy(
labels_broadcasted, logits_list, from_logits=True)
negative_log_likelihood = tf.reduce_mean(
-tf.reduce_logsumexp(log_likelihoods, axis=[0]) +
tf.math.log(float(FLAGS.num_dropout_samples)))
if dataset_name == 'clean':
metrics[
f'{dataset_split}/negative_log_likelihood'].update_state(
negative_log_likelihood)
metrics[f'{dataset_split}/accuracy'].update_state(
labels, probs)
metrics[f'{dataset_split}/ece'].add_batch(probs, label=labels)
elif dataset_name.startswith('ood/'):
ood_labels = 1 - inputs['is_in_distribution']
if FLAGS.dempster_shafer_ood:
ood_scores = ood_utils.DempsterShaferUncertainty(logits)
else:
ood_scores = 1 - tf.reduce_max(probs, axis=-1)
for name, metric in metrics.items():
if dataset_name in name:
metric.update_state(ood_labels, ood_scores)
else:
corrupt_metrics['test/nll_{}'.format(
dataset_name)].update_state(negative_log_likelihood)
corrupt_metrics['test/accuracy_{}'.format(
dataset_name)].update_state(labels, probs)
corrupt_metrics['test/ece_{}'.format(dataset_name)].add_batch(
probs, label=labels)
for _ in tf.range(tf.cast(num_steps, tf.int32)):
strategy.run(step_fn, args=(next(iterator), ))
metrics.update({'test/ms_per_example': tf.keras.metrics.Mean()})
train_iterator = iter(train_dataset)
start_time = time.time()
for epoch in range(initial_epoch, FLAGS.train_epochs):
logging.info('Starting to run epoch: %s', epoch)
train_step(train_iterator)
current_step = (epoch + 1) * steps_per_epoch
max_steps = steps_per_epoch * FLAGS.train_epochs
time_elapsed = time.time() - start_time
steps_per_sec = float(current_step) / time_elapsed
eta_seconds = (max_steps - current_step) / steps_per_sec
message = ('{:.1%} completion: epoch {:d}/{:d}. {:.1f} steps/s. '
'ETA: {:.0f} min. Time elapsed: {:.0f} min'.format(
current_step / max_steps, epoch + 1, FLAGS.train_epochs,
steps_per_sec, eta_seconds / 60, time_elapsed / 60))
logging.info(message)
if validation_dataset:
validation_iterator = iter(validation_dataset)
test_step(validation_iterator, 'validation', 'clean',
steps_per_validation)
datasets_to_evaluate = {'clean': test_datasets['clean']}
if (FLAGS.corruptions_interval > 0
and (epoch + 1) % FLAGS.corruptions_interval == 0):
datasets_to_evaluate = test_datasets
for dataset_name, test_dataset in datasets_to_evaluate.items():
test_iterator = iter(test_dataset)
logging.info('Testing on dataset %s', dataset_name)
logging.info('Starting to run eval at epoch: %s', epoch)
test_start_time = time.time()
test_step(test_iterator, 'test', dataset_name, steps_per_eval)
ms_per_example = (time.time() - test_start_time) * 1e6 / batch_size
metrics['test/ms_per_example'].update_state(ms_per_example)
logging.info('Done with testing on %s', dataset_name)
if FLAGS.eval_on_ood:
for ood_dataset_name, ood_dataset in ood_datasets.items():
ood_iterator = iter(ood_dataset)
logging.info('Calculating OOD on dataset %s', ood_dataset_name)
logging.info('Running OOD eval at epoch: %s', epoch)
test_step(ood_iterator, 'test', ood_dataset_name,
steps_per_ood[ood_dataset_name])
logging.info('Done with OOD eval on %s', dataset_name)
corrupt_results = {}
if (FLAGS.corruptions_interval > 0
and (epoch + 1) % FLAGS.corruptions_interval == 0):
corrupt_results = utils.aggregate_corrupt_metrics(
corrupt_metrics, corruption_types)
logging.info('Train Loss: %.4f, Accuracy: %.2f%%',
metrics['train/loss'].result(),
metrics['train/accuracy'].result() * 100)
logging.info('Test NLL: %.4f, Accuracy: %.2f%%',
metrics['test/negative_log_likelihood'].result(),
metrics['test/accuracy'].result() * 100)
total_results = {
name: metric.result()
for name, metric in metrics.items()
}
total_results.update(corrupt_results)
# Metrics from Robustness Metrics (like ECE) will return a dict with a
# single key/value, instead of a scalar.
total_results = {
k: (list(v.values())[0] if isinstance(v, dict) else v)
for k, v in total_results.items()
}
with summary_writer.as_default():
for name, result in total_results.items():
tf.summary.scalar(name, result, step=epoch + 1)
for metric in metrics.values():
metric.reset_states()
if FLAGS.corruptions_interval > 0:
for metric in corrupt_metrics.values():
metric.reset_states()
if (FLAGS.checkpoint_interval > 0
and (epoch + 1) % FLAGS.checkpoint_interval == 0):
checkpoint_name = checkpoint.save(
os.path.join(FLAGS.output_dir, 'checkpoint'))
logging.info('Saved checkpoint to %s', checkpoint_name)
final_checkpoint_name = checkpoint.save(
os.path.join(FLAGS.output_dir, 'checkpoint'))
logging.info('Saved last checkpoint to %s', final_checkpoint_name)
with summary_writer.as_default():
hp.hparams({
'base_learning_rate':
FLAGS.base_learning_rate,
'one_minus_momentum':
FLAGS.one_minus_momentum,
'l2':
FLAGS.l2,
'dropout_rate':
FLAGS.dropout_rate,
'num_dropout_samples':
FLAGS.num_dropout_samples,
'num_dropout_samples_training':
FLAGS.num_dropout_samples_training,
})
def main(argv):
del argv # unused arg
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.random.set_seed(FLAGS.seed)
tf.io.gfile.makedirs(FLAGS.output_dir)
ds_info = tfds.builder(FLAGS.dataset).info
batch_size = FLAGS.total_batch_size
steps_per_eval = ds_info.splits['test'].num_examples // batch_size
num_classes = ds_info.features['label'].num_classes
data_dir = FLAGS.data_dir
dataset_builder = ub.datasets.get(
FLAGS.dataset,
data_dir=data_dir,
download_data=FLAGS.download_data,
split=tfds.Split.TEST,
drop_remainder=FLAGS.drop_remainder_for_eval)
dataset = dataset_builder.load(batch_size=batch_size)
test_datasets = {'clean': dataset}
if FLAGS.eval_on_ood:
ood_dataset_names = FLAGS.ood_dataset
ood_datasets, steps_per_ood = ood_utils.load_ood_datasets(
ood_dataset_names,
dataset_builder,
1. - FLAGS.train_proportion,
batch_size,
drop_remainder=FLAGS.drop_remainder_for_eval)
test_datasets.update(ood_datasets)
extra_kwargs = {}
if FLAGS.dataset == 'cifar100':
data_dir = FLAGS.cifar100_c_path
corruption_types, _ = utils.load_corrupted_test_info(FLAGS.dataset)
for corruption_type in corruption_types:
for severity in range(1, 6):
dataset = ub.datasets.get(
f'{FLAGS.dataset}_corrupted',
corruption_type=corruption_type,
data_dir=data_dir,
severity=severity,
split=tfds.Split.TEST,
drop_remainder=FLAGS.drop_remainder_for_eval,
**extra_kwargs).load(batch_size=batch_size)
test_datasets[f'{corruption_type}_{severity}'] = dataset
model = ub.models.wide_resnet_sngp(
input_shape=ds_info.features['image'].shape,
batch_size=FLAGS.total_batch_size // FLAGS.num_cores,
depth=28,
width_multiplier=10,
num_classes=num_classes,
l2=0.,
use_mc_dropout=FLAGS.use_mc_dropout,
use_filterwise_dropout=FLAGS.use_filterwise_dropout,
dropout_rate=FLAGS.dropout_rate,
use_gp_layer=FLAGS.use_gp_layer,
gp_input_dim=FLAGS.gp_input_dim,
gp_hidden_dim=FLAGS.gp_hidden_dim,
gp_scale=FLAGS.gp_scale,
gp_bias=FLAGS.gp_bias,
gp_input_normalization=FLAGS.gp_input_normalization,
gp_random_feature_type=FLAGS.gp_random_feature_type,
gp_cov_discount_factor=FLAGS.gp_cov_discount_factor,
gp_cov_ridge_penalty=FLAGS.gp_cov_ridge_penalty,
use_spec_norm=FLAGS.use_spec_norm,
spec_norm_iteration=FLAGS.spec_norm_iteration,
spec_norm_bound=FLAGS.spec_norm_bound)
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())
# Search for checkpoints from their index file; then remove the index suffix.
ensemble_filenames = tf.io.gfile.glob(
os.path.join(FLAGS.checkpoint_dir, '**/*.index'))
# Only apply ensemble on the models with the same model architecture
ensemble_filenames0 = [
filename for filename in ensemble_filenames
if f'use_gp_layer:{FLAGS.use_gp_layer}' in filename
and f'use_spec_norm:{FLAGS.use_spec_norm}' in filename
]
np.random.seed(FLAGS.seed)
ensemble_filenames = np.random.choice(ensemble_filenames0,
FLAGS.ensemble_size,
replace=True)
ensemble_filenames = [filename[:-6] for filename in ensemble_filenames]
ensemble_size = len(ensemble_filenames)
logging.info('Ensemble size: %s', ensemble_size)
logging.info('Ensemble filenames: %s', ensemble_filenames)
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)
# Write model predictions to files.
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.replace('/', '_'),
member=m) # ood dataset name has '/'
filename = os.path.join(FLAGS.output_dir, filename)
if not tf.io.gfile.exists(filename):
logits = []
test_iterator = iter(test_dataset)
steps = steps_per_eval if 'ood/' not in name else steps_per_ood[
name]
for _ in range(steps):
features = next(test_iterator)['features'] # pytype: disable=unsupported-operands
logits_member = model(features, training=False)
if isinstance(logits_member, (list, tuple)):
# If model returns a tuple of (logits, covmat), extract both
logits_member, covmat_member = logits_member
logits_member = ed.layers.utils.mean_field_logits(
logits_member, covmat_member,
FLAGS.gp_mean_field_factor_ensemble)
logits.append(logits_member)
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':
rm.metrics.ExpectedCalibrationError(num_bins=FLAGS.num_bins),
}
if FLAGS.eval_on_ood:
ood_metrics = ood_utils.create_ood_metrics(ood_dataset_names)
metrics.update(ood_metrics)
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)] = (
rm.metrics.ExpectedCalibrationError(num_bins=FLAGS.num_bins))
# Evaluate model predictions.
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.replace('/', '_'),
member=m) # ood dataset name has '/'
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)
steps = steps_per_eval if 'ood/' not in name else steps_per_ood[name]
for step in range(steps):
inputs = next(test_iterator)
labels = inputs['labels'] # pytype: disable=unsupported-operands
logits = logits_dataset[:, (step * batch_size):((step + 1) *
batch_size)]
labels = tf.cast(labels, tf.int32)
negative_log_likelihood_metric = rm.metrics.EnsembleCrossEntropy()
negative_log_likelihood_metric.add_batch(logits, labels=labels)
negative_log_likelihood = list(
negative_log_likelihood_metric.result().values())[0]
per_probs = tf.nn.softmax(logits)
probs = tf.reduce_mean(per_probs, axis=0)
logits_mean = tf.reduce_mean(logits, axis=0)
if name == 'clean':
gibbs_ce_metric = rm.metrics.GibbsCrossEntropy()
gibbs_ce_metric.add_batch(logits, labels=labels)
gibbs_ce = list(gibbs_ce_metric.result().values())[0]
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'].add_batch(probs, label=labels)
elif name.startswith('ood/'):
ood_labels = 1 - inputs['is_in_distribution'] # pytype: disable=unsupported-operands
if FLAGS.dempster_shafer_ood:
ood_scores = ood_utils.DempsterShaferUncertainty(
logits_mean)
else:
ood_scores = 1 - tf.reduce_max(probs, axis=-1)
for metric_name, metric in metrics.items():
if name in metric_name:
metric.update_state(ood_labels, ood_scores)
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)].add_batch(
probs, label=labels)
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)
total_results = {name: metric.result() for name, metric in metrics.items()}
total_results.update(corrupt_results)
# Metrics from Robustness Metrics (like ECE) will return a dict with a
# single key/value, instead of a scalar.
total_results = {
k: (list(v.values())[0] if isinstance(v, dict) else v)
for k, v in total_results.items()
}
logging.info('Metrics: %s', total_results)