def _basic_metrics(self): metrics_basic = {} metrics_basic['logits'] = [] metrics_basic['probs'] = [tf.keras.metrics.SparseCategoricalAccuracy(name='acc'), um.ExpectedCalibrationError(num_bins=self.FLAGS.ece_bins,name='ece'), nll(name='nll'), BrierScore(name='brier')] metrics_basic['certs'] = [tf.keras.metrics.SparseCategoricalAccuracy(name='acc'), um.ExpectedCalibrationError(num_bins=self.FLAGS.ece_bins,name='ece'), nll(name='nll'), BrierScore(name='brier')] metrics_basic['logits_from_certs'] = [] return metrics_basic
def get_diabetic_retinopathy_base_metrics(use_tpu, num_bins): """Initialize base metrics for non-ensemble Diabetic Retinopathy predictors. Should be called within the distribution strategy scope (e.g. see deterministic.py script). Note: We disclude AUC in non-TPU case, which must be defined and added to this dict outside the strategy scope. We disclude ECE in TPU case, which currently throws an XLA error on TPU. Args: use_tpu: bool, is run using TPU. num_bins: number of ECE bins. Returns: dict, metrics """ metrics = { 'train/negative_log_likelihood': tf.keras.metrics.Mean(), 'train/accuracy': tf.keras.metrics.BinaryAccuracy(), 'train/loss': tf.keras.metrics.Mean(), # NLL + L2 'validation/negative_log_likelihood': tf.keras.metrics.Mean(), 'validation/accuracy': tf.keras.metrics.BinaryAccuracy(), 'test/negative_log_likelihood': tf.keras.metrics.Mean(), 'test/accuracy': tf.keras.metrics.BinaryAccuracy(), } if use_tpu: # AUC does not yet work within GPU strategy scope, but does for TPU metrics.update({ 'train/auc': tf.keras.metrics.AUC(), 'validation/auc': tf.keras.metrics.AUC(), 'test/auc': tf.keras.metrics.AUC(), }) else: # ECE does not yet work on TPU metrics.update({ 'train/ece': um.ExpectedCalibrationError(num_bins=num_bins), 'validation/ece': um.ExpectedCalibrationError(num_bins=num_bins), 'test/ece': um.ExpectedCalibrationError(num_bins=num_bins), }) return metrics
def testECEBinaryClassification(self): num_bins = 10 pred_probs = np.array([0.51, 0.45, 0.39, 0.66, 0.68, 0.29, 0.81, 0.85]) # max_pred_probs: [0.51, 0.55, 0.61, 0.66, 0.68, 0.71, 0.81, 0.85] # pred_class: [1, 0, 0, 1, 1, 0, 1, 1] labels = np.array([0., 0., 0., 1., 0., 1., 1., 1.]) n = len(pred_probs) # Bins for the max predicted probabilities are (0, 0.1), [0.1, 0.2), ..., # [0.9, 1) and are numbered starting at zero. bin_counts = np.array([0, 0, 0, 0, 0, 2, 3, 1, 2, 0]) bin_correct_sums = np.array([0, 0, 0, 0, 0, 1, 2, 0, 2, 0]) bin_prob_sums = np.array([0, 0, 0, 0, 0, 0.51 + 0.55, 0.61 + 0.66 + 0.68, 0.71, 0.81 + 0.85, 0]) correct_ece = 0. bin_accs = np.array([0.] * num_bins) bin_confs = np.array([0.] * num_bins) for i in range(num_bins): if bin_counts[i] > 0: bin_accs[i] = bin_correct_sums[i] / bin_counts[i] bin_confs[i] = bin_prob_sums[i] / bin_counts[i] correct_ece += bin_counts[i] / n * abs(bin_accs[i] - bin_confs[i]) metric = um.ExpectedCalibrationError( num_bins, name='ECE', dtype=tf.float64) self.assertLen(metric.variables, 3) ece1 = metric(labels, pred_probs) self.assertAllClose(ece1, correct_ece) actual_bin_counts = tf.convert_to_tensor(metric.counts) actual_bin_correct_sums = tf.convert_to_tensor(metric.correct_sums) actual_bin_prob_sums = tf.convert_to_tensor(metric.prob_sums) self.assertAllEqual(bin_counts, actual_bin_counts) self.assertAllEqual(bin_correct_sums, actual_bin_correct_sums) self.assertAllClose(bin_prob_sums, actual_bin_prob_sums) # Test various types of input shapes. metric.reset_states() metric.update_state(labels[:2], pred_probs[:2]) metric.update_state(labels[2:6].reshape(2, 2), pred_probs[2:6].reshape(2, 2)) metric.update_state(labels[6:7], pred_probs[6:7]) ece2 = metric(labels[7:, np.newaxis], pred_probs[7:, np.newaxis]) ece3 = metric.result() self.assertAllClose(ece2, ece3) self.assertAllClose(ece3, correct_ece) actual_bin_counts = tf.convert_to_tensor(metric.counts) actual_bin_correct_sums = tf.convert_to_tensor(metric.correct_sums) actual_bin_prob_sums = tf.convert_to_tensor(metric.prob_sums) self.assertAllEqual(bin_counts, actual_bin_counts) self.assertAllEqual(bin_correct_sums, actual_bin_correct_sums) self.assertAllClose(bin_prob_sums, actual_bin_prob_sums)
def testECEMulticlassClassification(self): num_bins = 10 pred_probs = [ [0.31, 0.32, 0.27], [0.37, 0.33, 0.30], [0.30, 0.31, 0.39], [0.61, 0.38, 0.01], [0.10, 0.65, 0.25], [0.91, 0.05, 0.04], ] # max_pred_probs: [0.32, 0.37, 0.39, 0.61, 0.65, 0.91] # pred_class: [1, 0, 2, 0, 1, 0] labels = [1., 0, 0., 1., 0., 0.] n = len(pred_probs) # Bins for the max predicted probabilities are (0, 0.1), [0.1, 0.2), ..., # [0.9, 1) and are numbered starting at zero. bin_counts = [0, 0, 0, 3, 0, 0, 2, 0, 0, 1] bin_correct_sums = [0, 0, 0, 2, 0, 0, 0, 0, 0, 1] bin_prob_sums = [0, 0, 0, 0.32 + 0.37 + 0.39, 0, 0, 0.61 + 0.65, 0, 0, 0.91] correct_ece = 0. bin_accs = [0.] * num_bins bin_confs = [0.] * num_bins for i in range(num_bins): if bin_counts[i] > 0: bin_accs[i] = bin_correct_sums[i] / bin_counts[i] bin_confs[i] = bin_prob_sums[i] / bin_counts[i] correct_ece += bin_counts[i] / n * abs(bin_accs[i] - bin_confs[i]) metric = um.ExpectedCalibrationError( num_bins, name='ECE', dtype=tf.float64) self.assertLen(metric.variables, 3) metric.update_state(labels[:4], pred_probs[:4]) ece1 = metric(labels[4:], pred_probs[4:]) ece2 = metric.result() self.assertAllClose(ece1, ece2) self.assertAllClose(ece2, correct_ece) actual_bin_counts = tf.convert_to_tensor(metric.counts) actual_bin_correct_sums = tf.convert_to_tensor(metric.correct_sums) actual_bin_prob_sums = tf.convert_to_tensor(metric.prob_sums) self.assertAllEqual(bin_counts, actual_bin_counts) self.assertAllEqual(bin_correct_sums, actual_bin_correct_sums) self.assertAllClose(bin_prob_sums, actual_bin_prob_sums)
def testECEBinaryClassificationKerasModel(self): num_bins = 10 pred_probs = np.array([0.51, 0.45, 0.39, 0.66, 0.68, 0.29, 0.81, 0.85]) # max_pred_probs: [0.51, 0.55, 0.61, 0.66, 0.68, 0.71, 0.81, 0.85] # pred_class: [1, 0, 0, 1, 1, 0, 1, 1] labels = np.array([0., 0., 0., 1., 0., 1., 1., 1.]) n = len(pred_probs) # Bins for the max predicted probabilities are (0, 0.1), [0.1, 0.2), ..., # [0.9, 1) and are numbered starting at zero. bin_counts = [0, 0, 0, 0, 0, 2, 3, 1, 2, 0] bin_correct_sums = [0, 0, 0, 0, 0, 1, 2, 0, 2, 0] bin_prob_sums = [0, 0, 0, 0, 0, 0.51 + 0.55, 0.61 + 0.66 + 0.68, 0.71, 0.81 + 0.85, 0] correct_ece = 0. bin_accs = [0.] * num_bins bin_confs = [0.] * num_bins for i in range(num_bins): if bin_counts[i] > 0: bin_accs[i] = bin_correct_sums[i] / bin_counts[i] bin_confs[i] = bin_prob_sums[i] / bin_counts[i] correct_ece += bin_counts[i] / n * abs(bin_accs[i] - bin_confs[i]) metric = um.ExpectedCalibrationError(num_bins, name='ECE') self.assertLen(metric.variables, 3) model = tf.keras.models.Sequential([tf.keras.layers.Lambda(lambda x: 1*x)]) model.compile(loss='binary_crossentropy', optimizer='sgd', metrics=[metric]) outputs = model.predict(pred_probs) self.assertAllClose(pred_probs.reshape([n, 1]), outputs) _, ece = model.evaluate(pred_probs, labels) self.assertAllClose(ece, correct_ece) actual_bin_counts = tf.convert_to_tensor(metric.counts) actual_bin_correct_sums = tf.convert_to_tensor(metric.correct_sums) actual_bin_prob_sums = tf.convert_to_tensor(metric.prob_sums) self.assertAllEqual(bin_counts, actual_bin_counts) self.assertAllEqual(bin_correct_sums, actual_bin_correct_sums) self.assertAllClose(bin_prob_sums, actual_bin_prob_sums)
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) batch_size = FLAGS.per_core_batch_size * FLAGS.num_cores steps_per_eval = IMAGENET_VALIDATION_IMAGES // batch_size builder = utils.ImageNetInput(data_dir=FLAGS.data_dir, use_bfloat16=False) clean_test_dataset = builder.as_dataset(split=tfds.Split.TEST, batch_size=batch_size) test_datasets = {'clean': clean_test_dataset} 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( corruption_name=name, corruption_intensity=intensity, batch_size=batch_size, drop_remainder=True, use_bfloat16=False) model = ub.models.resnet50_deterministic(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()) # 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')) 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) # 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, 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) # pytype: disable=attribute-error 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': um.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)] = um.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, 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) # pytype: disable=attribute-error logits = logits_dataset[:, (step*batch_size):((step+1)*batch_size)] labels = tf.cast(tf.reshape(labels, [-1]), tf.int32) negative_log_likelihood = um.ensemble_cross_entropy(labels, logits) per_probs = tf.nn.softmax(logits) probs = tf.reduce_mean(per_probs, axis=0) if name == 'clean': gibbs_ce = um.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 # 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) logging.info('Model checkpoint will be saved at %s', FLAGS.output_dir) tf.io.gfile.makedirs(FLAGS.output_dir) batch_size = FLAGS.per_core_batch_size * FLAGS.num_cores test_batch_size = batch_size data_buffer_size = batch_size * 10 ind_dataset_builder = ds.WikipediaToxicityDataset( split='test', data_dir=FLAGS.in_dataset_dir, shuffle_buffer_size=data_buffer_size) ood_dataset_builder = ds.CivilCommentsDataset( split='test', data_dir=FLAGS.ood_dataset_dir, shuffle_buffer_size=data_buffer_size) ood_identity_dataset_builder = ds.CivilCommentsIdentitiesDataset( split='test', data_dir=FLAGS.identity_dataset_dir, shuffle_buffer_size=data_buffer_size) test_dataset_builders = { 'ind': ind_dataset_builder, 'ood': ood_dataset_builder, 'ood_identity': ood_identity_dataset_builder, } class_weight = utils.create_class_weight( test_dataset_builders=test_dataset_builders) logging.info('class_weight: %s', str(class_weight)) ds_info = ind_dataset_builder.tfds_info # Positive and negative classes. num_classes = ds_info.metadata['num_classes'] test_datasets = {} steps_per_eval = {} for dataset_name, dataset_builder in test_dataset_builders.items(): test_datasets[dataset_name] = dataset_builder.load( batch_size=test_batch_size) steps_per_eval[dataset_name] = ( dataset_builder.num_examples // test_batch_size) logging.info('Building %s model', FLAGS.model_family) bert_config_dir, _ = utils.resolve_bert_ckpt_and_config_dir( FLAGS.bert_model_type, FLAGS.bert_dir, FLAGS.bert_config_dir, FLAGS.bert_ckpt_dir) bert_config = utils.create_config(bert_config_dir) gp_layer_kwargs = dict( num_inducing=FLAGS.gp_hidden_dim, gp_kernel_scale=FLAGS.gp_scale, gp_output_bias=FLAGS.gp_bias, normalize_input=FLAGS.gp_input_normalization, gp_cov_momentum=FLAGS.gp_cov_discount_factor, gp_cov_ridge_penalty=FLAGS.gp_cov_ridge_penalty) spec_norm_kwargs = dict( iteration=FLAGS.spec_norm_iteration, norm_multiplier=FLAGS.spec_norm_bound) model, _ = ub.models.SngpBertBuilder( num_classes=num_classes, bert_config=bert_config, gp_layer_kwargs=gp_layer_kwargs, spec_norm_kwargs=spec_norm_kwargs, use_gp_layer=FLAGS.use_gp_layer, use_spec_norm_att=FLAGS.use_spec_norm_att, use_spec_norm_ffn=FLAGS.use_spec_norm_ffn, use_layer_norm_att=FLAGS.use_layer_norm_att, use_layer_norm_ffn=FLAGS.use_layer_norm_ffn, use_spec_norm_plr=FLAGS.use_spec_norm_plr) 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')) ensemble_filenames = [filename[:-6] for filename in ensemble_filenames] if FLAGS.num_models > len(ensemble_filenames): raise ValueError('Number of models to be included in the ensemble ' 'should be less than total number of models in ' 'the checkpoint_dir.') ensemble_filenames = ensemble_filenames[:FLAGS.num_models] 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).assert_existing_objects_matched() for n, (dataset_name, test_dataset) in enumerate(test_datasets.items()): filename = '{dataset}_{member}.npy'.format(dataset=dataset_name, member=m) filename = os.path.join(FLAGS.output_dir, filename) if not tf.io.gfile.exists(filename): logits_list = [] test_iterator = iter(test_dataset) for step in range(steps_per_eval[dataset_name]): try: inputs = next(test_iterator) except StopIteration: continue features, labels, _ = utils.create_feature_and_label(inputs) logits = model(features, training=False) if isinstance(logits, (list, tuple)): # If model returns a tuple of (logits, covmat), extract both. logits, covmat = logits else: covmat = tf.eye(test_batch_size) if FLAGS.use_bfloat16: logits = tf.cast(logits, tf.float32) covmat = tf.cast(covmat, tf.float32) logits = ed.layers.utils.mean_field_logits( logits, covmat, mean_field_factor=FLAGS.gp_mean_field_factor_ensemble) logits_list.append(logits) logits_all = tf.concat(logits_list, axis=0) with tf.io.gfile.GFile(filename, 'w') as f: np.save(f, logits_all.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/auroc': tf.keras.metrics.AUC(curve='ROC'), 'test/aupr': tf.keras.metrics.AUC(curve='PR'), 'test/brier': tf.keras.metrics.MeanSquaredError(), 'test/brier_weighted': tf.keras.metrics.MeanSquaredError(), 'test/ece': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), 'test/acc': tf.keras.metrics.Accuracy(), 'test/acc_weighted': tf.keras.metrics.Accuracy(), 'test/precision': tf.keras.metrics.Precision(), 'test/recall': tf.keras.metrics.Recall(), 'test/f1': tfa_metrics.F1Score( num_classes=num_classes, average='micro', threshold=FLAGS.ece_label_threshold) } for fraction in FLAGS.fractions: metrics.update({ 'test_collab_acc/collab_acc_{}'.format(fraction): um.OracleCollaborativeAccuracy( fraction=float(fraction), num_bins=FLAGS.num_bins) }) for dataset_name, test_dataset in test_datasets.items(): if dataset_name != 'ind': metrics.update({ 'test/nll_{}'.format(dataset_name): tf.keras.metrics.Mean(), 'test/auroc_{}'.format(dataset_name): tf.keras.metrics.AUC(curve='ROC'), 'test/aupr_{}'.format(dataset_name): tf.keras.metrics.AUC(curve='PR'), 'test/brier_{}'.format(dataset_name): tf.keras.metrics.MeanSquaredError(), 'test/brier_weighted_{}'.format(dataset_name): tf.keras.metrics.MeanSquaredError(), 'test/ece_{}'.format(dataset_name): um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), 'test/acc_weighted_{}'.format(dataset_name): tf.keras.metrics.Accuracy(), 'test/acc_{}'.format(dataset_name): tf.keras.metrics.Accuracy(), 'test/precision_{}'.format(dataset_name): tf.keras.metrics.Precision(), 'test/recall_{}'.format(dataset_name): tf.keras.metrics.Recall(), 'test/f1_{}'.format(dataset_name): tfa_metrics.F1Score( num_classes=num_classes, average='micro', threshold=FLAGS.ece_label_threshold) }) for fraction in FLAGS.fractions: metrics.update({ 'test_collab_acc/collab_acc_{}_{}'.format(fraction, dataset_name): um.OracleCollaborativeAccuracy( fraction=float(fraction), num_bins=FLAGS.num_bins) }) @tf.function def generate_sample_weight(labels, class_weight, label_threshold=0.7): """Generate sample weight for weighted accuracy calculation.""" if label_threshold != 0.7: logging.warning('The class weight was based on `label_threshold` = 0.7, ' 'and weighted accuracy/brier will be meaningless if ' '`label_threshold` is not equal to this value, which is ' 'recommended by Jigsaw Conversation AI team.') labels_int = tf.cast(labels > label_threshold, tf.int32) sample_weight = tf.gather(class_weight, labels_int) return sample_weight # Evaluate model predictions. for n, (dataset_name, test_dataset) in enumerate(test_datasets.items()): logits_dataset = [] for m in range(ensemble_size): filename = '{dataset}_{member}.npy'.format(dataset=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) texts_list = [] logits_list = [] labels_list = [] # Use dict to collect additional labels specified by additional label names. # Here we use `OrderedDict` to get consistent ordering for this dict so # we can retrieve the predictions for each identity labels in Colab. additional_labels_dict = collections.OrderedDict() for step in range(steps_per_eval[dataset_name]): try: inputs = next(test_iterator) # type: Mapping[Text, tf.Tensor] # pytype: disable=annotation-type-mismatch except StopIteration: continue features, labels, additional_labels = ( utils.create_feature_and_label(inputs)) logits = logits_dataset[:, (step * batch_size):((step + 1) * batch_size)] loss_logits = tf.squeeze(logits, axis=-1) negative_log_likelihood = um.ensemble_cross_entropy( labels, loss_logits, binary=True) per_probs = tf.nn.sigmoid(logits) probs = tf.reduce_mean(per_probs, axis=0) # Cast labels to discrete for ECE computation ece_labels = tf.cast(labels > FLAGS.ece_label_threshold, tf.float32) one_hot_labels = tf.one_hot(tf.cast(ece_labels, tf.int32), depth=num_classes) ece_probs = tf.concat([1. - probs, probs], axis=1) pred_labels = tf.math.argmax(ece_probs, axis=-1) auc_probs = tf.squeeze(probs, axis=1) texts_list.append(inputs['input_ids']) logits_list.append(logits) labels_list.append(labels) if 'identity' in dataset_name: for identity_label_name in utils.IDENTITY_LABELS: if identity_label_name not in additional_labels_dict: additional_labels_dict[identity_label_name] = [] additional_labels_dict[identity_label_name].append( additional_labels[identity_label_name].numpy()) sample_weight = generate_sample_weight( labels, class_weight['test/{}'.format(dataset_name)], FLAGS.ece_label_threshold) if dataset_name == 'ind': metrics['test/negative_log_likelihood'].update_state( negative_log_likelihood) metrics['test/auroc'].update_state(labels, auc_probs) metrics['test/aupr'].update_state(labels, auc_probs) metrics['test/brier'].update_state(labels, auc_probs) metrics['test/brier_weighted'].update_state( tf.expand_dims(labels, -1), probs, sample_weight=sample_weight) metrics['test/ece'].add_batch(ece_probs, label=ece_labels) metrics['test/acc'].update_state(ece_labels, pred_labels) metrics['test/acc_weighted'].update_state( ece_labels, pred_labels, sample_weight=sample_weight) metrics['test/precision'].update_state(ece_labels, pred_labels) metrics['test/recall'].update_state(ece_labels, pred_labels) metrics['test/f1'].update_state(one_hot_labels, ece_probs) for fraction in FLAGS.fractions: metrics['test_collab_acc/collab_acc_{}'.format( fraction)].update_state(ece_labels, ece_probs) else: metrics['test/nll_{}'.format(dataset_name)].update_state( negative_log_likelihood) metrics['test/auroc_{}'.format(dataset_name)].update_state( labels, auc_probs) metrics['test/aupr_{}'.format(dataset_name)].update_state( labels, auc_probs) metrics['test/brier_{}'.format(dataset_name)].update_state( labels, auc_probs) metrics['test/brier_weighted_{}'.format(dataset_name)].update_state( tf.expand_dims(labels, -1), probs, sample_weight=sample_weight) metrics['test/ece_{}'.format(dataset_name)].add_batch( ece_probs, label=ece_labels) metrics['test/acc_{}'.format(dataset_name)].update_state( ece_labels, pred_labels) metrics['test/acc_weighted_{}'.format(dataset_name)].update_state( ece_labels, pred_labels, sample_weight=sample_weight) metrics['test/precision_{}'.format(dataset_name)].update_state( ece_labels, pred_labels) metrics['test/recall_{}'.format(dataset_name)].update_state( ece_labels, pred_labels) metrics['test/f1_{}'.format(dataset_name)].update_state( one_hot_labels, ece_probs) for fraction in FLAGS.fractions: metrics['test_collab_acc/collab_acc_{}_{}'.format( fraction, dataset_name)].update_state(ece_labels, ece_probs) texts_all = tf.concat(texts_list, axis=0) logits_all = tf.concat(logits_list, axis=1) labels_all = tf.concat(labels_list, axis=0) additional_labels_all = [] if additional_labels_dict: additional_labels_all = list(additional_labels_dict.values()) utils.save_prediction( texts_all.numpy(), path=os.path.join(FLAGS.output_dir, 'texts_{}'.format(dataset_name))) utils.save_prediction( labels_all.numpy(), path=os.path.join(FLAGS.output_dir, 'labels_{}'.format(dataset_name))) utils.save_prediction( logits_all.numpy(), path=os.path.join(FLAGS.output_dir, 'logits_{}'.format(dataset_name))) if 'identity' in dataset_name: utils.save_prediction( np.array(additional_labels_all), path=os.path.join(FLAGS.output_dir, 'additional_labels_{}'.format(dataset_name))) message = ('{:.1%} completion for evaluation: dataset {:d}/{:d}'.format( (n + 1) / num_datasets, n + 1, num_datasets)) logging.info(message) total_results = {name: metric.result() for name, metric in metrics.items()} # 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)
def main(argv): del argv # Unused arg. tf.random.set_seed(FLAGS.seed) if FLAGS.version2: per_core_bs_train = FLAGS.per_core_batch_size // (FLAGS.ensemble_size * FLAGS.num_train_samples) per_core_bs_eval = FLAGS.per_core_batch_size // (FLAGS.ensemble_size * FLAGS.num_eval_samples) else: per_core_bs_train = FLAGS.per_core_batch_size // FLAGS.num_train_samples per_core_bs_eval = FLAGS.per_core_batch_size // FLAGS.num_eval_samples batch_size_train = per_core_bs_train * FLAGS.num_cores batch_size_eval = per_core_bs_eval * FLAGS.num_cores logging.info('Saving checkpoints at %s', FLAGS.output_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) train_dataset = utils.load_dataset( split=tfds.Split.TRAIN, name=FLAGS.dataset, batch_size=batch_size_train, use_bfloat16=FLAGS.use_bfloat16, normalize=False) clean_test_dataset = utils.load_dataset( split=tfds.Split.TEST, name=FLAGS.dataset, batch_size=batch_size_eval, use_bfloat16=FLAGS.use_bfloat16, normalize=False) train_dataset = strategy.experimental_distribute_dataset(train_dataset) test_datasets = { 'clean': strategy.experimental_distribute_dataset(clean_test_dataset), } if FLAGS.corruptions_interval > 0: if FLAGS.dataset == 'cifar10': load_c_dataset = utils.load_cifar10_c else: load_c_dataset = functools.partial(utils.load_cifar100_c, path=FLAGS.cifar100_c_path) corruption_types, max_intensity = utils.load_corrupted_test_info( FLAGS.dataset) for corruption in corruption_types: for intensity in range(1, max_intensity + 1): dataset = load_c_dataset( corruption_name=corruption, corruption_intensity=intensity, batch_size=batch_size_eval, use_bfloat16=FLAGS.use_bfloat16, normalize=False) test_datasets['{0}_{1}'.format(corruption, intensity)] = ( strategy.experimental_distribute_dataset(dataset)) ds_info = tfds.builder(FLAGS.dataset).info train_dataset_size = ds_info.splits['train'].num_examples test_dataset_size = ds_info.splits['test'].num_examples num_classes = ds_info.features['label'].num_classes steps_per_epoch = train_dataset_size // batch_size_train steps_per_eval = test_dataset_size // batch_size_eval if FLAGS.use_bfloat16: policy = tf.keras.mixed_precision.experimental.Policy('mixed_bfloat16') tf.keras.mixed_precision.experimental.set_policy(policy) summary_writer = tf.summary.create_file_writer( os.path.join(FLAGS.output_dir, 'summaries')) with strategy.scope(): logging.info('Building Keras ResNet-32 model') model = resnet_cifar_model.rank1_resnet_v1( input_shape=ds_info.features['image'].shape, depth=32, num_classes=num_classes, width_multiplier=4, alpha_initializer=FLAGS.alpha_initializer, gamma_initializer=FLAGS.gamma_initializer, alpha_regularizer=FLAGS.alpha_regularizer, gamma_regularizer=FLAGS.gamma_regularizer, use_additive_perturbation=FLAGS.use_additive_perturbation, ensemble_size=FLAGS.ensemble_size, random_sign_init=FLAGS.random_sign_init, dropout_rate=FLAGS.dropout_rate) logging.info(model.summary()) base_lr = FLAGS.base_learning_rate * batch_size_train / 128 lr_decay_epochs = [(int(start_epoch_str) * FLAGS.train_epochs) // 200 for start_epoch_str in FLAGS.lr_decay_epochs] lr_schedule = utils.LearningRateSchedule( 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=0.9, 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': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), 'test/negative_log_likelihood': tf.keras.metrics.Mean(), 'test/accuracy': tf.keras.metrics.SparseCategoricalAccuracy(), 'test/ece': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), 'test/loss': tf.keras.metrics.Mean(), } if FLAGS.corruptions_interval > 0: corrupt_metrics = {} for intensity in range(1, max_intensity + 1): 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)] = ( um.ExpectedCalibrationError(num_bins=FLAGS.num_bins)) test_diversity = {} training_diversity = {} if FLAGS.ensemble_size > 1: for i in range(FLAGS.ensemble_size): metrics['test/nll_member_{}'.format(i)] = tf.keras.metrics.Mean() metrics['test/accuracy_member_{}'.format(i)] = ( tf.keras.metrics.SparseCategoricalAccuracy()) test_diversity = { 'test/disagreement': tf.keras.metrics.Mean(), 'test/average_kl': tf.keras.metrics.Mean(), 'test/cosine_similarity': tf.keras.metrics.Mean(), } training_diversity = { 'train/disagreement': tf.keras.metrics.Mean(), 'train/average_kl': tf.keras.metrics.Mean(), 'train/cosine_similarity': tf.keras.metrics.Mean(), } 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, labels = inputs if FLAGS.version2 and FLAGS.ensemble_size > 1: images = tf.tile(images, [FLAGS.ensemble_size, 1, 1, 1]) if not (FLAGS.member_sampling or FLAGS.expected_probs): labels = tf.tile(labels, [FLAGS.ensemble_size]) if FLAGS.num_train_samples > 1: images = tf.tile(images, [FLAGS.num_train_samples, 1, 1, 1]) with tf.GradientTape() as tape: logits = model(images, training=True) probs = tf.nn.softmax(logits) # Diversity evaluation. if FLAGS.version2 and FLAGS.ensemble_size > 1: per_probs = tf.reshape( probs, tf.concat([[FLAGS.ensemble_size, -1], probs.shape[1:]], 0)) diversity_results = um.average_pairwise_diversity( per_probs, FLAGS.ensemble_size) if FLAGS.num_train_samples > 1: probs = tf.reshape(probs, tf.concat([[FLAGS.num_train_samples, -1], probs.shape[1:]], 0)) probs = tf.reduce_mean(probs, 0) if FLAGS.member_sampling and FLAGS.version2 and FLAGS.ensemble_size > 1: idx = tf.random.uniform([], maxval=FLAGS.ensemble_size, dtype=tf.int64) idx_one_hot = tf.expand_dims(tf.one_hot(idx, FLAGS.ensemble_size, dtype=probs.dtype), 0) probs_shape = probs.shape probs = tf.reshape(probs, [FLAGS.ensemble_size, -1]) probs = tf.matmul(idx_one_hot, probs) probs = tf.reshape(probs, tf.concat([[-1], probs_shape[1:]], 0)) elif FLAGS.expected_probs and FLAGS.version2 and FLAGS.ensemble_size > 1: probs = tf.reshape(probs, tf.concat([[FLAGS.ensemble_size, -1], probs.shape[1:]], 0)) probs = tf.reduce_mean(probs, 0) negative_log_likelihood = tf.reduce_mean( tf.keras.losses.sparse_categorical_crossentropy(labels, probs)) filtered_variables = [] for var in model.trainable_variables: # Apply l2 on the slow weights and bias terms. This excludes BN # parameters and fast weight approximate posterior/prior parameters, # but pay caution to their naming scheme. if 'kernel' in var.name or 'bias' in var.name: filtered_variables.append(tf.reshape(var, (-1,))) l2_loss = FLAGS.l2 * 2 * tf.nn.l2_loss( tf.concat(filtered_variables, axis=0)) kl = sum(model.losses) / train_dataset_size kl_scale = tf.cast(optimizer.iterations + 1, kl.dtype) kl_scale /= FLAGS.kl_annealing_steps kl_scale = tf.minimum(1., kl_scale) kl_loss = kl_scale * kl # Scale the loss given the TPUStrategy will reduce sum all gradients. loss = negative_log_likelihood + l2_loss + kl_loss scaled_loss = loss / strategy.num_replicas_in_sync grads = tape.gradient(scaled_loss, model.trainable_variables) # Separate learning rate implementation. grad_list = [] if FLAGS.fast_weight_lr_multiplier != 1.0: grads_and_vars = list(zip(grads, model.trainable_variables)) for vec, var in grads_and_vars: # Apply different learning rate on the fast weight approximate # posterior/prior parameters. This is excludes BN and slow weights, # but pay caution to the naming scheme. if ('batch_norm' not in var.name and 'kernel' not in var.name): grad_list.append((vec * FLAGS.fast_weight_lr_multiplier, var)) else: grad_list.append((vec, var)) optimizer.apply_gradients(grad_list) else: optimizer.apply_gradients(zip(grads, model.trainable_variables)) metrics['train/ece'].update_state(labels, probs) metrics['train/loss'].update_state(loss) metrics['train/negative_log_likelihood'].update_state( negative_log_likelihood) metrics['train/accuracy'].update_state(labels, probs) if FLAGS.version2 and FLAGS.ensemble_size > 1: for k, v in diversity_results.items(): training_diversity['train/' + k].update_state(v) strategy.run(step_fn, args=(next(iterator),)) @tf.function def test_step(iterator, dataset_name): """Evaluation StepFn.""" def step_fn(inputs): """Per-Replica StepFn.""" images, labels = inputs if FLAGS.ensemble_size > 1: images = tf.tile(images, [FLAGS.ensemble_size, 1, 1, 1]) if FLAGS.num_eval_samples > 1: images = tf.tile(images, [FLAGS.num_eval_samples, 1, 1, 1]) logits = model(images, training=False) probs = tf.nn.softmax(logits) if FLAGS.num_eval_samples > 1: probs = tf.reshape(probs, tf.concat([[FLAGS.num_eval_samples, -1], probs.shape[1:]], 0)) probs = tf.reduce_mean(probs, 0) if FLAGS.ensemble_size > 1: per_probs = tf.split(probs, num_or_size_splits=FLAGS.ensemble_size, axis=0) if dataset_name == 'clean': per_probs_tensor = tf.reshape( probs, tf.concat([[FLAGS.ensemble_size, -1], probs.shape[1:]], 0)) diversity_results = um.average_pairwise_diversity( per_probs_tensor, FLAGS.ensemble_size) for k, v in diversity_results.items(): test_diversity['test/' + k].update_state(v) for i in range(FLAGS.ensemble_size): member_probs = per_probs[i] member_nll = tf.keras.losses.sparse_categorical_crossentropy( labels, member_probs) metrics['test/nll_member_{}'.format(i)].update_state(member_nll) metrics['test/accuracy_member_{}'.format(i)].update_state( labels, member_probs) probs = tf.reduce_mean(per_probs, axis=0) negative_log_likelihood = tf.reduce_mean( tf.keras.losses.sparse_categorical_crossentropy(labels, probs)) filtered_variables = [] for var in model.trainable_variables: if 'kernel' in var.name or 'bias' in var.name: filtered_variables.append(tf.reshape(var, (-1,))) kl = sum(model.losses) / test_dataset_size l2_loss = kl + FLAGS.l2 * 2 * tf.nn.l2_loss( tf.concat(filtered_variables, axis=0)) loss = negative_log_likelihood + l2_loss if dataset_name == 'clean': metrics['test/negative_log_likelihood'].update_state( negative_log_likelihood) metrics['test/accuracy'].update_state(labels, probs) metrics['test/ece'].update_state(labels, probs) metrics['test/loss'].update_state(loss) 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)].update_state( labels, probs) strategy.run(step_fn, args=(next(iterator),)) 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) for step in range(steps_per_epoch): train_step(train_iterator) 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)) work_unit.set_notes(message) if step % 20 == 0: logging.info(message) 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) for step in range(steps_per_eval): if step % 20 == 0: logging.info('Starting to run eval step %s of epoch: %s', step, epoch) test_step(test_iterator, dataset_name) logging.info('Done with testing 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, max_intensity) 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) for i in range(FLAGS.ensemble_size): logging.info('Member %d Test Loss: %.4f, Accuracy: %.2f%%', i, metrics['test/nll_member_{}'.format(i)].result(), metrics['test/accuracy_member_{}'.format(i)].result() * 100) total_metrics = itertools.chain(metrics.items(), training_diversity.items(), test_diversity.items()) total_results = {name: metric.result() for name, metric in total_metrics} total_results.update(corrupt_results) with summary_writer.as_default(): for name, result in total_results.items(): tf.summary.scalar(name, result, step=epoch + 1) for name, result in total_results.items(): name = name.replace('/', '_') if 'negative_log_likelihood' in name: # Plots sort WIDs from high-to-low so look at maximization objectives. name = name.replace('negative_log_likelihood', 'log_likelihood') result = -result objective = work_unit.get_measurement_series(name) objective.create_measurement(result, epoch + 1) for _, metric in total_metrics: metric.reset_states() summary_writer.flush() 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)
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) 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 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_dataset = utils.load_dataset( split=tfds.Split.TRAIN, name=FLAGS.dataset, batch_size=batch_size, use_bfloat16=FLAGS.use_bfloat16) clean_test_dataset = utils.load_dataset( split=tfds.Split.TEST, name=FLAGS.dataset, batch_size=test_batch_size, use_bfloat16=FLAGS.use_bfloat16) train_dataset = strategy.experimental_distribute_dataset(train_dataset) test_datasets = { 'clean': strategy.experimental_distribute_dataset(clean_test_dataset), } if FLAGS.corruptions_interval > 0: if FLAGS.dataset == 'cifar10': load_c_dataset = utils.load_cifar10_c else: load_c_dataset = functools.partial( utils.load_cifar100_c, path=FLAGS.cifar100_c_path) corruption_types, max_intensity = utils.load_corrupted_test_info( FLAGS.dataset) for corruption in corruption_types: for intensity in range(1, max_intensity + 1): dataset = load_c_dataset( corruption_name=corruption, corruption_intensity=intensity, batch_size=test_batch_size, use_bfloat16=FLAGS.use_bfloat16) test_datasets['{0}_{1}'.format(corruption, intensity)] = ( strategy.experimental_distribute_dataset(dataset)) if FLAGS.use_bfloat16: policy = tf.keras.mixed_precision.experimental.Policy('mixed_bfloat16') tf.keras.mixed_precision.experimental.set_policy(policy) 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_condconv( input_shape=ds_info.features['image'].shape, depth=28, width_multiplier=FLAGS.resnet_width_multiplier, num_classes=num_classes, num_experts=FLAGS.num_experts, per_core_batch_size=FLAGS.per_core_batch_size, use_cond_dense=FLAGS.use_cond_dense, reduce_dense_outputs=FLAGS.reduce_dense_outputs, cond_placement=FLAGS.cond_placement, routing_fn=FLAGS.routing_fn, normalize_routing=FLAGS.normalize_routing, normalize_dense_routing=FLAGS.normalize_dense_routing, top_k=FLAGS.top_k, routing_pooling=FLAGS.routing_pooling, l2=FLAGS.l2) # reuse_routing=FLAGS.reuse_routing, # shared_routing_type=FLAGS.shared_routing_type) 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 = utils.LearningRateSchedule( 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=0.9, 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': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), 'test/negative_log_likelihood': tf.keras.metrics.Mean(), 'test/accuracy': tf.keras.metrics.SparseCategoricalAccuracy(), 'test/ece': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), } if not FLAGS.reduce_dense_outputs and FLAGS.use_cond_dense: metrics.update({ 'test/nll_poe': tf.keras.metrics.Mean(), 'test/nll_moe': tf.keras.metrics.Mean(), 'test/nll_unweighted_poe': tf.keras.metrics.Mean(), 'test/nll_unweighted_moe': tf.keras.metrics.Mean(), 'test/unweighted_gibbs_ce': tf.keras.metrics.Mean(), 'test/ece_unweighted_moe': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), 'test/accuracy_unweighted_moe': tf.keras.metrics.SparseCategoricalAccuracy(), 'test/ece_poe': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), 'test/accuracy_poe': tf.keras.metrics.SparseCategoricalAccuracy(), 'test/ece_unweighted_poe': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), 'test/accuracy_unweighted_poe': tf.keras.metrics.SparseCategoricalAccuracy(), }) for idx in range(FLAGS.num_experts): metrics['test/dense_routing_weight_{}'.format( idx)] = tf.keras.metrics.Mean() metrics['test/dense_routing_weight_normalized_{}'.format( idx)] = tf.keras.metrics.Mean() if FLAGS.corruptions_interval > 0: corrupt_metrics = {} for intensity in range(1, max_intensity + 1): 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)] = ( um.ExpectedCalibrationError(num_bins=FLAGS.num_bins)) corrupt_metrics['test/nll_weighted_moe_{}'.format(dataset_name)] = ( tf.keras.metrics.Mean()) corrupt_metrics['test/accuracy_weighted_moe_{}'.format( dataset_name)] = ( tf.keras.metrics.SparseCategoricalAccuracy()) corrupt_metrics['test/ece_weighted_moe_{}'.format(dataset_name)] = ( um.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 def _process_3d_logits(logits, routing_weights, labels): routing_weights_3d = tf.expand_dims(routing_weights, axis=-1) weighted_logits = tf.math.reduce_mean(routing_weights_3d * logits, axis=1) unweighted_logits = tf.math.reduce_mean(logits, axis=1) probs = tf.nn.softmax(logits) unweighted_probs = tf.math.reduce_mean(probs, axis=1) weighted_probs = tf.math.reduce_sum(routing_weights_3d * probs, axis=1) labels_broadcasted = tf.tile( tf.reshape(labels, (-1, 1)), (1, FLAGS.num_experts)) neg_log_likelihoods = tf.keras.losses.sparse_categorical_crossentropy( labels_broadcasted, logits, from_logits=True) unweighted_gibbs_ce = tf.math.reduce_mean(neg_log_likelihoods) weighted_gibbs_ce = tf.math.reduce_mean( tf.math.reduce_sum(routing_weights * neg_log_likelihoods, axis=1)) return { 'weighted_logits': weighted_logits, 'unweighted_logits': unweighted_logits, 'unweighted_probs': unweighted_probs, 'weighted_probs': weighted_probs, 'neg_log_likelihoods': neg_log_likelihoods, 'unweighted_gibbs_ce': unweighted_gibbs_ce, 'weighted_gibbs_ce': weighted_gibbs_ce } def _process_3d_logits_train(logits, routing_weights, labels): processing_results = _process_3d_logits(logits, routing_weights, labels) if FLAGS.loss == 'gibbs_ce': probs = processing_results['weighted_probs'] negative_log_likelihood = processing_results['weighted_gibbs_ce'] elif FLAGS.loss == 'unweighted_gibbs_ce': probs = processing_results['unweighted_probs'] negative_log_likelihood = processing_results['unweighted_gibbs_ce'] elif FLAGS.loss == 'moe': probs = processing_results['weighted_probs'] negative_log_likelihood = tf.math.reduce_mean( tf.keras.losses.sparse_categorical_crossentropy( labels, probs, from_logits=False)) elif FLAGS.loss == 'unweighted_moe': probs = processing_results['unweighted_probs'] negative_log_likelihood = tf.math.reduce_mean( tf.keras.losses.sparse_categorical_crossentropy( labels, probs, from_logits=False)) elif FLAGS.loss == 'poe': probs = tf.softmax(processing_results['weighted_logits']) negative_log_likelihood = tf.math.reduce_mean( tf.keras.losses.sparse_categorical_crossentropy( labels, processing_results['weighted_logits'], from_logits=True)) elif FLAGS.loss == 'unweighted_poe': probs = tf.softmax(processing_results['unweighted_logits']) negative_log_likelihood = tf.math.reduce_mean( tf.keras.losses.sparse_categorical_crossentropy( labels, processing_results['unweighted_logits'], from_logits=True)) return probs, negative_log_likelihood def _process_3d_logits_test(routing_weights, logits, labels): processing_results = _process_3d_logits(logits, routing_weights, labels) nll_poe = tf.math.reduce_mean( tf.keras.losses.sparse_categorical_crossentropy( labels, processing_results['weighted_logits'], from_logits=True)) nll_unweighted_poe = tf.math.reduce_mean( tf.keras.losses.sparse_categorical_crossentropy( labels, processing_results['unweighted_logits'], from_logits=True)) nll_moe = tf.math.reduce_mean( tf.keras.losses.sparse_categorical_crossentropy( labels, processing_results['weighted_probs'], from_logits=False)) nll_unweighted_moe = tf.math.reduce_mean( tf.keras.losses.sparse_categorical_crossentropy( labels, processing_results['unweighted_probs'], from_logits=False)) return { 'nll_poe': nll_poe, 'nll_moe': nll_moe, 'nll_unweighted_poe': nll_unweighted_poe, 'nll_unweighted_moe': nll_unweighted_moe, 'unweighted_gibbs_ce': processing_results['unweighted_gibbs_ce'], 'weighted_gibbs_ce': processing_results['weighted_gibbs_ce'], 'weighted_probs': processing_results['weighted_probs'], 'unweighted_probs': processing_results['unweighted_probs'], 'weighted_logits': processing_results['weighted_logits'], 'unweighted_logits': processing_results['unweighted_logits'] } @tf.function def train_step(iterator): """Training StepFn.""" def step_fn(inputs): """Per-Replica StepFn.""" images, labels = inputs with tf.GradientTape() as tape: logits = model(images, training=True) if FLAGS.use_bfloat16: logits = tf.cast(logits, tf.float32) # if not FLAGS.reduce_dense_outputs and FLAGS.use_cond_dense: if not isinstance(logits, tuple): raise ValueError('Logits are not a tuple.') # logits is a `Tensor` of shape [batch_size, num_experts, num_classes] logits, all_routing_weights = logits # routing_weights is a `Tensor` of shape [batch_size, num_experts] routing_weights = all_routing_weights[-1] if not FLAGS.reduce_dense_outputs and FLAGS.use_cond_dense: probs, negative_log_likelihood = _process_3d_logits_train( logits, routing_weights, labels) else: probs = tf.nn.softmax(logits) # Prior to reduce_mean the NLLs are of the shape [batch, num_experts]. 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)) metrics['train/ece'].update_state(labels, probs) metrics['train/loss'].update_state(loss) metrics['train/negative_log_likelihood'].update_state( negative_log_likelihood) metrics['train/accuracy'].update_state(labels, probs) strategy.run(step_fn, args=(next(iterator),)) @tf.function def test_step(iterator, dataset_name): """Evaluation StepFn.""" def step_fn(inputs): """Per-Replica StepFn.""" images, labels = inputs logits = model(images, training=False) if FLAGS.use_bfloat16: logits = tf.cast(logits, tf.float32) if not isinstance(logits, tuple): raise ValueError('Logits not a tuple') # logits is a `Tensor` of shape [batch_size, num_experts, num_classes] # routing_weights is a `Tensor` of shape [batch_size, num_experts] logits, all_routing_weights = logits routing_weights = all_routing_weights[-1] if not FLAGS.reduce_dense_outputs and FLAGS.use_cond_dense: results = _process_3d_logits_test(routing_weights, logits, labels) else: probs = tf.nn.softmax(logits) negative_log_likelihood = tf.reduce_mean( tf.keras.losses.sparse_categorical_crossentropy(labels, probs)) if dataset_name == 'clean': if not FLAGS.reduce_dense_outputs and FLAGS.use_cond_dense: metrics['test/nll_poe'].update_state(results['nll_poe']) metrics['test/nll_moe'].update_state(results['nll_moe']) metrics['test/nll_unweighted_poe'].update_state( results['nll_unweighted_poe']) metrics['test/nll_unweighted_moe'].update_state( results['nll_unweighted_moe']) metrics['test/unweighted_gibbs_ce'].update_state( results['unweighted_gibbs_ce']) metrics['test/negative_log_likelihood'].update_state( results['weighted_gibbs_ce']) metrics['test/ece'].update_state(labels, results['weighted_probs']) metrics['test/accuracy'].update_state(labels, results['weighted_probs']) metrics['test/ece_unweighted_moe'].update_state( labels, results['unweighted_probs']) metrics['test/accuracy_unweighted_moe'].update_state( labels, results['unweighted_probs']) metrics['test/ece_poe'].update_state(labels, results['weighted_logits']) metrics['test/accuracy_poe'].update_state(labels, results['weighted_logits']) metrics['test/ece_unweighted_poe'].update_state( labels, results['unweighted_logits']) metrics['test/accuracy_unweighted_poe'].update_state( labels, results['unweighted_logits']) # TODO(ghassen): summarize all routing weights not only last layer's. average_routing_weights = tf.math.reduce_mean(routing_weights, axis=0) routing_weights_sum = tf.math.reduce_sum(average_routing_weights) for idx in range(FLAGS.num_experts): metrics['test/dense_routing_weight_{}'.format(idx)].update_state( average_routing_weights[idx]) metrics['test/dense_routing_weight_normalized_{}'.format( idx)].update_state(average_routing_weights[idx] / routing_weights_sum) # TODO(ghassen): add more metrics for expert utilization, # load loss and importance/balance loss. else: metrics['test/negative_log_likelihood'].update_state( negative_log_likelihood) metrics['test/accuracy'].update_state(labels, probs) metrics['test/ece'].update_state(labels, probs) else: # TODO(ghassen): figure out how to aggregate probs for the OOD case. if not FLAGS.reduce_dense_outputs and FLAGS.use_cond_dense: corrupt_metrics['test/nll_{}'.format(dataset_name)].update_state( results['unweighted_gibbs_ce']) corrupt_metrics['test/accuracy_{}'.format(dataset_name)].update_state( labels, results['unweighted_probs']) corrupt_metrics['test/ece_{}'.format(dataset_name)].update_state( labels, results['unweighted_probs']) corrupt_metrics['test/nll_weighted_moe{}'.format( dataset_name)].update_state(results['weighted_gibbs_ce']) corrupt_metrics['test/accuracy_weighted_moe_{}'.format( dataset_name)].update_state(labels, results['weighted_probs']) corrupt_metrics['test/ece_weighted_moe{}'.format( dataset_name)].update_state(labels, results['weighted_probs']) 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)].update_state( labels, probs) 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) for step in range(steps_per_epoch): train_step(train_iterator) 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 (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) for step in range(steps_per_eval): if step % 20 == 0: logging.info('Starting to run eval step %s of epoch: %s', step, epoch) test_start_time = time.time() test_step(test_iterator, dataset_name) 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) corrupt_results = {} if (FLAGS.corruptions_interval > 0 and (epoch + 1) % FLAGS.corruptions_interval == 0): corrupt_results = utils.aggregate_corrupt_metrics(corrupt_metrics, corruption_types, max_intensity) 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) 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.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)
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) batch_size = FLAGS.per_core_batch_size * FLAGS.num_cores steps_per_epoch = APPROX_IMAGENET_TRAIN_IMAGES // batch_size steps_per_eval = IMAGENET_VALIDATION_IMAGES // batch_size 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.experimental.TPUStrategy(resolver) imagenet_train = utils.ImageNetInput(is_training=True, data_dir=FLAGS.data_dir, batch_size=FLAGS.per_core_batch_size, use_bfloat16=FLAGS.use_bfloat16) imagenet_eval = utils.ImageNetInput(is_training=False, data_dir=FLAGS.data_dir, batch_size=FLAGS.per_core_batch_size, use_bfloat16=FLAGS.use_bfloat16) test_datasets = { 'clean': strategy.experimental_distribute_datasets_from_function( imagenet_eval.input_fn) } if FLAGS.corruptions_interval > 0: 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) corrupt_input_fn = utils.corrupt_test_input_fn( batch_size=FLAGS.per_core_batch_size, corruption_name=name, corruption_intensity=intensity, use_bfloat16=FLAGS.use_bfloat16) test_datasets[dataset_name] = ( strategy.experimental_distribute_datasets_from_function( corrupt_input_fn)) train_dataset = strategy.experimental_distribute_datasets_from_function( imagenet_train.input_fn) if FLAGS.use_bfloat16: policy = tf.keras.mixed_precision.experimental.Policy('mixed_bfloat16') tf.keras.mixed_precision.experimental.set_policy(policy) with strategy.scope(): logging.info('Building Keras ResNet-50 model') model = ub.models.resnet50_sngp( input_shape=(224, 224, 3), batch_size=None, num_classes=NUM_CLASSES, use_mc_dropout=FLAGS.use_mc_dropout, dropout_rate=FLAGS.dropout_rate, filterwise_dropout=FLAGS.filterwise_dropout, use_gp_layer=FLAGS.use_gp_layer, 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_cov_discount_factor=FLAGS.gp_cov_discount_factor, gp_cov_ridge_penalty=FLAGS.gp_cov_ridge_penalty, gp_output_imagenet_initializer=FLAGS. gp_output_imagenet_initializer, 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()) # Scale learning rate and decay epochs by vanilla settings. base_lr = FLAGS.base_learning_rate * batch_size / 256 learning_rate = utils.LearningRateSchedule(steps_per_epoch, base_lr, FLAGS.train_epochs, _LR_SCHEDULE) optimizer = tf.keras.optimizers.SGD(learning_rate=learning_rate, momentum=0.9, 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': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), 'test/negative_log_likelihood': tf.keras.metrics.Mean(), 'test/accuracy': tf.keras.metrics.SparseCategoricalAccuracy(), 'test/ece': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), 'test/stddev': tf.keras.metrics.Mean(), } if FLAGS.corruptions_interval > 0: corrupt_metrics = {} for intensity in range(1, max_intensity + 1): 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)] = ( um.ExpectedCalibrationError(num_bins=FLAGS.num_bins)) corrupt_metrics['test/stddev_{}'.format(dataset_name)] = ( tf.keras.metrics.Mean()) logging.info('Finished building Keras ResNet-50 model') 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 summary_writer = tf.summary.create_file_writer( os.path.join(FLAGS.output_dir, 'summaries')) @tf.function def train_step(iterator): """Training StepFn.""" def step_fn(inputs): """Per-Replica StepFn.""" images, labels = inputs with tf.GradientTape() as tape: logits = model(images, training=True) if isinstance(logits, tuple): # If model returns a tuple of (logits, covmat), extract logits logits, _ = logits if FLAGS.use_bfloat16: logits = tf.cast(logits, tf.float32) negative_log_likelihood = tf.reduce_mean( tf.keras.losses.sparse_categorical_crossentropy( labels, logits, from_logits=True)) filtered_variables = [] for var in model.trainable_variables: # Apply l2 on the weights. This excludes BN parameters and biases, but # pay caution to their naming scheme. if 'kernel' in var.name or 'bias' in var.name: filtered_variables.append(tf.reshape(var, (-1, ))) l2_loss = FLAGS.l2 * 2 * tf.nn.l2_loss( tf.concat(filtered_variables, axis=0)) # Scale the loss given the TPUStrategy will reduce sum all gradients. loss = negative_log_likelihood + l2_loss 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'].update_state(labels, probs) 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), )) @tf.function def test_step(iterator, dataset_name): """Evaluation StepFn.""" def step_fn(inputs): """Per-Replica StepFn.""" images, labels = inputs logits_list = [] stddev_list = [] for _ in range(FLAGS.num_dropout_samples): logits = model(images, training=False) if isinstance(logits, tuple): # If model returns a tuple of (logits, covmat), extract both logits, covmat = logits else: covmat = tf.eye(FLAGS.per_core_batch_size) 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) 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) labels_broadcasted = tf.broadcast_to( labels, [FLAGS.num_dropout_samples, labels.shape[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['test/negative_log_likelihood'].update_state( negative_log_likelihood) metrics['test/accuracy'].update_state(labels, probs) metrics['test/ece'].update_state(labels, probs) metrics['test/stddev'].update_state(stddev) 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)].update_state(labels, probs) corrupt_metrics['test/stddev_{}'.format( dataset_name)].update_state(stddev) 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) for step in range(steps_per_epoch): train_step(train_iterator) 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 (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) for step in range(steps_per_eval): if step % 20 == 0: logging.info('Starting to run eval step %s of epoch: %s', step, epoch) test_start_time = time.time() test_step(test_iterator, dataset_name) 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) corrupt_results = {} if (FLAGS.corruptions_interval > 0 and (epoch + 1) % FLAGS.corruptions_interval == 0): corrupt_results = utils.aggregate_corrupt_metrics( corrupt_metrics, corruption_types, max_intensity, FLAGS.alexnet_errors_path) 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) 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.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) # Save final checkpoint. final_checkpoint_name = checkpoint.save( os.path.join(FLAGS.output_dir, 'checkpoint')) logging.info('Saved last checkpoint to %s', final_checkpoint_name) # Export final model as SavedModel. final_save_name = os.path.join(FLAGS.output_dir, 'model') model.save(final_save_name) logging.info('Saved model to %s', final_save_name)
def __init__(self, dataset_info): metric = um.ExpectedCalibrationError() super().__init__( dataset_info, metric, "ece", take_argmax=False, one_hot=False)
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) 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) 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, 'label_smoothing': FLAGS.label_smoothing, 'ensemble_size': FLAGS.ensemble_size, 'mixup_alpha': FLAGS.mixup_alpha, 'random_augment': FLAGS.random_augment, 'adaptive_mixup': FLAGS.adaptive_mixup, 'forget_mixup': FLAGS.forget_mixup, 'num_cores': FLAGS.num_cores, 'threshold': FLAGS.forget_threshold, 'cutmix': FLAGS.cutmix, } batch_size = ((FLAGS.per_core_batch_size // FLAGS.ensemble_size) * FLAGS.num_cores) train_input_fn = data_utils.load_input_fn( split=tfds.Split.TRAIN, name=FLAGS.dataset, batch_size=batch_size, use_bfloat16=FLAGS.use_bfloat16, proportion=FLAGS.train_proportion, validation_set=FLAGS.validation, aug_params=aug_params) if FLAGS.validation: validation_input_fn = data_utils.load_input_fn( split=tfds.Split.VALIDATION, name=FLAGS.dataset, batch_size=FLAGS.per_core_batch_size, use_bfloat16=FLAGS.use_bfloat16, validation_set=True) val_dataset = strategy.experimental_distribute_datasets_from_function( validation_input_fn) clean_test_input_fn = data_utils.load_input_fn( split=tfds.Split.TEST, name=FLAGS.dataset, batch_size=FLAGS.per_core_batch_size // FLAGS.ensemble_size, use_bfloat16=FLAGS.use_bfloat16) train_dataset = strategy.experimental_distribute_dataset( train_input_fn()) test_datasets = { 'clean': strategy.experimental_distribute_datasets_from_function( clean_test_input_fn), } if FLAGS.corruptions_interval > 0: if FLAGS.dataset == 'cifar10': load_c_dataset = utils.load_cifar10_c else: load_c_dataset = functools.partial(utils.load_cifar100_c, path=FLAGS.cifar100_c_path) corruption_types, max_intensity = utils.load_corrupted_test_info( FLAGS.dataset) for corruption in corruption_types: for intensity in range(1, max_intensity + 1): dataset = load_c_dataset( corruption_name=corruption, corruption_intensity=intensity, batch_size=batch_size, use_bfloat16=FLAGS.use_bfloat16) test_datasets['{0}_{1}'.format(corruption, intensity)] = ( strategy.experimental_distribute_dataset(dataset)) ds_info = tfds.builder(FLAGS.dataset).info num_train_examples = ds_info.splits['train'].num_examples # Train_proportion is a float so need to convert steps_per_epoch to int. if FLAGS.validation: # TODO(ywenxu): Remove hard-coding validation images. steps_per_epoch = int((num_train_examples * FLAGS.train_proportion - 2500) // batch_size) steps_per_val = 2500 // (FLAGS.per_core_batch_size * FLAGS.num_cores) else: steps_per_epoch = int( num_train_examples * FLAGS.train_proportion) // batch_size steps_per_eval = ds_info.splits['test'].num_examples // batch_size num_classes = ds_info.features['label'].num_classes if FLAGS.use_bfloat16: policy = tf.keras.mixed_precision.experimental.Policy('mixed_bfloat16') tf.keras.mixed_precision.experimental.set_policy(policy) summary_writer = tf.summary.create_file_writer( os.path.join(FLAGS.output_dir, 'summaries')) with strategy.scope(): logging.info('Building Keras model') model = batchensemble_model.wide_resnet( input_shape=ds_info.features['image'].shape, depth=28, width_multiplier=10, num_classes=num_classes, ensemble_size=FLAGS.ensemble_size, random_sign_init=FLAGS.random_sign_init, l2=FLAGS.l2, use_ensemble_bn=FLAGS.use_ensemble_bn) 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 = utils.LearningRateSchedule( 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=0.9, nesterov=True) diversity_schedule = tf.keras.optimizers.schedules.ExponentialDecay( FLAGS.diversity_coeff, FLAGS.diversity_decay_epoch * steps_per_epoch, decay_rate=0.97, staircase=True) metrics = { 'train/negative_log_likelihood': tf.keras.metrics.Mean(), 'train/accuracy': tf.keras.metrics.SparseCategoricalAccuracy(), 'train/loss': tf.keras.metrics.Mean(), 'train/similarity': tf.keras.metrics.Mean(), 'train/l2': tf.keras.metrics.Mean(), 'train/ece': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), 'test/negative_log_likelihood': tf.keras.metrics.Mean(), 'test/accuracy': tf.keras.metrics.SparseCategoricalAccuracy(), 'test/member_accuracy_mean': ( tf.keras.metrics.SparseCategoricalAccuracy()), 'test/ece': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), 'test/member_ece_mean': um.ExpectedCalibrationError( num_bins=FLAGS.num_bins) } for i in range(FLAGS.ensemble_size): metrics['test/nll_member_{}'.format(i)] = tf.keras.metrics.Mean() metrics['test/accuracy_member_{}'.format(i)] = ( tf.keras.metrics.SparseCategoricalAccuracy()) metrics['test/ece_member_{}'.format(i)] = ( um.ExpectedCalibrationError(num_bins=FLAGS.num_bins)) test_diversity = {} training_diversity = {} corrupt_diversity = {} if FLAGS.ensemble_size > 1: test_diversity = { 'test/disagreement': tf.keras.metrics.Mean(), 'test/average_kl': tf.keras.metrics.Mean(), 'test/cosine_similarity': tf.keras.metrics.Mean(), } training_diversity = { 'train/disagreement': tf.keras.metrics.Mean(), 'train/average_kl': tf.keras.metrics.Mean(), 'train/cosine_similarity': tf.keras.metrics.Mean(), } if FLAGS.corruptions_interval > 0: corrupt_metrics = {} for intensity in range(1, max_intensity + 1): 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)] = ( um.ExpectedCalibrationError(num_bins=FLAGS.num_bins)) corrupt_metrics['test/member_acc_mean_{}'.format(dataset_name)] = ( tf.keras.metrics.SparseCategoricalAccuracy()) corrupt_metrics['test/member_ece_mean_{}'.format(dataset_name)] = ( um.ExpectedCalibrationError(num_bins=FLAGS.num_bins)) corrupt_diversity['corrupt_diversity/average_kl_{}'.format( dataset_name)] = tf.keras.metrics.Mean() corrupt_diversity['corrupt_diversity/cosine_similarity_{}'.format( dataset_name)] = tf.keras.metrics.Mean() corrupt_diversity['corrupt_diversity/disagreement_{}'.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 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.""" if FLAGS.forget_mixup: images, labels, idx = inputs else: images, labels = inputs if FLAGS.adaptive_mixup or FLAGS.forget_mixup: images = tf.identity(images) elif FLAGS.augmix or FLAGS.random_augment: images_shape = tf.shape(images) images = tf.reshape(tf.transpose( images, [1, 0, 2, 3, 4]), [-1, images_shape[2], images_shape[3], images_shape[4]]) else: images = tf.tile(images, [FLAGS.ensemble_size, 1, 1, 1]) # Augmix, adaptive mixup, forget mixup preprocessing gives tiled labels. if FLAGS.mixup_alpha > 0 or FLAGS.label_smoothing > 0 or FLAGS.cutmix: if FLAGS.augmix or FLAGS.adaptive_mixup or FLAGS.forget_mixup: labels = tf.identity(labels) else: labels = tf.tile(labels, [FLAGS.ensemble_size, 1]) else: labels = tf.tile(labels, [FLAGS.ensemble_size]) def _is_batch_norm(v): """Decide whether a variable belongs to `batch_norm`.""" keywords = ['batchnorm', 'batch_norm', 'bn'] return any([k in v.name.lower() for k in keywords]) def _normalize(x): """Normalize an input with l2 norm.""" l2 = tf.norm(x, ord=2, axis=-1) return x / tf.expand_dims(l2, axis=-1) # Taking the sum of upper triangular of XX^T and divided by ensemble size. def pairwise_cosine_distance(x): """Compute the pairwise distance in a matrix.""" normalized_x = _normalize(x) return (tf.reduce_sum( tf.matmul(normalized_x, normalized_x, transpose_b=True)) - FLAGS.ensemble_size) / (2.0 * FLAGS.ensemble_size) with tf.GradientTape() as tape: logits = model(images, training=True) if FLAGS.use_bfloat16: logits = tf.cast(logits, tf.float32) if FLAGS.mixup_alpha > 0 or FLAGS.label_smoothing > 0 or FLAGS.cutmix: 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) fast_weights = [var for var in model.trainable_variables if not _is_batch_norm(var) and ( 'alpha' in var.name or 'gamma' in var.name)] pairwise_distance_loss = tf.add_n( [pairwise_cosine_distance(var) for var in fast_weights]) diversity_start_iter = steps_per_epoch * FLAGS.diversity_start_epoch diversity_iterations = optimizer.iterations - diversity_start_iter if diversity_iterations > 0: diversity_coeff = diversity_schedule(diversity_iterations) diversity_loss = diversity_coeff * pairwise_distance_loss loss = negative_log_likelihood + l2_loss + diversity_loss else: 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) # Separate learning rate implementation. if FLAGS.fast_weight_lr_multiplier != 1.0: grads_and_vars = [] for grad, var in zip(grads, model.trainable_variables): # Apply different learning rate on the fast weight approximate # posterior/prior parameters. This is excludes BN and slow weights, # but pay caution to the naming scheme. if (not _is_batch_norm(var) and 'kernel' not in var.name): grads_and_vars.append((grad * FLAGS.fast_weight_lr_multiplier, var)) else: grads_and_vars.append((grad, var)) optimizer.apply_gradients(grads_and_vars) else: optimizer.apply_gradients(zip(grads, model.trainable_variables)) probs = tf.nn.softmax(logits) if FLAGS.ensemble_size > 1: per_probs = tf.reshape( probs, tf.concat([[FLAGS.ensemble_size, -1], probs.shape[1:]], 0)) diversity_results = um.average_pairwise_diversity( per_probs, FLAGS.ensemble_size) for k, v in diversity_results.items(): training_diversity['train/' + k].update_state(v) if FLAGS.mixup_alpha > 0 or FLAGS.label_smoothing > 0 or FLAGS.cutmix: labels = tf.argmax(labels, axis=-1) metrics['train/ece'].update_state(labels, probs) metrics['train/similarity'].update_state(pairwise_distance_loss) metrics['train/l2'].update_state(l2_loss) metrics['train/loss'].update_state(loss) metrics['train/negative_log_likelihood'].update_state( negative_log_likelihood) metrics['train/accuracy'].update_state(labels, logits) if FLAGS.forget_mixup: train_predictions = tf.argmax(probs, -1) labels = tf.cast(labels, train_predictions.dtype) # For each ensemble member, we accumulate the accuracy counts. accuracy_counts = tf.cast(tf.reshape( (train_predictions == labels), [FLAGS.ensemble_size, -1]), tf.float32) return accuracy_counts, idx if FLAGS.forget_mixup: return strategy.run(step_fn, args=(next(iterator),)) else: strategy.run(step_fn, args=(next(iterator),)) @tf.function def test_step(iterator, dataset_name): """Evaluation StepFn.""" def step_fn(inputs): """Per-Replica StepFn.""" images, labels = inputs images = tf.tile(images, [FLAGS.ensemble_size, 1, 1, 1]) logits = model(images, training=False) if FLAGS.use_bfloat16: logits = tf.cast(logits, tf.float32) probs = tf.nn.softmax(logits) per_probs = tf.split(probs, num_or_size_splits=FLAGS.ensemble_size, axis=0) for i in range(FLAGS.ensemble_size): member_probs = per_probs[i] if dataset_name == 'clean': 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/member_accuracy_mean'].update_state( labels, member_probs) metrics['test/ece_member_{}'.format(i)].update_state(labels, member_probs) metrics['test/member_ece_mean'].update_state(labels, member_probs) elif dataset_name != 'validation': corrupt_metrics['test/member_acc_mean_{}'.format( dataset_name)].update_state(labels, member_probs) corrupt_metrics['test/member_ece_mean_{}'.format( dataset_name)].update_state(labels, member_probs) if FLAGS.ensemble_size > 1: per_probs_tensor = tf.reshape( probs, tf.concat([[FLAGS.ensemble_size, -1], probs.shape[1:]], 0)) diversity_results = um.average_pairwise_diversity( per_probs_tensor, FLAGS.ensemble_size) if dataset_name == 'clean': for k, v in diversity_results.items(): test_diversity['test/' + k].update_state(v) elif dataset_name != 'validation': for k, v in diversity_results.items(): corrupt_diversity['corrupt_diversity/{}_{}'.format( k, dataset_name)].update_state(v) probs = tf.reduce_mean(per_probs, axis=0) negative_log_likelihood = tf.reduce_mean( tf.keras.losses.sparse_categorical_crossentropy(labels, probs)) if dataset_name == 'clean': metrics['test/negative_log_likelihood'].update_state( negative_log_likelihood) metrics['test/accuracy'].update_state(labels, probs) metrics['test/ece'].update_state(labels, probs) elif dataset_name != 'validation': 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)].update_state( labels, probs) if dataset_name == 'validation': return per_probs_tensor, labels if dataset_name == 'validation': return strategy.run(step_fn, args=(next(iterator),)) else: strategy.run(step_fn, args=(next(iterator),)) train_iterator = iter(train_dataset) start_time = time.time() forget_counts_history = [] for epoch in range(initial_epoch, FLAGS.train_epochs): logging.info('Starting to run epoch: %s', epoch) acc_counts_list = [] idx_list = [] for step in range(steps_per_epoch): if FLAGS.forget_mixup: temp_accuracy_counts, temp_idx = train_step(train_iterator) acc_counts_list.append(temp_accuracy_counts) idx_list.append(temp_idx) else: train_step(train_iterator) 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) # Only one of the forget_mixup and adaptive_mixup can be true. if FLAGS.forget_mixup: current_acc = [tf.concat(list(acc_counts_list[i].values), axis=1) for i in range(len(acc_counts_list))] total_idx = [tf.concat(list(idx_list[i].values), axis=0) for i in range(len(idx_list))] current_acc = tf.cast(tf.concat(current_acc, axis=1), tf.int32) total_idx = tf.concat(total_idx, axis=0) current_forget_path = os.path.join(FLAGS.output_dir, 'forget_counts.npy') last_acc_path = os.path.join(FLAGS.output_dir, 'last_acc.npy') if epoch == 0: forget_counts = tf.zeros( [FLAGS.ensemble_size, num_train_examples], dtype=tf.int32) last_acc = tf.zeros( [FLAGS.ensemble_size, num_train_examples], dtype=tf.int32) else: if 'last_acc' not in locals(): with tf.io.gfile.GFile(last_acc_path, 'rb') as f: last_acc = np.load(f) last_acc = tf.cast(tf.convert_to_tensor(last_acc), tf.int32) if 'forget_counts' not in locals(): with tf.io.gfile.GFile(current_forget_path, 'rb') as f: forget_counts = np.load(f) forget_counts = tf.cast(tf.convert_to_tensor(forget_counts), tf.int32) selected_last_acc = tf.gather(last_acc, total_idx, axis=1) forget_this_epoch = tf.cast(current_acc < selected_last_acc, tf.int32) forget_this_epoch = tf.transpose(forget_this_epoch) target_shape = tf.constant([num_train_examples, FLAGS.ensemble_size]) current_forget_counts = tf.scatter_nd(tf.reshape(total_idx, [-1, 1]), forget_this_epoch, target_shape) current_forget_counts = tf.transpose(current_forget_counts) acc_this_epoch = tf.transpose(current_acc) last_acc = tf.scatter_nd(tf.reshape(total_idx, [-1, 1]), acc_this_epoch, target_shape) # This is lower bound of true acc. last_acc = tf.transpose(last_acc) # TODO(ywenxu): We count the dropped examples as forget. Fix this later. forget_counts += current_forget_counts forget_counts_history.append(forget_counts) logging.info('forgetting counts') logging.info(tf.stack(forget_counts_history, 0)) with tf.io.gfile.GFile(os.path.join( FLAGS.output_dir, 'forget_counts_history.npy'), 'wb') as f: np.save(f, tf.stack(forget_counts_history, 0).numpy()) with tf.io.gfile.GFile(current_forget_path, 'wb') as f: np.save(f, forget_counts.numpy()) with tf.io.gfile.GFile(last_acc_path, 'wb') as f: np.save(f, last_acc.numpy()) aug_params['forget_counts_dir'] = current_forget_path train_input_fn = data_utils.load_input_fn( split=tfds.Split.TRAIN, name=FLAGS.dataset, batch_size=FLAGS.num_cores * ( FLAGS.per_core_batch_size // FLAGS.ensemble_size), use_bfloat16=FLAGS.use_bfloat16, validation_set=FLAGS.validation, aug_params=aug_params) train_dataset = strategy.experimental_distribute_dataset( train_input_fn()) train_iterator = iter(train_dataset) if FLAGS.adaptive_mixup: val_iterator = iter(val_dataset) logging.info('Testing on validation dataset') predictions_list = [] labels_list = [] for step in range(steps_per_val): temp_predictions, temp_labels = test_step(val_iterator, 'validation') predictions_list.append(temp_predictions) labels_list.append(temp_labels) predictions = [tf.concat(list(predictions_list[i].values), axis=1) for i in range(len(predictions_list))] labels = [tf.concat(list(labels_list[i].values), axis=0) for i in range(len(labels_list))] predictions = tf.concat(predictions, axis=1) labels = tf.cast(tf.concat(labels, axis=0), tf.int64) def compute_acc_conf(preds, label, focus_class): class_preds = tf.boolean_mask(preds, label == focus_class, axis=1) class_pred_labels = tf.argmax(class_preds, axis=-1) confidence = tf.reduce_mean(tf.reduce_max(class_preds, axis=-1), -1) accuracy = tf.reduce_mean(tf.cast( class_pred_labels == focus_class, tf.float32), axis=-1) return accuracy - confidence calibration_per_class = [compute_acc_conf( predictions, labels, i) for i in range(num_classes)] calibration_per_class = tf.stack(calibration_per_class, axis=1) logging.info('calibration per class') logging.info(calibration_per_class) mixup_coeff = tf.where(calibration_per_class > 0, 1.0, FLAGS.mixup_alpha) mixup_coeff = tf.clip_by_value(mixup_coeff, 0, 1) logging.info('mixup coeff') logging.info(mixup_coeff) aug_params['mixup_coeff'] = mixup_coeff train_input_fn = data_utils.load_input_fn( split=tfds.Split.TRAIN, name=FLAGS.dataset, batch_size=FLAGS.per_core_batch_size // FLAGS.ensemble_size, use_bfloat16=FLAGS.use_bfloat16, validation_set=True, aug_params=aug_params) train_dataset = strategy.experimental_distribute_datasets_from_function( train_input_fn) train_iterator = iter(train_dataset) 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) for step in range(steps_per_eval): if step % 20 == 0: logging.info('Starting to run eval step %s of epoch: %s', step, epoch) test_step(test_iterator, dataset_name) logging.info('Done with testing on %s', dataset_name) corrupt_results = {} if (FLAGS.corruptions_interval > 0 and (epoch + 1) % FLAGS.corruptions_interval == 0): # This includes corrupt_diversity whose disagreement normalized by its # corrupt mean error rate. corrupt_results = utils.aggregate_corrupt_metrics( corrupt_metrics, corruption_types, max_intensity, corrupt_diversity=corrupt_diversity, output_dir=FLAGS.output_dir) 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) for i in range(FLAGS.ensemble_size): logging.info('Member %d Test Loss: %.4f, Accuracy: %.2f%%', i, metrics['test/nll_member_{}'.format(i)].result(), metrics['test/accuracy_member_{}'.format(i)].result() * 100) total_metrics = metrics.copy() total_metrics.update(training_diversity) total_metrics.update(test_diversity) total_results = {name: metric.result() for name, metric in total_metrics.items()} total_results.update(corrupt_results) # Normalize all disagreement metrics (training, testing) by test accuracy. # Disagreement on corrupt dataset is normalized by their own error rate. test_acc = total_metrics['test/accuracy'].result() for name, metric in total_metrics.items(): if 'disagreement' in name: total_results[name] = metric.result() / test_acc with summary_writer.as_default(): for name, result in total_results.items(): tf.summary.scalar(name, result, step=epoch + 1) if FLAGS.forget_mixup: tf.summary.histogram('forget_counts', forget_counts, step=epoch + 1) for metric in total_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) # Need to store the last but one checkpoint in adaptive mixup setup. if FLAGS.adaptive_mixup and epoch == (FLAGS.train_epochs - 2): checkpoint_name = checkpoint.save( os.path.join(FLAGS.output_dir, 'last_but_one_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)
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) 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 steps_per_epoch = train_dataset_size // batch_size steps_per_eval = ds_info.splits['test'].num_examples // batch_size num_classes = ds_info.features['label'].num_classes train_dataset = ub.datasets.get( FLAGS.dataset, split=tfds.Split.TRAIN).load(batch_size=batch_size) clean_test_dataset = ub.datasets.get( FLAGS.dataset, split=tfds.Split.TEST).load(batch_size=batch_size) train_dataset = strategy.experimental_distribute_dataset(train_dataset) test_datasets = { 'clean': strategy.experimental_distribute_dataset(clean_test_dataset), } if FLAGS.corruptions_interval > 0: extra_kwargs = {} if FLAGS.dataset == 'cifar100': extra_kwargs['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, **extra_kwargs).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_variational( input_shape=ds_info.features['image'].shape, depth=28, width_multiplier=10, num_classes=num_classes, prior_stddev=FLAGS.prior_stddev, dataset_size=train_dataset_size, stddev_init=FLAGS.stddev_init) 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 = utils.LearningRateSchedule( 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=0.9, 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': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), 'train/kl': tf.keras.metrics.Mean(), 'train/kl_scale': tf.keras.metrics.Mean(), 'test/negative_log_likelihood': tf.keras.metrics.Mean(), 'test/accuracy': tf.keras.metrics.SparseCategoricalAccuracy(), 'test/ece': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), } 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)] = ( um.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'] 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)) filtered_variables = [] for var in model.trainable_variables: # Apply l2 on the BN parameters and bias terms. This # excludes only fast weight approximate posterior/prior parameters, # but pay caution to their naming scheme. if 'batch_norm' in var.name or 'bias' in var.name: filtered_variables.append(tf.reshape(var, (-1, ))) l2_loss = FLAGS.l2 * 2 * tf.nn.l2_loss( tf.concat(filtered_variables, axis=0)) kl = sum(model.losses) kl_scale = tf.cast(optimizer.iterations + 1, kl.dtype) kl_scale /= steps_per_epoch * FLAGS.kl_annealing_epochs kl_scale = tf.minimum(1., kl_scale) kl_loss = kl_scale * kl # Scale the loss given the TPUStrategy will reduce sum all gradients. loss = negative_log_likelihood + l2_loss + kl_loss 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'].update_state(labels, probs) metrics['train/loss'].update_state(loss) metrics['train/negative_log_likelihood'].update_state( negative_log_likelihood) metrics['train/kl'].update_state(kl) metrics['train/kl_scale'].update_state(kl_scale) metrics['train/accuracy'].update_state(labels, logits) strategy.run(step_fn, args=(next(iterator), )) @tf.function def test_step(iterator, dataset_name): """Evaluation StepFn.""" def step_fn(inputs): """Per-Replica StepFn.""" images = inputs['features'] labels = inputs['labels'] # TODO(trandustin): Use more eval samples only on corrupted predictions; # it's expensive but a one-time compute if scheduled post-training. if FLAGS.num_eval_samples > 1 and dataset_name != 'clean': logits = tf.stack([ model(images, training=False) for _ in range(FLAGS.num_eval_samples) ], axis=0) else: logits = model(images, training=False) probs = tf.nn.softmax(logits) if FLAGS.num_eval_samples > 1 and dataset_name != 'clean': probs = tf.reduce_mean(probs, axis=0) negative_log_likelihood = tf.reduce_mean( tf.keras.losses.sparse_categorical_crossentropy(labels, probs)) if dataset_name == 'clean': metrics['test/negative_log_likelihood'].update_state( negative_log_likelihood) metrics['test/accuracy'].update_state(labels, probs) metrics['test/ece'].update_state(labels, probs) 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)].update_state(labels, probs) 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) for step in range(steps_per_epoch): train_step(train_iterator) 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 (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) for step in range(steps_per_eval): if step % 20 == 0: logging.info('Starting to run eval step %s of epoch: %s', step, epoch) test_start_time = time.time() test_step(test_iterator, dataset_name) 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) 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) 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.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)
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) 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.experimental.TPUStrategy(resolver) per_core_batch_size = FLAGS.per_core_batch_size // FLAGS.ensemble_size check_bool = FLAGS.train_proportion > 0 and FLAGS.train_proportion <= 1 assert check_bool, 'Proportion of train set has to meet 0 < prop <= 1.' drop_remainder_validation = True if not FLAGS.use_gpu: # This has to be True for TPU traing, otherwise the batchsize of images in # the validation set can't be determined by TPU compile. assert drop_remainder_validation, 'drop_remainder must be True in TPU mode.' train_input_fn = utils.load_input_fn(split=tfds.Split.TRAIN, name=FLAGS.dataset, batch_size=per_core_batch_size, use_bfloat16=FLAGS.use_bfloat16, repeat=True, proportion=FLAGS.train_proportion) validation_proportion = 1 - FLAGS.train_proportion validation_input_fn = utils.load_input_fn( split=tfds.Split.VALIDATION, name=FLAGS.dataset, batch_size=per_core_batch_size, use_bfloat16=FLAGS.use_bfloat16, repeat=True, proportion=validation_proportion, drop_remainder=drop_remainder_validation) clean_test_input_fn = utils.load_input_fn(split=tfds.Split.TEST, name=FLAGS.dataset, batch_size=per_core_batch_size, use_bfloat16=FLAGS.use_bfloat16) train_dataset = strategy.experimental_distribute_datasets_from_function( train_input_fn) validation_dataset = strategy.experimental_distribute_datasets_from_function( validation_input_fn) test_datasets = { 'clean': strategy.experimental_distribute_datasets_from_function( clean_test_input_fn), } if FLAGS.corruptions_interval > 0: if FLAGS.dataset == 'cifar10': load_c_input_fn = utils.load_cifar10_c_input_fn else: load_c_input_fn = functools.partial(utils.load_cifar100_c_input_fn, path=FLAGS.cifar100_c_path) corruption_types, max_intensity = utils.load_corrupted_test_info( FLAGS.dataset) for corruption in corruption_types: for intensity in range(1, max_intensity + 1): input_fn = load_c_input_fn(corruption_name=corruption, corruption_intensity=intensity, batch_size=per_core_batch_size, use_bfloat16=FLAGS.use_bfloat16) test_datasets['{0}_{1}'.format(corruption, intensity)] = ( strategy.experimental_distribute_datasets_from_function( input_fn)) ds_info = tfds.builder(FLAGS.dataset).info batch_size = per_core_batch_size * FLAGS.num_cores train_sample_size = ds_info.splits[ 'train'].num_examples * FLAGS.train_proportion steps_per_epoch = int(train_sample_size / batch_size) train_sample_size = int(train_sample_size) steps_per_eval = ds_info.splits['test'].num_examples // batch_size num_classes = ds_info.features['label'].num_classes if FLAGS.use_bfloat16: policy = tf.keras.mixed_precision.experimental.Policy('mixed_bfloat16') tf.keras.mixed_precision.experimental.set_policy(policy) summary_writer = tf.summary.create_file_writer( os.path.join(FLAGS.output_dir, 'summaries')) logging.info('Building Keras model.') depth = 28 width = 10 dict_ranges = {'min': FLAGS.min_l2_range, 'max': FLAGS.max_l2_range} ranges = [dict_ranges for _ in range(6)] # 6 independent l2 parameters model_config = { 'key_to_index': { 'input_conv_l2_kernel': 0, 'group_l2_kernel': 1, 'group_1_l2_kernel': 2, 'group_2_l2_kernel': 3, 'dense_l2_kernel': 4, 'dense_l2_bias': 5, }, 'ranges': ranges, 'test': None } lambdas_config = LambdaConfig(model_config['ranges'], model_config['key_to_index']) if FLAGS.e_body_hidden_units > 0: e_body_arch = '({},)'.format(FLAGS.e_body_hidden_units) else: e_body_arch = '()' e_shared_arch = '()' e_activation = 'tanh' filters_resnet = [16] for i in range(0, 3): # 3 groups of blocks filters_resnet.extend([16 * width * 2**i] * 9) # 9 layers in each block # e_head dim for conv2d is just the number of filters (only # kernel) and twice num of classes for the last dense layer (kernel + bias) e_head_dims = [x for x in filters_resnet] + [2 * num_classes] with strategy.scope(): e_models = e_factory( lambdas_config.input_shape, e_head_dims=e_head_dims, e_body_arch=eval(e_body_arch), # pylint: disable=eval-used e_shared_arch=eval(e_shared_arch), # pylint: disable=eval-used activation=e_activation, use_bias=FLAGS.e_model_use_bias, e_head_init=FLAGS.init_emodels_stddev) model = wide_resnet_hyperbatchensemble( input_shape=ds_info.features['image'].shape, depth=depth, width_multiplier=width, num_classes=num_classes, ensemble_size=FLAGS.ensemble_size, random_sign_init=FLAGS.random_sign_init, config=lambdas_config, e_models=e_models, l2_batchnorm_layer=FLAGS.l2_batchnorm, regularize_fast_weights=FLAGS.regularize_fast_weights, fast_weights_eq_contraint=FLAGS.fast_weights_eq_contraint, 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()) # build hyper-batchensemble complete ------------------------- # Initialize Lambda distributions for tuning lambdas_mean = tf.reduce_mean( log_uniform_mean([lambdas_config.log_min, lambdas_config.log_max])) lambdas0 = tf.random.normal((FLAGS.ensemble_size, lambdas_config.dim), lambdas_mean, 0.1 * FLAGS.ens_init_delta_bounds) lower0 = lambdas0 - tf.constant(FLAGS.ens_init_delta_bounds) lower0 = tf.maximum(lower0, 1e-8) upper0 = lambdas0 + tf.constant(FLAGS.ens_init_delta_bounds) log_lower = tf.Variable(tf.math.log(lower0)) log_upper = tf.Variable(tf.math.log(upper0)) lambda_parameters = [log_lower, log_upper] # these variables are tuned clip_lambda_parameters(lambda_parameters, lambdas_config) # Optimizer settings to train model weights # Linearly scale learning rate and the decay epochs by vanilla settings. # Note: Here, we don't divide the epochs by 200 as for the other uncertainty # baselines. base_lr = FLAGS.base_learning_rate * batch_size / 128 lr_decay_epochs = [int(l) for l in FLAGS.lr_decay_epochs] lr_schedule = utils.LearningRateSchedule( 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=0.9, nesterov=True) # tuner used for optimizing lambda_parameters tuner = tf.keras.optimizers.Adam(FLAGS.lr_tuning) metrics = { 'train/negative_log_likelihood': tf.keras.metrics.Mean(), 'train/accuracy': tf.keras.metrics.SparseCategoricalAccuracy(), 'train/loss': tf.keras.metrics.Mean(), 'train/ece': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), 'train/disagreement': tf.keras.metrics.Mean(), 'train/average_kl': tf.keras.metrics.Mean(), 'train/cosine_similarity': tf.keras.metrics.Mean(), 'test/negative_log_likelihood': tf.keras.metrics.Mean(), 'test/accuracy': tf.keras.metrics.SparseCategoricalAccuracy(), 'test/ece': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), 'test/gibbs_nll': tf.keras.metrics.Mean(), 'test/gibbs_accuracy': tf.keras.metrics.SparseCategoricalAccuracy(), 'test/disagreement': tf.keras.metrics.Mean(), 'test/average_kl': tf.keras.metrics.Mean(), 'test/cosine_similarity': tf.keras.metrics.Mean(), 'validation/loss': tf.keras.metrics.Mean(), 'validation/loss_entropy': tf.keras.metrics.Mean(), 'validation/loss_ce': tf.keras.metrics.Mean() } corrupt_metrics = {} for i in range(FLAGS.ensemble_size): metrics['test/nll_member_{}'.format(i)] = tf.keras.metrics.Mean() metrics['test/accuracy_member_{}'.format(i)] = ( tf.keras.metrics.SparseCategoricalAccuracy()) if FLAGS.corruptions_interval > 0: for intensity in range(1, max_intensity + 1): 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)] = ( um.ExpectedCalibrationError(num_bins=FLAGS.num_bins)) checkpoint = tf.train.Checkpoint(model=model, lambda_parameters=lambda_parameters, optimizer=optimizer) latest_checkpoint = tf.train.latest_checkpoint(FLAGS.output_dir) initial_epoch = 0 if latest_checkpoint and FLAGS.restore_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, labels = inputs images = tf.tile(images, [FLAGS.ensemble_size, 1, 1, 1]) # generate lambdas lambdas = log_uniform_sample(per_core_batch_size, lambda_parameters) lambdas = tf.reshape(lambdas, (FLAGS.ensemble_size * per_core_batch_size, lambdas_config.dim)) with tf.GradientTape() as tape: logits = model([images, lambdas], training=True) if FLAGS.use_bfloat16: logits = tf.cast(logits, tf.float32) if FLAGS.use_gibbs_ce: # Average of single model CEs # tiling of labels should be only done for Gibbs CE loss labels = tf.tile(labels, [FLAGS.ensemble_size]) negative_log_likelihood = tf.reduce_mean( tf.keras.losses.sparse_categorical_crossentropy( labels, logits, from_logits=True)) else: # Ensemble CE uses no tiling of the labels negative_log_likelihood = ensemble_crossentropy( labels, logits, FLAGS.ensemble_size) # Note: Divide l2_loss by sample_size (this differs from uncertainty_ # baselines implementation.) l2_loss = sum(model.losses) / train_sample_size 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) # Separate learning rate for fast weights. grads_and_vars = [] for grad, var in zip(grads, model.trainable_variables): if (('alpha' in var.name or 'gamma' in var.name) and 'batch_norm' not in var.name): grads_and_vars.append( (grad * FLAGS.fast_weight_lr_multiplier, var)) else: grads_and_vars.append((grad, var)) optimizer.apply_gradients(grads_and_vars) probs = tf.nn.softmax(logits) per_probs = tf.split(probs, num_or_size_splits=FLAGS.ensemble_size, axis=0) per_probs_stacked = tf.stack(per_probs, axis=0) metrics['train/ece'].update_state(labels, probs) metrics['train/loss'].update_state(loss) metrics['train/negative_log_likelihood'].update_state( negative_log_likelihood) metrics['train/accuracy'].update_state(labels, logits) diversity_results = um.average_pairwise_diversity( per_probs_stacked, FLAGS.ensemble_size) for k, v in diversity_results.items(): metrics['train/' + k].update_state(v) if grads_and_vars: grads, _ = zip(*grads_and_vars) strategy.run(step_fn, args=(next(iterator), )) @tf.function def tuning_step(iterator): """Tuning StepFn.""" def step_fn(inputs): """Per-Replica StepFn.""" images, labels = inputs images = tf.tile(images, [FLAGS.ensemble_size, 1, 1, 1]) with tf.GradientTape(watch_accessed_variables=False) as tape: tape.watch(lambda_parameters) # sample lambdas if FLAGS.sample_and_tune: lambdas = log_uniform_sample(per_core_batch_size, lambda_parameters) else: lambdas = log_uniform_mean(lambda_parameters) lambdas = tf.repeat(lambdas, per_core_batch_size, axis=0) lambdas = tf.reshape(lambdas, (FLAGS.ensemble_size * per_core_batch_size, lambdas_config.dim)) # ensemble CE logits = model([images, lambdas], training=False) ce = ensemble_crossentropy(labels, logits, FLAGS.ensemble_size) # entropy penalty for lambda distribution entropy = FLAGS.tau * log_uniform_entropy(lambda_parameters) loss = ce - entropy scaled_loss = loss / strategy.num_replicas_in_sync gradients = tape.gradient(loss, lambda_parameters) tuner.apply_gradients(zip(gradients, lambda_parameters)) metrics['validation/loss_ce'].update_state( ce / strategy.num_replicas_in_sync) metrics['validation/loss_entropy'].update_state( entropy / strategy.num_replicas_in_sync) metrics['validation/loss'].update_state(scaled_loss) strategy.run(step_fn, args=(next(iterator), )) @tf.function def test_step(iterator, dataset_name): """Evaluation StepFn.""" def step_fn(inputs): """Per-Replica StepFn.""" # Note that we don't use tf.tile for labels here images, labels = inputs images = tf.tile(images, [FLAGS.ensemble_size, 1, 1, 1]) # get lambdas lambdas = log_uniform_mean(lambda_parameters) rep_lambdas = tf.repeat(lambdas, per_core_batch_size, axis=0) # eval on testsets logits = model([images, rep_lambdas], training=False) if FLAGS.use_bfloat16: logits = tf.cast(logits, tf.float32) probs = tf.nn.softmax(logits) per_probs = tf.split(probs, num_or_size_splits=FLAGS.ensemble_size, axis=0) # per member performance and gibbs performance (average per member perf) if dataset_name == 'clean': for i in range(FLAGS.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) labels_tile = tf.tile(labels, [FLAGS.ensemble_size]) metrics['test/gibbs_nll'].update_state( tf.reduce_mean( tf.keras.losses.sparse_categorical_crossentropy( labels_tile, logits, from_logits=True))) metrics['test/gibbs_accuracy'].update_state(labels_tile, probs) # ensemble performance negative_log_likelihood = ensemble_crossentropy( labels, logits, FLAGS.ensemble_size) probs = tf.reduce_mean(per_probs, axis=0) if dataset_name == 'clean': metrics['test/negative_log_likelihood'].update_state( negative_log_likelihood) metrics['test/accuracy'].update_state(labels, probs) metrics['test/ece'].update_state(labels, probs) 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)].update_state(labels, probs) if dataset_name == 'clean': per_probs_stacked = tf.stack(per_probs, axis=0) diversity_results = um.average_pairwise_diversity( per_probs_stacked, FLAGS.ensemble_size) for k, v in diversity_results.items(): metrics['test/' + k].update_state(v) strategy.run(step_fn, args=(next(iterator), )) logging.info('--- Starting training using %d examples. ---', train_sample_size) train_iterator = iter(train_dataset) validation_iterator = iter(validation_dataset) start_time = time.time() for epoch in range(initial_epoch, FLAGS.train_epochs): logging.info('Starting to run epoch: %s', epoch) for step in range(steps_per_epoch): train_step(train_iterator) do_tuning = (epoch >= FLAGS.tuning_warmup_epochs) if do_tuning and ((step + 1) % FLAGS.tuning_every_x_step == 0): tuning_step(validation_iterator) # clip lambda parameters if outside of range clip_lambda_parameters(lambda_parameters, lambdas_config) 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) # evaluate on test data 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) for step in range(steps_per_eval): if step % 20 == 0: logging.info('Starting to run eval step %s of epoch: %s', step, epoch) test_step(test_iterator, dataset_name) logging.info('Done with testing 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, max_intensity) logging.info('Train Loss: %.4f, Accuracy: %.2f%%', metrics['train/loss'].result(), metrics['train/accuracy'].result() * 100) logging.info('Validation Loss: %.4f, CE: %.4f, Entropy: %.4f', metrics['validation/loss'].result(), metrics['validation/loss_ce'].result(), metrics['validation/loss_entropy'].result()) logging.info('Test NLL: %.4f, Accuracy: %.2f%%', metrics['test/negative_log_likelihood'].result(), metrics['test/accuracy'].result() * 100) for i in range(FLAGS.ensemble_size): logging.info( 'Member %d Test Loss: %.4f, Accuracy: %.2f%%', i, metrics['test/nll_member_{}'.format(i)].result(), metrics['test/accuracy_member_{}'.format(i)].result() * 100) total_results = { name: metric.result() for name, metric in metrics.items() } total_results.update({ name: metric.result() for name, metric in corrupt_metrics.items() }) total_results.update(corrupt_results) 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() # save checkpoint and lambdas config if (FLAGS.checkpoint_interval > 0 and (epoch + 1) % FLAGS.checkpoint_interval == 0): checkpoint_name = checkpoint.save( os.path.join(FLAGS.output_dir, 'checkpoint')) lambdas_cf = lambdas_config.get_config() filepath = os.path.join(FLAGS.output_dir, 'lambdas_config.p') with tf.io.gfile.GFile(filepath, 'wb') as fp: pickle.dump(lambdas_cf, fp, protocol=pickle.HIGHEST_PROTOCOL) logging.info('Saved checkpoint to %s', checkpoint_name)
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) 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) 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, 'adaptive_mixup': FLAGS.adaptive_mixup, 'random_augment': FLAGS.random_augment, 'forget_mixup': FLAGS.forget_mixup, 'num_cores': FLAGS.num_cores, 'threshold': FLAGS.forget_threshold, } batch_size = (FLAGS.per_core_batch_size * FLAGS.num_cores // FLAGS.num_dropout_samples_training) train_input_fn = data_utils.load_input_fn( split=tfds.Split.TRAIN, name=FLAGS.dataset, batch_size=batch_size, use_bfloat16=FLAGS.use_bfloat16, proportion=FLAGS.train_proportion, validation_set=FLAGS.validation, aug_params=aug_params) if FLAGS.validation: validation_input_fn = data_utils.load_input_fn( split=tfds.Split.VALIDATION, name=FLAGS.dataset, batch_size=FLAGS.per_core_batch_size, use_bfloat16=FLAGS.use_bfloat16, validation_set=True) val_dataset = strategy.experimental_distribute_datasets_from_function( validation_input_fn) clean_test_input_fn = data_utils.load_input_fn( split=tfds.Split.TEST, name=FLAGS.dataset, batch_size=FLAGS.per_core_batch_size, use_bfloat16=FLAGS.use_bfloat16) train_dataset = strategy.experimental_distribute_dataset(train_input_fn()) test_datasets = { 'clean': strategy.experimental_distribute_datasets_from_function( clean_test_input_fn), } if FLAGS.corruptions_interval > 0: if FLAGS.dataset == 'cifar10': load_c_dataset = utils.load_cifar10_c else: load_c_dataset = functools.partial(utils.load_cifar100_c, path=FLAGS.cifar100_c_path) corruption_types, max_intensity = utils.load_corrupted_test_info( FLAGS.dataset) for corruption in corruption_types: for intensity in range(1, max_intensity + 1): dataset = load_c_dataset(corruption_name=corruption, corruption_intensity=intensity, batch_size=FLAGS.per_core_batch_size * FLAGS.num_cores, use_bfloat16=FLAGS.use_bfloat16) test_datasets['{0}_{1}'.format(corruption, intensity)] = ( strategy.experimental_distribute_dataset(dataset)) ds_info = tfds.builder(FLAGS.dataset).info num_train_examples = ds_info.splits['train'].num_examples # Train_proportion is a float so need to convert steps_per_epoch to int. if FLAGS.validation: # TODO(ywenxu): Remove hard-coding validation images. steps_per_epoch = int( (num_train_examples * FLAGS.train_proportion - 2500) // batch_size) steps_per_val = 2500 // (FLAGS.per_core_batch_size * FLAGS.num_cores) else: steps_per_epoch = int( num_train_examples * FLAGS.train_proportion) // batch_size steps_per_eval = ds_info.splits['test'].num_examples // batch_size num_classes = ds_info.features['label'].num_classes if FLAGS.use_bfloat16: policy = tf.keras.mixed_precision.experimental.Policy('mixed_bfloat16') tf.keras.mixed_precision.experimental.set_policy(policy) 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=ds_info.features['image'].shape, 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 = 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=0.9, 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': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), 'test/negative_log_likelihood': tf.keras.metrics.Mean(), 'test/accuracy': tf.keras.metrics.SparseCategoricalAccuracy(), 'test/ece': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), } if FLAGS.corruptions_interval > 0: corrupt_metrics = {} for intensity in range(1, max_intensity + 1): 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)] = ( um.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.""" if FLAGS.forget_mixup: images, labels, idx = inputs else: images, labels = inputs 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 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'].update_state(labels, probs) metrics['train/loss'].update_state(loss) metrics['train/negative_log_likelihood'].update_state( negative_log_likelihood) metrics['train/accuracy'].update_state(labels, logits) if FLAGS.forget_mixup: train_predictions = tf.argmax(probs, -1) labels = tf.cast(labels, train_predictions.dtype) # For each ensemble member (1 here), we accumulate the accuracy counts. accuracy_counts = tf.cast( tf.reshape((train_predictions == labels), [1, -1]), tf.float32) return accuracy_counts, idx if FLAGS.forget_mixup: return strategy.run(step_fn, args=(next(iterator), )) else: strategy.run(step_fn, args=(next(iterator), )) @tf.function def test_step(iterator, dataset_name): """Evaluation StepFn.""" def step_fn(inputs): """Per-Replica StepFn.""" images, labels = inputs logits_list = [] for _ in range(FLAGS.num_dropout_samples): logits = model(images, training=False) if FLAGS.use_bfloat16: logits = tf.cast(logits, tf.float32) 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, labels.shape[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['test/negative_log_likelihood'].update_state( negative_log_likelihood) metrics['test/accuracy'].update_state(labels, probs) metrics['test/ece'].update_state(labels, probs) elif dataset_name != 'validation': 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)].update_state(labels, probs) if dataset_name == 'validation': return tf.reshape(probs, [1, -1, num_classes]), labels if dataset_name == 'validation': return strategy.run(step_fn, args=(next(iterator), )) else: strategy.run(step_fn, args=(next(iterator), )) metrics.update({'test/ms_per_example': tf.keras.metrics.Mean()}) train_iterator = iter(train_dataset) forget_counts_history = [] start_time = time.time() for epoch in range(initial_epoch, FLAGS.train_epochs): logging.info('Starting to run epoch: %s', epoch) acc_counts_list = [] idx_list = [] for step in range(steps_per_epoch): if FLAGS.forget_mixup: temp_accuracy_counts, temp_idx = train_step(train_iterator) acc_counts_list.append(temp_accuracy_counts) idx_list.append(temp_idx) else: train_step(train_iterator) 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) # Only one of the forget_mixup and adaptive_mixup can be true. if FLAGS.forget_mixup: current_acc = [ tf.concat(list(acc_counts_list[i].values), axis=1) for i in range(len(acc_counts_list)) ] total_idx = [ tf.concat(list(idx_list[i].values), axis=0) for i in range(len(idx_list)) ] current_acc = tf.cast(tf.concat(current_acc, axis=1), tf.int32) total_idx = tf.concat(total_idx, axis=0) current_forget_path = os.path.join(FLAGS.output_dir, 'forget_counts.npy') last_acc_path = os.path.join(FLAGS.output_dir, 'last_acc.npy') if epoch == 0: forget_counts = tf.zeros([1, num_train_examples], dtype=tf.int32) last_acc = tf.zeros([1, num_train_examples], dtype=tf.int32) else: if 'last_acc' not in locals(): with tf.io.gfile.GFile(last_acc_path, 'rb') as f: last_acc = np.load(f) last_acc = tf.cast(tf.convert_to_tensor(last_acc), tf.int32) if 'forget_counts' not in locals(): with tf.io.gfile.GFile(current_forget_path, 'rb') as f: forget_counts = np.load(f) forget_counts = tf.cast( tf.convert_to_tensor(forget_counts), tf.int32) selected_last_acc = tf.gather(last_acc, total_idx, axis=1) forget_this_epoch = tf.cast(current_acc < selected_last_acc, tf.int32) forget_this_epoch = tf.transpose(forget_this_epoch) target_shape = tf.constant([num_train_examples, 1]) current_forget_counts = tf.scatter_nd( tf.reshape(total_idx, [-1, 1]), forget_this_epoch, target_shape) current_forget_counts = tf.transpose(current_forget_counts) acc_this_epoch = tf.transpose(current_acc) last_acc = tf.scatter_nd(tf.reshape(total_idx, [-1, 1]), acc_this_epoch, target_shape) # This is lower bound of true acc. last_acc = tf.transpose(last_acc) # TODO(ywenxu): We count the dropped examples as forget. Fix this later. forget_counts += current_forget_counts forget_counts_history.append(forget_counts) logging.info('forgetting counts') logging.info(tf.stack(forget_counts_history, 0)) with tf.io.gfile.GFile( os.path.join(FLAGS.output_dir, 'forget_counts_history.npy'), 'wb') as f: np.save(f, tf.stack(forget_counts_history, 0).numpy()) with tf.io.gfile.GFile(current_forget_path, 'wb') as f: np.save(f, forget_counts.numpy()) with tf.io.gfile.GFile(last_acc_path, 'wb') as f: np.save(f, last_acc.numpy()) aug_params['forget_counts_dir'] = current_forget_path train_input_fn = data_utils.load_input_fn( split=tfds.Split.TRAIN, name=FLAGS.dataset, batch_size=batch_size, use_bfloat16=FLAGS.use_bfloat16, validation_set=FLAGS.validation, aug_params=aug_params) train_dataset = strategy.experimental_distribute_dataset( train_input_fn()) train_iterator = iter(train_dataset) elif FLAGS.adaptive_mixup: val_iterator = iter(val_dataset) logging.info('Testing on validation dataset') predictions_list = [] labels_list = [] for step in range(steps_per_val): temp_predictions, temp_labels = test_step( val_iterator, 'validation') predictions_list.append(temp_predictions) labels_list.append(temp_labels) predictions = [ tf.concat(list(predictions_list[i].values), axis=1) for i in range(len(predictions_list)) ] labels = [ tf.concat(list(labels_list[i].values), axis=0) for i in range(len(labels_list)) ] predictions = tf.concat(predictions, axis=1) labels = tf.cast(tf.concat(labels, axis=0), tf.int64) def compute_acc_conf(preds, label, focus_class): class_preds = tf.boolean_mask(preds, label == focus_class, axis=1) class_pred_labels = tf.argmax(class_preds, axis=-1) confidence = tf.reduce_mean( tf.reduce_max(class_preds, axis=-1), -1) accuracy = tf.reduce_mean(tf.cast( class_pred_labels == focus_class, tf.float32), axis=-1) return accuracy - confidence calibration_per_class = [ compute_acc_conf(predictions, labels, i) for i in range(num_classes) ] calibration_per_class = tf.stack(calibration_per_class, axis=1) logging.info('calibration per class') logging.info(calibration_per_class) mixup_coeff = tf.where(calibration_per_class > 0, 1.0, FLAGS.mixup_alpha) mixup_coeff = tf.clip_by_value(mixup_coeff, 0, 1) logging.info('mixup coeff') logging.info(mixup_coeff) aug_params['mixup_coeff'] = mixup_coeff train_input_fn = data_utils.load_input_fn( split=tfds.Split.TRAIN, name=FLAGS.dataset, batch_size=batch_size, use_bfloat16=FLAGS.use_bfloat16, validation_set=True, aug_params=aug_params) train_dataset = strategy.experimental_distribute_dataset( train_input_fn()) train_iterator = iter(train_dataset) 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) for step in range(steps_per_eval): if step % 20 == 0: logging.info('Starting to run eval step %s of epoch: %s', step, epoch) test_start_time = time.time() test_step(test_iterator, dataset_name) 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) corrupt_results = {} if (FLAGS.corruptions_interval > 0 and (epoch + 1) % FLAGS.corruptions_interval == 0): corrupt_results = utils.aggregate_corrupt_metrics( corrupt_metrics, corruption_types, max_intensity) 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) 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.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)
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) ind_dataset_builder = ub.datasets.ClincIntentDetectionDataset( split='test', data_dir=FLAGS.data_dir, data_mode='ind') ood_dataset_builder = ub.datasets.ClincIntentDetectionDataset( split='test', data_dir=FLAGS.data_dir, data_mode='ood') all_dataset_builder = ub.datasets.ClincIntentDetectionDataset( split='test', data_dir=FLAGS.data_dir, data_mode='all') dataset_builders = { 'clean': ind_dataset_builder, 'ood': ood_dataset_builder, 'all': all_dataset_builder } ds_info = ind_dataset_builder.tfds_info feature_size = ds_info.metadata['feature_size'] # num_classes is number of valid intents plus out-of-scope intent num_classes = ds_info.features['intent_label'].num_classes + 1 # vocab_size is total number of valid tokens plus the out-of-vocabulary token. vocab_size = ind_dataset_builder.tokenizer.num_words + 1 batch_size = FLAGS.per_core_batch_size * FLAGS.num_cores test_datasets = {} steps_per_eval = {} for dataset_name, dataset_builder in dataset_builders.items(): test_datasets[dataset_name] = dataset_builder.load( batch_size=batch_size) steps_per_eval[ dataset_name] = dataset_builder.num_examples // batch_size bert_config_dir, _ = sngp.resolve_bert_ckpt_and_config_dir( FLAGS.bert_dir, FLAGS.bert_config_dir, FLAGS.bert_ckpt_dir) bert_config = bert_utils.create_config(bert_config_dir) gp_layer_kwargs = dict(num_inducing=FLAGS.gp_hidden_dim, gp_kernel_scale=FLAGS.gp_scale, gp_output_bias=FLAGS.gp_bias, normalize_input=FLAGS.gp_input_normalization, gp_cov_momentum=FLAGS.gp_cov_discount_factor, gp_cov_ridge_penalty=FLAGS.gp_cov_ridge_penalty) spec_norm_kwargs = dict(iteration=FLAGS.spec_norm_iteration, norm_multiplier=FLAGS.spec_norm_bound) model, bert_encoder = ub.models.SngpBertBuilder( num_classes=num_classes, bert_config=bert_config, gp_layer_kwargs=gp_layer_kwargs, spec_norm_kwargs=spec_norm_kwargs, use_gp_layer=FLAGS.use_gp_layer, use_spec_norm_att=FLAGS.use_spec_norm_att, use_spec_norm_ffn=FLAGS.use_spec_norm_ffn, use_layer_norm_att=FLAGS.use_layer_norm_att, use_layer_norm_ffn=FLAGS.use_layer_norm_ffn, use_spec_norm_plr=FLAGS.use_spec_norm_plr) 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')) 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) # 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, member=m) filename = os.path.join(FLAGS.output_dir, filename) if not tf.io.gfile.exists(filename): logits_list = [] test_iterator = iter(test_dataset) for _ in range(steps_per_eval[name]): inputs = next(test_iterator) features, _ = bert_utils.create_feature_and_label( inputs, feature_size) logits, covmat = model(features, training=False) logits = ed.layers.utils.mean_field_logits( logits, covmat, mean_field_factor=FLAGS.gp_mean_field_factor) logits_list.append(logits) logits_list = tf.concat(logits_list, axis=0) with tf.io.gfile.GFile(filename, 'w') as f: np.save(f, logits_list.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': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), } for dataset_name, test_dataset in test_datasets.items(): if dataset_name != 'clean': metrics.update({ 'test/nll_{}'.format(dataset_name): tf.keras.metrics.Mean(), 'test/accuracy_{}'.format(dataset_name): tf.keras.metrics.SparseCategoricalAccuracy(), 'test/ece_{}'.format(dataset_name): um.ExpectedCalibrationError(num_bins=FLAGS.num_bins) }) # Finally, define OOD metrics for the combined IND and OOD dataset. metrics.update({ 'test/auroc_all': tf.keras.metrics.AUC(curve='ROC'), 'test/auprc_all': tf.keras.metrics.AUC(curve='PR') }) # 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, 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[name]): inputs = next(test_iterator) _, labels = bert_utils.create_feature_and_label( inputs, feature_size) logits = logits_dataset[:, (step * batch_size):((step + 1) * batch_size)] labels = tf.cast(labels, tf.int32) negative_log_likelihood = um.ensemble_cross_entropy(labels, logits) per_probs = tf.nn.softmax(logits) probs = tf.reduce_mean(per_probs, axis=0) if name == 'clean': gibbs_ce = um.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: metrics['test/nll_{}'.format(name)].update_state( negative_log_likelihood) metrics['test/accuracy_{}'.format(name)].update_state( labels, probs) metrics['test/ece_{}'.format(name)].update_state(labels, probs) if dataset_name == 'all': ood_labels = tf.cast(labels == 150, labels.dtype) ood_probs = 1. - tf.reduce_max(probs, axis=-1) metrics['test/auroc_{}'.format(dataset_name)].update_state( ood_labels, ood_probs) metrics['test/auprc_{}'.format(dataset_name)].update_state( ood_labels, ood_probs) message = ( '{:.1%} completion for evaluation: dataset {:d}/{:d}'.format( (n + 1) / num_datasets, n + 1, num_datasets)) logging.info(message) total_results = {name: metric.result() for name, metric in metrics.items()} logging.info('Metrics: %s', total_results)
def load_pivot_df_dict_googlefig( hdf5file='metrics_googlefig/cifar_model_predictions.hdf5', clean_key='test', instance=0): ''' load pivoted df_dict from google paper ''' file = h5py.File(hdf5file, 'r') corruptions = [ 'corrupt-static-brightness-1', 'corrupt-static-brightness-2', 'corrupt-static-brightness-3', 'corrupt-static-brightness-4', 'corrupt-static-brightness-5', 'corrupt-static-contrast-1', 'corrupt-static-contrast-2', 'corrupt-static-contrast-3', 'corrupt-static-contrast-4', 'corrupt-static-contrast-5', 'corrupt-static-defocus_blur-1', 'corrupt-static-defocus_blur-2', 'corrupt-static-defocus_blur-3', 'corrupt-static-defocus_blur-4', 'corrupt-static-defocus_blur-5', 'corrupt-static-elastic_transform-1', 'corrupt-static-elastic_transform-2', 'corrupt-static-elastic_transform-3', 'corrupt-static-elastic_transform-4', 'corrupt-static-elastic_transform-5', 'corrupt-static-fog-1', 'corrupt-static-fog-2', 'corrupt-static-fog-3', 'corrupt-static-fog-4', 'corrupt-static-fog-5', 'corrupt-static-frost-1', 'corrupt-static-frost-2', 'corrupt-static-frost-3', 'corrupt-static-frost-4', 'corrupt-static-frost-5', 'corrupt-static-gaussian_blur-1', 'corrupt-static-gaussian_blur-2', 'corrupt-static-gaussian_blur-3', 'corrupt-static-gaussian_blur-4', 'corrupt-static-gaussian_blur-5', 'corrupt-static-gaussian_noise-1', 'corrupt-static-gaussian_noise-2', 'corrupt-static-gaussian_noise-3', 'corrupt-static-gaussian_noise-4', 'corrupt-static-gaussian_noise-5', 'corrupt-static-glass_blur-1', 'corrupt-static-glass_blur-2', 'corrupt-static-glass_blur-3', 'corrupt-static-glass_blur-4', 'corrupt-static-glass_blur-5', 'corrupt-static-impulse_noise-1', 'corrupt-static-impulse_noise-2', 'corrupt-static-impulse_noise-3', 'corrupt-static-impulse_noise-4', 'corrupt-static-impulse_noise-5', 'corrupt-static-pixelate-1', 'corrupt-static-pixelate-2', 'corrupt-static-pixelate-3', 'corrupt-static-pixelate-4', 'corrupt-static-pixelate-5', 'corrupt-static-saturate-1', 'corrupt-static-saturate-2', 'corrupt-static-saturate-3', 'corrupt-static-saturate-4', 'corrupt-static-saturate-5', 'corrupt-static-shot_noise-1', 'corrupt-static-shot_noise-2', 'corrupt-static-shot_noise-3', 'corrupt-static-shot_noise-4', 'corrupt-static-shot_noise-5', 'corrupt-static-spatter-1', 'corrupt-static-spatter-2', 'corrupt-static-spatter-3', 'corrupt-static-spatter-4', 'corrupt-static-spatter-5', 'corrupt-static-speckle_noise-1', 'corrupt-static-speckle_noise-2', 'corrupt-static-speckle_noise-3', 'corrupt-static-speckle_noise-4', 'corrupt-static-speckle_noise-5', 'corrupt-static-zoom_blur-1', 'corrupt-static-zoom_blur-2', 'corrupt-static-zoom_blur-3', 'corrupt-static-zoom_blur-4', 'corrupt-static-zoom_blur-5' ] clean = [clean_key] corruptions += clean gfig_dict = dict() metrics = [ tf.keras.metrics.SparseCategoricalAccuracy(name='acc'), um.ExpectedCalibrationError(num_bins=15, name='ece'), nll(name='nll'), BrierScore(name='brier') ] for model_name, model_data in file.items(): #print(model_name) df = pd.DataFrame(columns=['acc', 'brier', 'ece', 'nll']) df.columns.name = 'metrics' df.index.name = 'dataset' for key in model_data.keys(): if key in corruptions: #model_data[key] y = model_data[key]['labels'] p = model_data[key]['probs'] if 'roll' in key: key = 'roll' for metric in metrics: metric.update_state(y[instance, :].astype(np.int32), p[instance, :].astype(np.float32)) d0 = { 'acc': metrics[0].result().numpy(), 'ece': metrics[1].result().numpy(), 'nll': metrics[2].result().numpy(), 'brier': metrics[3].result().numpy() } record = pd.Series(d0, index=df.columns, name=key) df = df.append(record) for metric in metrics: metric.reset_states() gfig_dict[model_name] = df.copy() return gfig_dict
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.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 dataset_input_fn = utils.load_input_fn( split=tfds.Split.TEST, name=FLAGS.dataset, batch_size=FLAGS.per_core_batch_size, use_bfloat16=FLAGS.use_bfloat16) test_datasets = {'clean': dataset_input_fn()} corruption_types, max_intensity = utils.load_corrupted_test_info( FLAGS.dataset) for name in corruption_types: for intensity in range(1, max_intensity + 1): dataset_name = '{0}_{1}'.format(name, intensity) if FLAGS.dataset == 'cifar10': load_c_dataset = utils.load_cifar10_c_input_fn else: load_c_dataset = functools.partial( utils.load_cifar100_c_input_fn, path=FLAGS.cifar100_c_path) corrupted_input_fn = load_c_dataset( corruption_name=name, corruption_intensity=intensity, batch_size=FLAGS.per_core_batch_size, use_bfloat16=FLAGS.use_bfloat16) test_datasets[dataset_name] = corrupted_input_fn() 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 = 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) # 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, 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) # pytype: disable=attribute-error 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': um.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)] = ( um.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()) test_diversity = { 'test/disagreement': tf.keras.metrics.Mean(), 'test/average_kl': tf.keras.metrics.Mean(), 'test/cosine_similarity': tf.keras.metrics.Mean(), } metrics.update(test_diversity) # 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, 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) # pytype: disable=attribute-error logits = logits_dataset[:, (step*batch_size):((step+1)*batch_size)] labels = tf.cast(labels, tf.int32) negative_log_likelihood = um.ensemble_cross_entropy(labels, logits) per_probs = tf.nn.softmax(logits) probs = tf.reduce_mean(per_probs, axis=0) if name == 'clean': gibbs_ce = um.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) 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) diversity_results = um.average_pairwise_diversity( per_probs, ensemble_size) for k, v in diversity_results.items(): test_diversity['test/' + k].update_state(v) 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) 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 # unused arg tf.io.gfile.makedirs(FLAGS.output_dir) logging.info('Saving checkpoints at %s', FLAGS.output_dir) tf.random.set_seed(FLAGS.seed) train_batch_size = (FLAGS.per_core_batch_size * FLAGS.num_cores // FLAGS.batch_repetitions) test_batch_size = FLAGS.per_core_batch_size * FLAGS.num_cores steps_per_epoch = APPROX_IMAGENET_TRAIN_IMAGES // train_batch_size steps_per_eval = IMAGENET_VALIDATION_IMAGES // test_batch_size 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) train_builder = ub.datasets.ImageNetDataset( split=tfds.Split.TRAIN, use_bfloat16=FLAGS.use_bfloat16) train_dataset = train_builder.load( batch_size=train_batch_size, strategy=strategy) test_builder = ub.datasets.ImageNetDataset( split=tfds.Split.TEST, use_bfloat16=FLAGS.use_bfloat16) test_dataset = test_builder.load( batch_size=test_batch_size, strategy=strategy) if FLAGS.use_bfloat16: tf.keras.mixed_precision.set_global_policy('mixed_bfloat16') with strategy.scope(): logging.info('Building Keras ResNet-50 model') model = imagenet_model.resnet50( input_shape=(FLAGS.ensemble_size, 224, 224, 3), num_classes=NUM_CLASSES, ensemble_size=FLAGS.ensemble_size, width_multiplier=FLAGS.width_multiplier) 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()) # Scale learning rate and decay epochs by vanilla settings. base_lr = FLAGS.base_learning_rate * train_batch_size / 256 decay_epochs = [ (FLAGS.train_epochs * int(FLAGS.lr_decay_epochs[0])) // 90, (FLAGS.train_epochs * int(FLAGS.lr_decay_epochs[1])) // 90, (FLAGS.train_epochs * int(FLAGS.lr_decay_epochs[2])) // 90, ] learning_rate = ub.schedules.WarmUpPiecewiseConstantSchedule( steps_per_epoch=steps_per_epoch, base_learning_rate=base_lr, decay_ratio=0.1, decay_epochs=decay_epochs, warmup_epochs=FLAGS.lr_warmup_epochs) optimizer = tf.keras.optimizers.SGD(learning_rate=learning_rate, momentum=0.9, 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': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), 'test/negative_log_likelihood': tf.keras.metrics.Mean(), 'test/accuracy': tf.keras.metrics.SparseCategoricalAccuracy(), 'test/ece': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), } for i in range(FLAGS.ensemble_size): metrics['test/nll_member_{}'.format(i)] = tf.keras.metrics.Mean() metrics['test/accuracy_member_{}'.format(i)] = ( tf.keras.metrics.SparseCategoricalAccuracy()) test_diversity = { 'test/disagreement': tf.keras.metrics.Mean(), 'test/average_kl': tf.keras.metrics.Mean(), 'test/cosine_similarity': tf.keras.metrics.Mean(), } logging.info('Finished building Keras ResNet-50 model') 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 summary_writer = tf.summary.create_file_writer( os.path.join(FLAGS.output_dir, 'summaries')) @tf.function def train_step(iterator): """Training StepFn.""" def step_fn(inputs): """Per-Replica StepFn.""" images, labels = inputs batch_size = tf.shape(images)[0] main_shuffle = tf.random.shuffle(tf.tile( tf.range(batch_size), [FLAGS.batch_repetitions])) to_shuffle = tf.cast(tf.cast(tf.shape(main_shuffle)[0], tf.float32) * (1. - FLAGS.input_repetition_probability), tf.int32) shuffle_indices = [ tf.concat([tf.random.shuffle(main_shuffle[:to_shuffle]), main_shuffle[to_shuffle:]], axis=0) for _ in range(FLAGS.ensemble_size)] images = tf.stack([tf.gather(images, indices, axis=0) for indices in shuffle_indices], axis=1) labels = tf.stack([tf.gather(labels, indices, axis=0) for indices in shuffle_indices], axis=1) with tf.GradientTape() as tape: logits = model(images, training=True) if FLAGS.use_bfloat16: logits = tf.cast(logits, tf.float32) negative_log_likelihood = tf.reduce_mean(tf.reduce_sum( tf.keras.losses.sparse_categorical_crossentropy(labels, logits, from_logits=True), axis=1)) filtered_variables = [] for var in model.trainable_variables: # Apply l2 on the weights. This excludes BN parameters and biases, but # pay caution to their naming scheme. if 'kernel' in var.name or 'bias' in var.name: filtered_variables.append(tf.reshape(var, (-1,))) l2_loss = FLAGS.l2 * 2 * tf.nn.l2_loss( tf.concat(filtered_variables, axis=0)) # Scale the loss given the TPUStrategy will reduce sum all gradients. loss = negative_log_likelihood + l2_loss 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(tf.reshape(logits, [-1, NUM_CLASSES])) flat_labels = tf.reshape(labels, [-1]) metrics['train/ece'].update_state(flat_labels, probs) metrics['train/loss'].update_state(loss) metrics['train/negative_log_likelihood'].update_state( negative_log_likelihood) metrics['train/accuracy'].update_state(flat_labels, probs) strategy.run(step_fn, args=(next(iterator),)) @tf.function def test_step(iterator): """Evaluation StepFn.""" def step_fn(inputs): """Per-Replica StepFn.""" images, labels = inputs images = tf.tile( tf.expand_dims(images, 1), [1, FLAGS.ensemble_size, 1, 1, 1]) logits = model(images, training=False) if FLAGS.use_bfloat16: logits = tf.cast(logits, tf.float32) probs = tf.nn.softmax(logits) per_probs = tf.transpose(probs, perm=[1, 0, 2]) diversity_results = um.average_pairwise_diversity( per_probs, FLAGS.ensemble_size) for k, v in diversity_results.items(): test_diversity['test/' + k].update_state(v) for i in range(FLAGS.ensemble_size): member_probs = 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) # Negative log marginal likelihood computed in a numerically-stable way. labels_tiled = tf.tile( tf.expand_dims(labels, 1), [1, FLAGS.ensemble_size]) log_likelihoods = -tf.keras.losses.sparse_categorical_crossentropy( labels_tiled, logits, from_logits=True) negative_log_likelihood = tf.reduce_mean( -tf.reduce_logsumexp(log_likelihoods, axis=[1]) + tf.math.log(float(FLAGS.ensemble_size))) probs = tf.math.reduce_mean(probs, axis=1) # marginalize metrics['test/negative_log_likelihood'].update_state( negative_log_likelihood) metrics['test/accuracy'].update_state(labels, probs) metrics['test/ece'].update_state(labels, probs) 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) for step in range(steps_per_epoch): train_step(train_iterator) 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) test_iterator = iter(test_dataset) for step in range(steps_per_eval): if step % 20 == 0: logging.info('Starting to run eval step %s of epoch: %s', step, epoch) test_start_time = time.time() test_step(test_iterator) ms_per_example = (time.time() - test_start_time) * 1e6 / test_batch_size metrics['test/ms_per_example'].update_state(ms_per_example) 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) for i in range(FLAGS.ensemble_size): logging.info('Member %d Test Loss: %.4f, Accuracy: %.2f%%', i, metrics['test/nll_member_{}'.format(i)].result(), metrics['test/accuracy_member_{}'.format(i)].result() * 100) total_metrics = metrics.copy() total_metrics.update(test_diversity) total_results = {name: metric.result() for name, metric in total_metrics.items()} with summary_writer.as_default(): for name, result in total_results.items(): tf.summary.scalar(name, result, step=epoch + 1) for _, metric in total_metrics.items(): 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_save_name = os.path.join(FLAGS.output_dir, 'model') model.save(final_save_name) logging.info('Saved model to %s', final_save_name)
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) 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 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, 'ensemble_size': 1, 'mixup_alpha': FLAGS.mixup_alpha, } validation_proportion = 1. - FLAGS.train_proportion use_validation_set = validation_proportion > 0. train_dataset = data_utils.load_dataset( split=tfds.Split.TRAIN, name=FLAGS.dataset, batch_size=batch_size, use_bfloat16=FLAGS.use_bfloat16, aug_params=aug_params, validation_set=use_validation_set, validation_proportion=validation_proportion) train_sample_size = ds_info.splits[ 'train'].num_examples * FLAGS.train_proportion val_sample_size = ds_info.splits['train'].num_examples - train_sample_size if use_validation_set: validation_dataset = data_utils.load_dataset( split=tfds.Split.VALIDATION, name=FLAGS.dataset, batch_size=batch_size, use_bfloat16=FLAGS.use_bfloat16, aug_params=aug_params, validation_set=use_validation_set, validation_proportion=validation_proportion) validation_dataset = strategy.experimental_distribute_dataset( validation_dataset) steps_per_val = steps_per_epoch = int(val_sample_size / batch_size) clean_test_dataset = utils.load_dataset(split=tfds.Split.TEST, name=FLAGS.dataset, batch_size=test_batch_size, use_bfloat16=FLAGS.use_bfloat16) train_sample_size = ds_info.splits[ 'train'].num_examples * FLAGS.train_proportion steps_per_epoch = int(train_sample_size / batch_size) steps_per_epoch = ds_info.splits['train'].num_examples // batch_size steps_per_eval = ds_info.splits['test'].num_examples // batch_size train_dataset = strategy.experimental_distribute_dataset(train_dataset) test_datasets = { 'clean': strategy.experimental_distribute_dataset(clean_test_dataset), } if FLAGS.corruptions_interval > 0: if FLAGS.dataset == 'cifar10': load_c_dataset = utils.load_cifar10_c else: load_c_dataset = functools.partial(utils.load_cifar100_c, path=FLAGS.cifar100_c_path) corruption_types, max_intensity = utils.load_corrupted_test_info( FLAGS.dataset) for corruption in corruption_types: for intensity in range(1, max_intensity + 1): dataset = load_c_dataset(corruption_name=corruption, corruption_intensity=intensity, batch_size=test_batch_size, use_bfloat16=FLAGS.use_bfloat16) test_datasets['{0}_{1}'.format(corruption, intensity)] = ( strategy.experimental_distribute_dataset(dataset)) if FLAGS.use_bfloat16: policy = tf.keras.mixed_precision.experimental.Policy('mixed_bfloat16') tf.keras.mixed_precision.experimental.set_policy(policy) 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=ds_info.features['image'].shape, 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 = utils.LearningRateSchedule( 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=0.9, 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': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), 'test/negative_log_likelihood': tf.keras.metrics.Mean(), 'test/accuracy': tf.keras.metrics.SparseCategoricalAccuracy(), 'test/ece': um.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': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), 'val/stddev': tf.keras.metrics.Mean(), }) if FLAGS.corruptions_interval > 0: corrupt_metrics = {} for intensity in range(1, max_intensity + 1): 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)] = ( um.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 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, step): """Training StepFn.""" def step_fn(inputs, step): """Per-Replica StepFn.""" images, labels = inputs if tf.equal(step, 0) and FLAGS.gp_cov_discount_factor < 0: # Resetting covaraince estimator at the begining of a new epoch. 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, 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'].update_state(labels, probs) 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): """Evaluation StepFn.""" def step_fn(inputs): """Per-Replica StepFn.""" images, labels = inputs logits_list = [] stddev_list = [] for _ in range(FLAGS.num_dropout_samples): logits = model(images, training=False) if isinstance(logits, tuple): # If model returns a tuple of (logits, covmat), extract both logits, covmat = logits else: covmat = tf.eye(FLAGS.per_core_batch_size) 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) 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) labels_broadcasted = tf.broadcast_to( labels, [FLAGS.num_dropout_samples, labels.shape[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['test/negative_log_likelihood'].update_state( negative_log_likelihood) metrics['test/accuracy'].update_state(labels, probs) metrics['test/ece'].update_state(labels, probs) 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'].update_state(labels, probs) metrics['val/stddev'].update_state(stddev) 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)].update_state(labels, probs) corrupt_metrics['test/stddev_{}'.format( dataset_name)].update_state(stddev) 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) 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 for step in range(steps_per_eval): if step % 20 == 0: logging.info('Starting to run eval step %s of epoch: %s', step, epoch) test_start_time = time.time() test_step(test_iterator, dataset_name) 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) corrupt_results = {} if (FLAGS.corruptions_interval > 0 and (epoch + 1) % FLAGS.corruptions_interval == 0): corrupt_results = utils.aggregate_corrupt_metrics( corrupt_metrics, corruption_types, max_intensity) 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) 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.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)
def run(trial_dir: str, flag_string: Optional[str]): """Run the experiment. Args: trial_dir: String to the dir to write checkpoints to and read them from. flag_string: Optional string used to record what flags the job was run with. """ tf.random.set_seed(FLAGS.seed) np.random.seed(FLAGS.seed) if not FLAGS.eval_frequency: FLAGS.eval_frequency = FLAGS.log_frequency if FLAGS.eval_frequency % FLAGS.log_frequency != 0: raise ValueError( 'log_frequency ({}) must evenly divide eval_frequency ' '({}).'.format(FLAGS.log_frequency, FLAGS.eval_frequency)) strategy = ub.strategy_utils.get_strategy(FLAGS.tpu, use_tpu=not FLAGS.use_cpu and not FLAGS.use_gpu) with strategy.scope(): _maybe_setup_trial_dir(strategy, trial_dir, flag_string) # TODO(znado): pass all dataset and model kwargs. train_dataset_builder = ub.datasets.get( dataset_name=FLAGS.dataset_name, split='train', validation_percent=FLAGS.validation_percent, shuffle_buffer_size=FLAGS.shuffle_buffer_size) if FLAGS.validation_percent > 0: validation_dataset_builder = ub.datasets.get( dataset_name=FLAGS.dataset_name, split='validation', validation_percent=FLAGS.validation_percent, shuffle_buffer_size=FLAGS.shuffle_buffer_size) else: validation_dataset_builder = None test_dataset_builder = ub.datasets.get( dataset_name=FLAGS.dataset_name, split='test', validation_percent=FLAGS.validation_percent, shuffle_buffer_size=FLAGS.shuffle_buffer_size) if FLAGS.use_spec_norm: logging.info('Use spectral normalization.') spec_norm_hparams = { 'spec_norm_bound': FLAGS.spec_norm_bound, 'spec_norm_iteration': FLAGS.spec_norm_iteration } else: spec_norm_hparams = None if FLAGS.use_gp_layer: logging.info('Use GP for output layer.') gp_layer_hparams = { '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_cov_discount_factor': FLAGS.gp_cov_discount_factor, 'gp_cov_ridge_penalty': FLAGS.gp_cov_ridge_penalty } else: gp_layer_hparams = None model = ub_smu_models.get( FLAGS.model_name, num_classes=FLAGS.num_classes, batch_size=FLAGS.batch_size, len_seqs=FLAGS.len_seqs, num_motifs=FLAGS.num_motifs, len_motifs=FLAGS.len_motifs, num_denses=FLAGS.num_denses, depth=FLAGS.wide_resnet_depth, width=FLAGS.wide_resnet_depth, l2_weight=FLAGS.l2_regularization, depth_multiplier=FLAGS.wide_resnet_width_multiplier, dropout_rate=FLAGS.dropout_rate, before_conv_dropout=FLAGS.before_conv_dropout, use_mc_dropout=FLAGS.use_mc_dropout, spec_norm_hparams=spec_norm_hparams, gp_layer_hparams=gp_layer_hparams) metrics = { 'accuracy': tf.keras.metrics.SparseCategoricalAccuracy(), 'brier_score': BrierScore(name='brier_score'), 'ece': um.ExpectedCalibrationError(num_bins=10, name='ece'), 'loss': tf.keras.metrics.SparseCategoricalCrossentropy(), } # Record all non-default hparams in tensorboard. hparams = _get_hparams() ood_dataset_builder = None ood_metrics = None if FLAGS.run_ood: if 'cifar' in FLAGS.dataset_name and FLAGS.ood_dataset_name == 'svhn': svhn_normalize_by_cifar = True else: svhn_normalize_by_cifar = False ood_dataset_builder_cls = ub.datasets.DATASETS[ FLAGS.ood_dataset_name] ood_dataset_builder_cls = ub.datasets.make_ood_dataset( ood_dataset_builder_cls) ood_dataset_builder = ood_dataset_builder_cls( in_distribution_dataset=test_dataset_builder, name=FLAGS.ood_dataset_name, split='test', validation_percent=FLAGS.validation_percent, normalize_by_cifar=svhn_normalize_by_cifar, data_mode='ood') _check_batch_replica_divisible(FLAGS.eval_batch_size, strategy) ood_metrics = { 'auroc': tf.keras.metrics.AUC(curve='ROC', summation_method='interpolation'), 'auprc': tf.keras.metrics.AUC(curve='PR', summation_method='interpolation') } aux_metrics = [ ('spec_at_sen', tf.keras.metrics.SpecificityAtSensitivity, FLAGS.sensitivity_thresholds), ('sen_at_spec', tf.keras.metrics.SensitivityAtSpecificity, FLAGS.specificity_thresholds), ('prec_at_rec', tf.keras.metrics.PrecisionAtRecall, FLAGS.recall_thresholds), ('rec_at_prec', tf.keras.metrics.RecallAtPrecision, FLAGS.precision_thresholds) ] for metric_name, metric_fn, threshold_vals in aux_metrics: vals = [float(x) for x in threshold_vals] thresholds = np.linspace(vals[0], vals[1], int(vals[2])) for thresh in thresholds: name = f'{metric_name}_{thresh:.2f}' ood_metrics[name] = metric_fn(thresh) if FLAGS.mode == 'eval': _check_batch_replica_divisible(FLAGS.eval_batch_size, strategy) eval_lib.run_eval_loop( validation_dataset_builder=validation_dataset_builder, test_dataset_builder=test_dataset_builder, batch_size=FLAGS.eval_batch_size, model=model, trial_dir=trial_dir, train_steps=FLAGS.train_steps, strategy=strategy, metrics=metrics, checkpoint_step=FLAGS.checkpoint_step, hparams=hparams, ood_dataset_builder=ood_dataset_builder, ood_metrics=ood_metrics) return if FLAGS.mode == 'train_and_eval': _check_batch_replica_divisible(FLAGS.eval_batch_size, strategy) steps_per_epoch = train_dataset_builder.num_examples // FLAGS.batch_size optimizer_kwargs = { k[len('optimizer_hparams_'):]: FLAGS[k].value for k in FLAGS if k.startswith('optimizer_hparams_') } optimizer_kwargs.update({ k[len('schedule_hparams_'):]: FLAGS[k].value for k in FLAGS if k.startswith('schedule_hparams_') }) optimizer = ub.optimizers.get( optimizer_name=FLAGS.optimizer, learning_rate_schedule=FLAGS.learning_rate_schedule, learning_rate=FLAGS.learning_rate, weight_decay=FLAGS.weight_decay, steps_per_epoch=steps_per_epoch, model=model, **optimizer_kwargs) train_lib.run_train_loop( train_dataset_builder=train_dataset_builder, validation_dataset_builder=validation_dataset_builder, test_dataset_builder=test_dataset_builder, batch_size=FLAGS.batch_size, eval_batch_size=FLAGS.eval_batch_size, model=model, optimizer=optimizer, eval_frequency=FLAGS.eval_frequency, log_frequency=FLAGS.log_frequency, trial_dir=trial_dir, train_steps=FLAGS.train_steps, mode=FLAGS.mode, strategy=strategy, metrics=metrics, hparams=hparams, ood_dataset_builder=ood_dataset_builder, ood_metrics=ood_metrics, focal_loss_gamma=FLAGS.focal_loss_gamma)
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) batch_size = FLAGS.per_core_batch_size * FLAGS.num_cores steps_per_epoch = APPROX_IMAGENET_TRAIN_IMAGES // batch_size steps_per_eval = IMAGENET_VALIDATION_IMAGES // batch_size 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) enable_mixup = (FLAGS.mixup_alpha > 0.0) mixup_params = { 'mixup_alpha': FLAGS.mixup_alpha, 'adaptive_mixup': False, 'same_mix_weight_per_batch': FLAGS.same_mix_weight_per_batch, 'use_random_shuffling': FLAGS.use_random_shuffling, 'use_truncated_beta': FLAGS.use_truncated_beta } train_builder = utils.ImageNetInput( data_dir=FLAGS.data_dir, use_bfloat16=FLAGS.use_bfloat16, one_hot=True, mixup_params=mixup_params) test_builder = utils.ImageNetInput( data_dir=FLAGS.data_dir, use_bfloat16=FLAGS.use_bfloat16) train_dataset = train_builder.as_dataset( split=tfds.Split.TRAIN, batch_size=batch_size) test_dataset = test_builder.as_dataset( split=tfds.Split.TEST, batch_size=batch_size) train_dataset = strategy.experimental_distribute_dataset(train_dataset) test_dataset = strategy.experimental_distribute_dataset(test_dataset) if enable_mixup: mean_theta = mean_truncated_beta_distribution(FLAGS.mixup_alpha) # Train set to compute the means of the images and of the (one-hot) labels imagenet_train_no_mixup = utils.ImageNetInput( data_dir=FLAGS.data_dir, use_bfloat16=FLAGS.use_bfloat16, one_hot=True) imagenet_train_no_mixup = imagenet_train_no_mixup.as_dataset( split=tfds.Split.TRAIN, batch_size=batch_size) tr_data_no_mixup = strategy.experimental_distribute_dataset( imagenet_train_no_mixup) if FLAGS.use_bfloat16: policy = tf.keras.mixed_precision.experimental.Policy('mixed_bfloat16') tf.keras.mixed_precision.experimental.set_policy(policy) with strategy.scope(): if enable_mixup: # Variables used to track the means of the images and the (one-hot) labels count = tf.Variable(tf.zeros((1,), dtype=tf.float32)) mean_images = tf.Variable(tf.zeros(IMAGE_SHAPE, dtype=tf.float32)) mean_labels = tf.Variable(tf.zeros((NUM_CLASSES,), dtype=tf.float32)) logging.info('Building Keras ResNet-50 model') model = ub.models.resnet50_deterministic(input_shape=IMAGE_SHAPE, 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()) # Scale learning rate and decay epochs by vanilla settings. base_lr = FLAGS.base_learning_rate * batch_size / 256 learning_rate = utils.LearningRateSchedule(steps_per_epoch, base_lr, FLAGS.train_epochs, _LR_SCHEDULE) optimizer = tf.keras.optimizers.SGD(learning_rate=learning_rate, momentum=0.9, 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': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), 'test/negative_log_likelihood': tf.keras.metrics.Mean(), 'test/accuracy': tf.keras.metrics.SparseCategoricalAccuracy(), 'test/ece': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), } logging.info('Finished building Keras ResNet-50 model') if enable_mixup: # With mixup enabled, we log the predictions with the rescaling from [2] metrics['test/negative_log_likelihood+rescaling'] = (tf.keras.metrics .Mean()) metrics['test/accuracy+rescaling'] = (tf.keras.metrics .SparseCategoricalAccuracy()) metrics['test/ece+rescaling'] = um.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 summary_writer = tf.summary.create_file_writer( os.path.join(FLAGS.output_dir, 'summaries')) @tf.function def moving_average_step(iterator): """Training StepFn to compute the means of the images and labels.""" def step_fn_labels(labels): return tf.reduce_mean(labels, axis=0) def step_fn_images(images): return tf.reduce_mean(tf.cast(images, tf.float32), axis=0) new_count = count + 1. count.assign(new_count) images, labels = next(iterator) per_replica_means = strategy.run(step_fn_labels, args=(labels,)) cr_replica_means = strategy.reduce('mean', per_replica_means, axis=0) mean_labels.assign(cr_replica_means/count + (count-1.)/count * mean_labels) per_replica_means = strategy.run(step_fn_images, args=(images,)) cr_replica_means = strategy.reduce('mean', per_replica_means, axis=0) mean_images.assign(cr_replica_means/count + (count-1.)/count * mean_images) @tf.function def train_step(iterator): """Training StepFn.""" def step_fn(inputs): """Per-Replica StepFn.""" images, labels = inputs with tf.GradientTape() as tape: logits = model(images, training=True) if FLAGS.use_bfloat16: logits = tf.cast(logits, tf.float32) negative_log_likelihood = tf.reduce_mean( tf.keras.losses.categorical_crossentropy( labels, logits, from_logits=True)) filtered_variables = [] for var in model.trainable_variables: # Apply l2 on the weights. This excludes BN parameters and biases, but # pay caution to their naming scheme. if 'kernel' in var.name or 'bias' in var.name: filtered_variables.append(tf.reshape(var, (-1,))) l2_loss = FLAGS.l2 * 2 * tf.nn.l2_loss( tf.concat(filtered_variables, axis=0)) # Scale the loss given the TPUStrategy will reduce sum all gradients. loss = negative_log_likelihood + l2_loss 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) # We go back from one-hot labels to integers labels = tf.argmax(labels, axis=-1) metrics['train/ece'].update_state(labels, probs) 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),)) @tf.function def update_test_metrics(labels, logits, metric_suffix=''): negative_log_likelihood = tf.reduce_mean( tf.keras.losses.sparse_categorical_crossentropy( labels, logits, from_logits=True)) probs = tf.nn.softmax(logits) metrics['test/negative_log_likelihood' + metric_suffix].update_state( negative_log_likelihood) metrics['test/accuracy' + metric_suffix].update_state(labels, probs) metrics['test/ece' + metric_suffix].update_state(labels, probs) @tf.function def test_step(iterator): """Evaluation StepFn.""" def step_fn(inputs): """Per-Replica StepFn.""" images, labels = inputs logits = model(images, training=False) if FLAGS.use_bfloat16: logits = tf.cast(logits, tf.float32) update_test_metrics(labels, logits) # Rescaling logic in Eq.(15) from [2] if enable_mixup: images *= mean_theta images += (1.-mean_theta) * tf.cast(mean_images, images.dtype) scaled_logits = model(images, training=False) if FLAGS.use_bfloat16: scaled_logits = tf.cast(scaled_logits, tf.float32) scaled_logits *= 1./mean_theta scaled_logits += (1.-1./mean_theta) * tf.cast(mean_labels, logits.dtype) update_test_metrics(labels, scaled_logits, '+rescaling') strategy.run(step_fn, args=(next(iterator),)) metrics.update({'test/ms_per_example': tf.keras.metrics.Mean()}) if enable_mixup: logging.info('Starting to compute the means of labels and images') tr_iterator_no_mixup = iter(tr_data_no_mixup) for step in range(steps_per_epoch): moving_average_step(tr_iterator_no_mixup) # Save stats required by the mixup rescaling [2] for subsequent predictions mixup_rescaling_stats = { 'mean_labels': mean_labels.numpy(), 'mean_images': mean_images.numpy(), 'mean_theta': mean_theta } output_dir = os.path.join(FLAGS.output_dir, 'mixup_rescaling_stats.npz') with tf.io.gfile.GFile(output_dir, 'wb') as f: np.save(f, list(mixup_rescaling_stats.items())) logging.info('Finished to compute the means of labels and images') 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) for step in range(steps_per_epoch): train_step(train_iterator) 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) test_iterator = iter(test_dataset) for step in range(steps_per_eval): if step % 20 == 0: logging.info('Starting to run eval step %s of epoch: %s', step, epoch) test_start_time = time.time() test_step(test_iterator) ms_per_example = (time.time() - test_start_time) * 1e6 / batch_size metrics['test/ms_per_example'].update_state(ms_per_example) 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) if enable_mixup: logging.info( 'Test NLL (+ rescaling): %.4f, Accuracy (+ rescaling): %.2f%%', metrics['test/negative_log_likelihood+rescaling'].result(), metrics['test/accuracy+rescaling'].result() * 100) total_results = {name: metric.result() for name, metric in metrics.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.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_save_name = os.path.join(FLAGS.output_dir, 'model') model.save(final_save_name) logging.info('Saved model to %s', final_save_name)
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) per_core_batch_size = FLAGS.per_core_batch_size // FLAGS.ensemble_size batch_size = per_core_batch_size * FLAGS.num_cores steps_per_epoch = APPROX_IMAGENET_TRAIN_IMAGES // batch_size steps_per_eval = IMAGENET_VALIDATION_IMAGES // batch_size 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) mixup_params = { 'ensemble_size': FLAGS.ensemble_size, 'mixup_alpha': FLAGS.mixup_alpha, 'adaptive_mixup': FLAGS.adaptive_mixup, 'num_classes': NUM_CLASSES, } train_builder = utils.ImageNetInput(data_dir=FLAGS.data_dir, one_hot=(FLAGS.mixup_alpha > 0), use_bfloat16=FLAGS.use_bfloat16, mixup_params=mixup_params, ensemble_size=FLAGS.ensemble_size) test_builder = utils.ImageNetInput(data_dir=FLAGS.data_dir, use_bfloat16=FLAGS.use_bfloat16) train_dataset = train_builder.as_dataset(split=tfds.Split.TRAIN, batch_size=batch_size) clean_test_dataset = test_builder.as_dataset(split=tfds.Split.TEST, batch_size=batch_size) train_dataset = strategy.experimental_distribute_dataset(train_dataset) test_datasets = { 'clean': strategy.experimental_distribute_dataset(clean_test_dataset) } if FLAGS.adaptive_mixup: imagenet_confidence_dataset = test_builder.as_dataset( split=tfds.Split.VALIDATION, batch_size=FLAGS.per_core_batch_size * FLAGS.num_cores) imagenet_confidence_dataset = ( strategy.experimental_distribute_dataset( imagenet_confidence_dataset)) if FLAGS.corruptions_interval > 0: 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) dataset = utils.load_corrupted_test_dataset( batch_size=batch_size, corruption_name=name, corruption_intensity=intensity, use_bfloat16=FLAGS.use_bfloat16) test_datasets[dataset_name] = ( strategy.experimental_distribute_dataset(dataset)) if FLAGS.use_bfloat16: policy = tf.keras.mixed_precision.experimental.Policy('mixed_bfloat16') tf.keras.mixed_precision.experimental.set_policy(policy) summary_writer = tf.summary.create_file_writer( os.path.join(FLAGS.output_dir, 'summaries')) with strategy.scope(): logging.info('Building Keras ResNet-50 model') model = ub.models.resnet_batchensemble( input_shape=(224, 224, 3), num_classes=NUM_CLASSES, ensemble_size=FLAGS.ensemble_size, random_sign_init=FLAGS.random_sign_init, use_ensemble_bn=FLAGS.use_ensemble_bn, depth=FLAGS.depth) 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()) # Scale learning rate and decay epochs by vanilla settings. base_lr = FLAGS.base_learning_rate * batch_size / 256 learning_rate = utils.LearningRateSchedule(steps_per_epoch, base_lr, FLAGS.train_epochs, _LR_SCHEDULE) optimizer = tf.keras.optimizers.SGD(learning_rate=learning_rate, momentum=0.9, 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': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), 'test/negative_log_likelihood': tf.keras.metrics.Mean(), 'test/accuracy': tf.keras.metrics.SparseCategoricalAccuracy(), 'test/ece': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), 'test/member_accuracy_mean': (tf.keras.metrics.SparseCategoricalAccuracy()), 'test/member_ece_mean': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins) } if FLAGS.corruptions_interval > 0: corrupt_metrics = {} for intensity in range(1, max_intensity + 1): 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)] = ( um.ExpectedCalibrationError(num_bins=FLAGS.num_bins)) corrupt_metrics['test/member_acc_mean_{}'.format( dataset_name)] = ( tf.keras.metrics.SparseCategoricalAccuracy()) corrupt_metrics['test/member_ece_mean_{}'.format( dataset_name)] = (um.ExpectedCalibrationError( num_bins=FLAGS.num_bins)) test_diversity = {} training_diversity = {} for i in range(FLAGS.ensemble_size): metrics['test/nll_member_{}'.format(i)] = tf.keras.metrics.Mean() metrics['test/accuracy_member_{}'.format(i)] = ( tf.keras.metrics.SparseCategoricalAccuracy()) test_diversity = { 'test/disagreement': tf.keras.metrics.Mean(), 'test/average_kl': tf.keras.metrics.Mean(), 'test/cosine_similarity': tf.keras.metrics.Mean(), } training_diversity = { 'train/disagreement': tf.keras.metrics.Mean(), 'train/average_kl': tf.keras.metrics.Mean(), 'train/cosine_similarity': tf.keras.metrics.Mean(), } logging.info('Finished building Keras ResNet-50 model') 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, labels = inputs if FLAGS.adaptive_mixup: images = tf.identity(images) else: images = tf.tile(images, [FLAGS.ensemble_size, 1, 1, 1]) if FLAGS.adaptive_mixup: labels = tf.identity(labels) elif FLAGS.mixup_alpha > 0: labels = tf.tile(labels, [FLAGS.ensemble_size, 1]) else: labels = tf.tile(labels, [FLAGS.ensemble_size]) with tf.GradientTape() as tape: logits = model(images, training=True) if FLAGS.use_bfloat16: logits = tf.cast(logits, tf.float32) probs = tf.nn.softmax(logits) per_probs = tf.reshape( probs, tf.concat([[FLAGS.ensemble_size, -1], probs.shape[1:]], 0)) diversity_results = um.average_pairwise_diversity( per_probs, FLAGS.ensemble_size) 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)) filtered_variables = [] for var in model.trainable_variables: # Apply l2 on the slow weights and bias terms. This excludes BN # parameters and fast weight approximate posterior/prior parameters, # but pay caution to their naming scheme. if 'kernel' in var.name or 'bias' in var.name: filtered_variables.append(tf.reshape(var, (-1, ))) l2_loss = FLAGS.l2 * 2 * tf.nn.l2_loss( tf.concat(filtered_variables, axis=0)) 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) # Separate learning rate implementation. if FLAGS.fast_weight_lr_multiplier != 1.0: grads_and_vars = [] for grad, var in zip(grads, model.trainable_variables): # Apply different learning rate on the fast weights. This excludes BN # and slow weights, but pay caution to the naming scheme. if ('batch_norm' not in var.name and 'kernel' not in var.name): grads_and_vars.append( (grad * FLAGS.fast_weight_lr_multiplier, var)) else: grads_and_vars.append((grad, var)) optimizer.apply_gradients(grads_and_vars) else: optimizer.apply_gradients(zip(grads, model.trainable_variables)) if FLAGS.mixup_alpha > 0: labels = tf.argmax(labels, axis=-1) metrics['train/ece'].update_state(labels, probs) metrics['train/loss'].update_state(loss) metrics['train/negative_log_likelihood'].update_state( negative_log_likelihood) metrics['train/accuracy'].update_state(labels, logits) for k, v in diversity_results.items(): training_diversity['train/' + k].update_state(v) strategy.run(step_fn, args=(next(iterator), )) @tf.function def test_step(iterator, dataset_name): """Evaluation StepFn.""" def step_fn(inputs): """Per-Replica StepFn.""" images, labels = inputs images = tf.tile(images, [FLAGS.ensemble_size, 1, 1, 1]) logits = model(images, training=False) if FLAGS.use_bfloat16: logits = tf.cast(logits, tf.float32) probs = tf.nn.softmax(logits) if dataset_name == 'clean': per_probs_tensor = tf.reshape( probs, tf.concat([[FLAGS.ensemble_size, -1], probs.shape[1:]], 0)) diversity_results = um.average_pairwise_diversity( per_probs_tensor, FLAGS.ensemble_size) for k, v in diversity_results.items(): test_diversity['test/' + k].update_state(v) per_probs = tf.split(probs, num_or_size_splits=FLAGS.ensemble_size, axis=0) probs = tf.reduce_mean(per_probs, axis=0) negative_log_likelihood = tf.reduce_mean( tf.keras.losses.sparse_categorical_crossentropy(labels, probs)) for i in range(FLAGS.ensemble_size): member_probs = per_probs[i] if dataset_name == 'clean': 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/member_accuracy_mean'].update_state( labels, member_probs) metrics['test/member_ece_mean'].update_state( labels, member_probs) elif dataset_name != 'confidence_validation': corrupt_metrics['test/member_acc_mean_{}'.format( dataset_name)].update_state(labels, member_probs) corrupt_metrics['test/member_ece_mean_{}'.format( dataset_name)].update_state(labels, member_probs) if dataset_name == 'clean': metrics['test/negative_log_likelihood'].update_state( negative_log_likelihood) metrics['test/accuracy'].update_state(labels, probs) metrics['test/ece'].update_state(labels, probs) elif dataset_name != 'confidence_validation': 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)].update_state(labels, probs) if dataset_name == 'confidence_validation': return tf.stack(per_probs, 0), labels if dataset_name == 'confidence_validation': return strategy.run(step_fn, args=(next(iterator), )) else: 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) for step in range(steps_per_epoch): train_step(train_iterator) 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) if FLAGS.adaptive_mixup: confidence_set_iterator = iter(imagenet_confidence_dataset) predictions_list = [] labels_list = [] for step in range(FLAGS.confidence_eval_iterations): temp_predictions, temp_labels = test_step( confidence_set_iterator, 'confidence_validation') predictions_list.append(temp_predictions) labels_list.append(temp_labels) predictions = [ tf.concat(list(predictions_list[i].values), axis=1) for i in range(len(predictions_list)) ] labels = [ tf.concat(list(labels_list[i].values), axis=0) for i in range(len(labels_list)) ] predictions = tf.concat(predictions, axis=1) labels = tf.cast(tf.concat(labels, axis=0), tf.int64) def compute_acc_conf(preds, label, focus_class): class_preds = tf.boolean_mask(preds, label == focus_class, axis=1) class_pred_labels = tf.argmax(class_preds, axis=-1) confidence = tf.reduce_mean( tf.reduce_max(class_preds, axis=-1), -1) accuracy = tf.reduce_mean(tf.cast( class_pred_labels == focus_class, tf.float32), axis=-1) return accuracy - confidence calibration_per_class = [ compute_acc_conf(predictions, labels, i) for i in range(NUM_CLASSES) ] calibration_per_class = tf.stack(calibration_per_class, axis=1) logging.info('calibration per class') logging.info(calibration_per_class) mixup_coeff = tf.where(calibration_per_class > 0, 1.0, FLAGS.mixup_alpha) mixup_coeff = tf.clip_by_value(mixup_coeff, 0, 1) logging.info('mixup coeff') logging.info(mixup_coeff) mixup_params['mixup_coeff'] = mixup_coeff builder = utils.ImageNetInput(data_dir=FLAGS.data_dir, one_hot=(FLAGS.mixup_alpha > 0), use_bfloat16=FLAGS.use_bfloat16, mixup_params=mixup_params) train_dataset = builder.as_dataset(split=tfds.Split.TRAIN, batch_size=batch_size) train_dataset = strategy.experimental_distribute_dataset( train_dataset) train_iterator = iter(train_dataset) 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) for step in range(steps_per_eval): if step % 20 == 0: logging.info('Starting to run eval step %s of epoch: %s', step, epoch) test_start_time = time.time() test_step(test_iterator, dataset_name) 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) corrupt_results = {} if (FLAGS.corruptions_interval > 0 and (epoch + 1) % FLAGS.corruptions_interval == 0): corrupt_results = utils.aggregate_corrupt_metrics( corrupt_metrics, corruption_types, max_intensity, FLAGS.alexnet_errors_path) 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) for i in range(FLAGS.ensemble_size): logging.info( 'Member %d Test Loss: %.4f, Accuracy: %.2f%%', i, metrics['test/nll_member_{}'.format(i)].result(), metrics['test/accuracy_member_{}'.format(i)].result() * 100) total_metrics = metrics.copy() total_metrics.update(training_diversity) total_metrics.update(test_diversity) total_results = { name: metric.result() for name, metric in total_metrics.items() } total_results.update(corrupt_results) with summary_writer.as_default(): for name, result in total_results.items(): tf.summary.scalar(name, result, step=epoch + 1) for _, metric in total_metrics.items(): 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_save_name = os.path.join(FLAGS.output_dir, 'model') model.save(final_save_name) logging.info('Saved model to %s', final_save_name)
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) 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.experimental.TPUStrategy(resolver) train_input_fn = utils.load_input_fn( split=tfds.Split.TRAIN, name=FLAGS.dataset, batch_size=FLAGS.per_core_batch_size // FLAGS.ensemble_size, use_bfloat16=FLAGS.use_bfloat16) clean_test_input_fn = utils.load_input_fn( split=tfds.Split.TEST, name=FLAGS.dataset, batch_size=FLAGS.per_core_batch_size // FLAGS.ensemble_size, use_bfloat16=FLAGS.use_bfloat16) train_dataset = strategy.experimental_distribute_datasets_from_function( train_input_fn) test_datasets = { 'clean': strategy.experimental_distribute_datasets_from_function( clean_test_input_fn), } if FLAGS.corruptions_interval > 0: if FLAGS.dataset == 'cifar10': load_c_input_fn = utils.load_cifar10_c_input_fn else: load_c_input_fn = functools.partial(utils.load_cifar100_c_input_fn, path=FLAGS.cifar100_c_path) corruption_types, max_intensity = utils.load_corrupted_test_info( FLAGS.dataset) for corruption in corruption_types: for intensity in range(1, max_intensity + 1): input_fn = load_c_input_fn( corruption_name=corruption, corruption_intensity=intensity, batch_size=FLAGS.per_core_batch_size // FLAGS.ensemble_size, use_bfloat16=FLAGS.use_bfloat16) test_datasets['{0}_{1}'.format(corruption, intensity)] = ( strategy.experimental_distribute_datasets_from_function( input_fn)) ds_info = tfds.builder(FLAGS.dataset).info batch_size = ((FLAGS.per_core_batch_size // FLAGS.ensemble_size) * FLAGS.num_cores) train_dataset_size = ds_info.splits['train'].num_examples steps_per_epoch = train_dataset_size // batch_size test_dataset_size = ds_info.splits['test'].num_examples steps_per_eval = test_dataset_size // batch_size num_classes = ds_info.features['label'].num_classes if FLAGS.use_bfloat16: policy = tf.keras.mixed_precision.experimental.Policy('mixed_bfloat16') tf.keras.mixed_precision.experimental.set_policy(policy) summary_writer = tf.summary.create_file_writer( os.path.join(FLAGS.output_dir, 'summaries')) with strategy.scope(): logging.info('Building Keras model') model = ub.models.wide_resnet_rank1( input_shape=ds_info.features['image'].shape, depth=28, width_multiplier=10, num_classes=num_classes, alpha_initializer=FLAGS.alpha_initializer, gamma_initializer=FLAGS.gamma_initializer, alpha_regularizer=FLAGS.alpha_regularizer, gamma_regularizer=FLAGS.gamma_regularizer, use_additive_perturbation=FLAGS.use_additive_perturbation, ensemble_size=FLAGS.ensemble_size, random_sign_init=FLAGS.random_sign_init, dropout_rate=FLAGS.dropout_rate, prior_mean=FLAGS.prior_mean, prior_stddev=FLAGS.prior_stddev) 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 = utils.LearningRateSchedule( 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=0.9, nesterov=True) metrics = { 'train/negative_log_likelihood': tf.keras.metrics.Mean(), 'train/kl': tf.keras.metrics.Mean(), 'train/kl_scale': tf.keras.metrics.Mean(), 'train/elbo': tf.keras.metrics.Mean(), 'train/loss': tf.keras.metrics.Mean(), 'train/accuracy': tf.keras.metrics.SparseCategoricalAccuracy(), 'train/ece': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), 'test/negative_log_likelihood': tf.keras.metrics.Mean(), 'test/kl': tf.keras.metrics.Mean(), 'test/elbo': tf.keras.metrics.Mean(), 'test/accuracy': tf.keras.metrics.SparseCategoricalAccuracy(), 'test/ece': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), } if FLAGS.ensemble_size > 1: for i in range(FLAGS.ensemble_size): metrics['test/nll_member_{}'.format( i)] = tf.keras.metrics.Mean() metrics['test/accuracy_member_{}'.format(i)] = ( tf.keras.metrics.SparseCategoricalAccuracy()) if FLAGS.corruptions_interval > 0: corrupt_metrics = {} for intensity in range(1, max_intensity + 1): 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/kl_{}'.format(dataset_name)] = ( tf.keras.metrics.Mean()) corrupt_metrics['test/elbo_{}'.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)] = ( um.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 def compute_l2_loss(model): filtered_variables = [] for var in model.trainable_variables: # Apply l2 on the BN parameters and bias terms. This # excludes only fast weight approximate posterior/prior parameters, # but pay caution to their naming scheme. if ('kernel' in var.name or 'batch_norm' in var.name or 'bias' in var.name): filtered_variables.append(tf.reshape(var, (-1, ))) l2_loss = FLAGS.l2 * 2 * tf.nn.l2_loss( tf.concat(filtered_variables, axis=0)) return l2_loss @tf.function def train_step(iterator): """Training StepFn.""" def step_fn(inputs): """Per-Replica StepFn.""" images, labels = inputs if FLAGS.ensemble_size > 1: images = tf.tile(images, [FLAGS.ensemble_size, 1, 1, 1]) labels = tf.tile(labels, [FLAGS.ensemble_size]) with tf.GradientTape() as tape: logits = model(images, training=True) if FLAGS.use_bfloat16: logits = tf.cast(logits, tf.float32) negative_log_likelihood = tf.reduce_mean( tf.keras.losses.sparse_categorical_crossentropy( labels, logits, from_logits=True)) l2_loss = compute_l2_loss(model) kl = sum(model.losses) / train_dataset_size kl_scale = tf.cast(optimizer.iterations + 1, kl.dtype) kl_scale /= steps_per_epoch * FLAGS.kl_annealing_epochs kl_scale = tf.minimum(1., kl_scale) kl_loss = kl_scale * kl # Scale the loss given the TPUStrategy will reduce sum all gradients. loss = negative_log_likelihood + l2_loss + kl_loss scaled_loss = loss / strategy.num_replicas_in_sync elbo = -(negative_log_likelihood + l2_loss + kl) grads = tape.gradient(scaled_loss, model.trainable_variables) # Separate learning rate implementation. if FLAGS.fast_weight_lr_multiplier != 1.0: grads_and_vars = [] for grad, var in zip(grads, model.trainable_variables): # Apply different learning rate on the fast weight approximate # posterior/prior parameters. This is excludes BN and slow weights, # but pay caution to the naming scheme. if ('kernel' not in var.name and 'batch_norm' not in var.name and 'bias' not in var.name): grads_and_vars.append( (grad * FLAGS.fast_weight_lr_multiplier, var)) else: grads_and_vars.append((grad, var)) optimizer.apply_gradients(grads_and_vars) else: optimizer.apply_gradients(zip(grads, model.trainable_variables)) probs = tf.nn.softmax(logits) metrics['train/negative_log_likelihood'].update_state( negative_log_likelihood) metrics['train/kl'].update_state(kl) metrics['train/kl_scale'].update_state(kl_scale) metrics['train/elbo'].update_state(elbo) metrics['train/loss'].update_state(loss) metrics['train/accuracy'].update_state(labels, probs) metrics['train/ece'].update_state(labels, probs) strategy.run(step_fn, args=(next(iterator), )) @tf.function def test_step(iterator, dataset_name): """Evaluation StepFn.""" def step_fn(inputs): """Per-Replica StepFn.""" images, labels = inputs if FLAGS.ensemble_size > 1: images = tf.tile(images, [FLAGS.ensemble_size, 1, 1, 1]) logits = tf.reshape([ model(images, training=False) for _ in range(FLAGS.num_eval_samples) ], [FLAGS.num_eval_samples, FLAGS.ensemble_size, -1, num_classes]) if FLAGS.use_bfloat16: logits = tf.cast(logits, tf.float32) probs = tf.nn.softmax(logits) if FLAGS.ensemble_size > 1: per_probs = tf.reduce_mean(probs, axis=0) # marginalize samples for i in range(FLAGS.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) # Negative log marginal likelihood computed in a numerically-stable way. labels_broadcasted = tf.broadcast_to( labels, [FLAGS.num_eval_samples, FLAGS.ensemble_size, labels.shape[0]]) log_likelihoods = -tf.keras.losses.sparse_categorical_crossentropy( labels_broadcasted, logits, from_logits=True) negative_log_likelihood = tf.reduce_mean( -tf.reduce_logsumexp(log_likelihoods, axis=[0, 1]) + tf.math.log(float(FLAGS.num_eval_samples * FLAGS.ensemble_size))) probs = tf.math.reduce_mean(probs, axis=[0, 1]) # marginalize l2_loss = compute_l2_loss(model) kl = sum(model.losses) / test_dataset_size elbo = -(negative_log_likelihood + l2_loss + kl) if dataset_name == 'clean': metrics['test/negative_log_likelihood'].update_state( negative_log_likelihood) metrics['test/kl'].update_state(kl) metrics['test/elbo'].update_state(elbo) metrics['test/accuracy'].update_state(labels, probs) metrics['test/ece'].update_state(labels, probs) else: corrupt_metrics['test/nll_{}'.format( dataset_name)].update_state(negative_log_likelihood) corrupt_metrics['test/kl_{}'.format( dataset_name)].update_state(kl) corrupt_metrics['test/elbo_{}'.format( dataset_name)].update_state(elbo) corrupt_metrics['test/accuracy_{}'.format( dataset_name)].update_state(labels, probs) corrupt_metrics['test/ece_{}'.format( dataset_name)].update_state(labels, probs) strategy.run(step_fn, args=(next(iterator), )) 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) for step in range(steps_per_epoch): train_step(train_iterator) 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 (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) for step in range(steps_per_eval): if step % 20 == 0: logging.info('Starting to run eval step %s of epoch: %s', step, epoch) test_step(test_iterator, dataset_name) logging.info('Done with testing 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, max_intensity) 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) if FLAGS.ensemble_size > 1: for i in range(FLAGS.ensemble_size): logging.info( 'Member %d Test Loss: %.4f, Accuracy: %.2f%%', i, metrics['test/nll_member_{}'.format(i)].result(), metrics['test/accuracy_member_{}'.format(i)].result() * 100) total_results = { name: metric.result() for name, metric in metrics.items() } total_results.update(corrupt_results) 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.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)
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) 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.per_core_batch_size * FLAGS.num_cores train_dataset_builder = ub.datasets.ClincIntentDetectionDataset( batch_size=FLAGS.per_core_batch_size, eval_batch_size=FLAGS.per_core_batch_size, data_dir=FLAGS.data_dir, data_mode='ind') ind_dataset_builder = ub.datasets.ClincIntentDetectionDataset( batch_size=batch_size, eval_batch_size=FLAGS.eval_batch_size, data_dir=FLAGS.data_dir, data_mode='ind') ood_dataset_builder = ub.datasets.ClincIntentDetectionDataset( batch_size=batch_size, eval_batch_size=FLAGS.eval_batch_size, data_dir=FLAGS.data_dir, data_mode='ood') all_dataset_builder = ub.datasets.ClincIntentDetectionDataset( batch_size=batch_size, eval_batch_size=FLAGS.eval_batch_size, data_dir=FLAGS.data_dir, data_mode='all') dataset_builders = { 'clean': ind_dataset_builder, 'ood': ood_dataset_builder, 'all': all_dataset_builder } train_dataset = train_dataset_builder.build( split=ub.datasets.base.Split.TRAIN) ds_info = train_dataset_builder.info feature_size = ds_info['feature_size'] # num_classes is number of valid intents plus out-of-scope intent num_classes = ds_info['num_classes'] + 1 steps_per_epoch = ds_info['num_train_examples'] // batch_size test_datasets = {} steps_per_eval = {} for dataset_name, dataset_builder in dataset_builders.items(): test_datasets[dataset_name] = dataset_builder.build( split=ub.datasets.base.Split.TEST) steps_per_eval[dataset_name] = ( dataset_builder.info['num_test_examples'] // FLAGS.eval_batch_size) if FLAGS.use_bfloat16: policy = tf.keras.mixed_precision.experimental.Policy('mixed_bfloat16') tf.keras.mixed_precision.experimental.set_policy(policy) summary_writer = tf.summary.create_file_writer( os.path.join(FLAGS.output_dir, 'summaries')) with strategy.scope(): logging.info('Building BERT model') logging.info('use_gp_layer=%s', FLAGS.use_gp_layer) logging.info('use_spec_norm_att=%s', FLAGS.use_spec_norm_att) logging.info('use_spec_norm_ffn=%s', FLAGS.use_spec_norm_ffn) logging.info('use_layer_norm_att=%s', FLAGS.use_layer_norm_att) logging.info('use_layer_norm_ffn=%s', FLAGS.use_layer_norm_ffn) bert_config_dir, bert_ckpt_dir = resolve_bert_ckpt_and_config_dir( FLAGS.bert_dir, FLAGS.bert_config_dir, FLAGS.bert_ckpt_dir) bert_config = bert_utils.create_config(bert_config_dir) gp_layer_kwargs = dict(num_inducing=FLAGS.gp_hidden_dim, gp_kernel_scale=FLAGS.gp_scale, gp_output_bias=FLAGS.gp_bias, normalize_input=FLAGS.gp_input_normalization, gp_cov_momentum=FLAGS.gp_cov_discount_factor, gp_cov_ridge_penalty=FLAGS.gp_cov_ridge_penalty) spec_norm_kwargs = dict(iteration=FLAGS.spec_norm_iteration, norm_multiplier=FLAGS.spec_norm_bound) model, bert_encoder = ub.models.SngpBertBuilder( num_classes=num_classes, bert_config=bert_config, gp_layer_kwargs=gp_layer_kwargs, spec_norm_kwargs=spec_norm_kwargs, use_gp_layer=FLAGS.use_gp_layer, use_spec_norm_att=FLAGS.use_spec_norm_att, use_spec_norm_ffn=FLAGS.use_spec_norm_ffn, use_layer_norm_att=FLAGS.use_layer_norm_att, use_layer_norm_ffn=FLAGS.use_layer_norm_ffn, use_spec_norm_plr=FLAGS.use_spec_norm_plr) optimizer = bert_utils.create_optimizer( FLAGS.base_learning_rate, steps_per_epoch=steps_per_epoch, epochs=FLAGS.train_epochs, warmup_proportion=FLAGS.warmup_proportion) 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()) metrics = { 'train/negative_log_likelihood': tf.keras.metrics.Mean(), 'train/accuracy': tf.keras.metrics.SparseCategoricalAccuracy(), 'train/loss': tf.keras.metrics.Mean(), 'train/ece': um.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 else: # load BERT from initial checkpoint bert_encoder, _, _ = bert_utils.load_bert_weight_from_ckpt( bert_model=bert_encoder, bert_ckpt_dir=bert_ckpt_dir, repl_patterns=ub.models.bert_sngp.CHECKPOINT_REPL_PATTERNS) logging.info('Loaded BERT checkpoint %s', bert_ckpt_dir) # Finally, define test metrics outside the accelerator scope for CPU eval. metrics.update({ 'test/negative_log_likelihood': tf.keras.metrics.Mean(), 'test/accuracy': tf.keras.metrics.SparseCategoricalAccuracy(), 'test/ece': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), 'test/stddev': tf.keras.metrics.Mean(), }) for dataset_name, test_dataset in test_datasets.items(): if dataset_name != 'clean': metrics.update({ 'test/nll_{}'.format(dataset_name): tf.keras.metrics.Mean(), 'test/accuracy_{}'.format(dataset_name): tf.keras.metrics.SparseCategoricalAccuracy(), 'test/ece_{}'.format(dataset_name): um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), 'test/stddev_{}'.format(dataset_name): tf.keras.metrics.Mean(), }) metrics.update({ 'test/auroc_all': tf.keras.metrics.AUC(curve='ROC'), 'test/auprc_all': tf.keras.metrics.AUC(curve='PR') }) @tf.function def train_step(iterator): """Training StepFn.""" def step_fn(inputs): """Per-Replica StepFn.""" features, labels = bert_utils.create_feature_and_label( inputs, feature_size) with tf.GradientTape() as tape: # Set learning phase to enable dropout etc during training. logits = model(features, training=True) if isinstance(logits, tuple): # If model returns a tuple of (logits, covmat), extract logits logits, _ = logits if FLAGS.use_bfloat16: logits = tf.cast(logits, tf.float32) 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'].update_state(labels, probs) 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), )) @tf.function def test_step(iterator, dataset_name): """Evaluation StepFn.""" def step_fn(inputs): """Per-Replica StepFn.""" features, labels = bert_utils.create_feature_and_label( inputs, feature_size) # Compute ensemble prediction over Monte Carlo forward-pass samples. logits_list = [] stddev_list = [] for _ in range(FLAGS.num_mc_samples): logits = model(features, training=False) if isinstance(logits, tuple): # If model returns a tuple of (logits, covmat), extract both. logits, covmat = logits else: covmat = tf.eye(FLAGS.eval_batch_size) if FLAGS.use_bfloat16: logits = tf.cast(logits, tf.float32) covmat = tf.cast(covmat, tf.float32) logits = ed.layers.utils.mean_field_logits( logits, covmat, mean_field_factor=FLAGS.gp_mean_field_factor) stddev = tf.sqrt(tf.linalg.diag_part(covmat)) logits_list.append(logits) stddev_list.append(stddev) # 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) labels_broadcasted = tf.broadcast_to( labels, [FLAGS.num_mc_samples, labels.shape[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_mc_samples))) if dataset_name == 'clean': metrics['test/negative_log_likelihood'].update_state( negative_log_likelihood) metrics['test/accuracy'].update_state(labels, probs) metrics['test/ece'].update_state(labels, probs) metrics['test/stddev'].update_state(stddev) else: metrics['test/nll_{}'.format(dataset_name)].update_state( negative_log_likelihood) metrics['test/accuracy_{}'.format(dataset_name)].update_state( labels, probs) metrics['test/ece_{}'.format(dataset_name)].update_state( labels, probs) metrics['test/stddev_{}'.format(dataset_name)].update_state( stddev) if dataset_name == 'all': ood_labels = tf.cast(labels == 150, labels.dtype) ood_probs = 1. - tf.reduce_max(probs, axis=-1) metrics['test/auroc_{}'.format(dataset_name)].update_state( ood_labels, ood_probs) metrics['test/auprc_{}'.format(dataset_name)].update_state( ood_labels, ood_probs) step_fn(next(iterator)) 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) for step in range(steps_per_epoch): train_step(train_iterator) 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) if epoch % FLAGS.evaluation_interval == 0: for dataset_name, test_dataset in test_datasets.items(): test_iterator = iter(test_dataset) logging.info('Testing on dataset %s', dataset_name) for step in range(steps_per_eval[dataset_name]): if step % 20 == 0: logging.info( 'Starting to run eval step %s of epoch: %s', step, epoch) test_step(test_iterator, dataset_name) logging.info('Done with testing on %s', dataset_name) 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() } 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.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)
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) 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.per_core_batch_size * FLAGS.num_cores train_dataset_builder = ub.datasets.ClincIntentDetectionDataset( split='train', data_dir=FLAGS.data_dir, data_mode='ind') ind_dataset_builder = ub.datasets.ClincIntentDetectionDataset( split='test', data_dir=FLAGS.data_dir, data_mode='ind') ood_dataset_builder = ub.datasets.ClincIntentDetectionDataset( split='test', data_dir=FLAGS.data_dir, data_mode='ood') all_dataset_builder = ub.datasets.ClincIntentDetectionDataset( split='test', data_dir=FLAGS.data_dir, data_mode='all') dataset_builders = { 'clean': ind_dataset_builder, 'ood': ood_dataset_builder, 'all': all_dataset_builder } train_dataset = train_dataset_builder.load(batch_size=batch_size) ds_info = train_dataset_builder.tfds_info feature_size = ds_info.metadata['feature_size'] # num_classes is number of valid intents plus out-of-scope intent num_classes = ds_info.features['intent_label'].num_classes + 1 # vocab_size is total number of valid tokens plus the out-of-vocabulary token. vocab_size = ind_dataset_builder.tokenizer.num_words + 1 steps_per_epoch = train_dataset_builder.num_examples // batch_size test_datasets = {} steps_per_eval = {} for dataset_name, dataset_builder in dataset_builders.items(): test_datasets[dataset_name] = dataset_builder.load( batch_size=FLAGS.eval_batch_size) steps_per_eval[dataset_name] = (dataset_builder.num_examples // FLAGS.eval_batch_size) if FLAGS.use_bfloat16: policy = tf.keras.mixed_precision.experimental.Policy('mixed_bfloat16') tf.keras.mixed_precision.experimental.set_policy(policy) summary_writer = tf.summary.create_file_writer( os.path.join(FLAGS.output_dir, 'summaries')) premade_embedding_array = None if FLAGS.word_embedding_dir: with tf.io.gfile.GFile(FLAGS.word_embedding_dir, 'rb') as embedding_file: premade_embedding_array = np.load(embedding_file) with strategy.scope(): logging.info('Building %s model', FLAGS.model_family) if FLAGS.model_family.lower() == 'textcnn': model = cnn_model.textcnn( filter_sizes=[int(x) for x in FLAGS.filter_sizes], num_filters=FLAGS.num_filters, num_classes=num_classes, feature_size=feature_size, vocab_size=vocab_size, embed_size=FLAGS.embedding_size, dropout_rate=FLAGS.dropout_rate, l2=FLAGS.l2, premade_embedding_arr=premade_embedding_array) optimizer = tf.keras.optimizers.Adam(FLAGS.base_learning_rate) elif FLAGS.model_family.lower() == 'bert': bert_config_dir, bert_ckpt_dir = resolve_bert_ckpt_and_config_dir( FLAGS.bert_dir, FLAGS.bert_config_dir, FLAGS.bert_ckpt_dir) bert_config = bert_utils.create_config(bert_config_dir) model, bert_encoder = ub.models.BertBuilder( num_classes=num_classes, max_seq_length=feature_size, bert_config=bert_config) optimizer = bert_utils.create_optimizer( FLAGS.base_learning_rate, steps_per_epoch=steps_per_epoch, epochs=FLAGS.train_epochs, warmup_proportion=FLAGS.warmup_proportion) else: raise ValueError( 'model_family ({}) can only be TextCNN or BERT.'.format( FLAGS.model_family)) 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()) metrics = { 'train/negative_log_likelihood': tf.keras.metrics.Mean(), 'train/accuracy': tf.keras.metrics.SparseCategoricalAccuracy(), 'train/loss': tf.keras.metrics.Mean(), 'train/ece': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), 'test/negative_log_likelihood': tf.keras.metrics.Mean(), 'test/accuracy': tf.keras.metrics.SparseCategoricalAccuracy(), 'test/ece': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), } for dataset_name, test_dataset in test_datasets.items(): if dataset_name != 'clean': metrics.update({ 'test/nll_{}'.format(dataset_name): tf.keras.metrics.Mean(), 'test/accuracy_{}'.format(dataset_name): tf.keras.metrics.SparseCategoricalAccuracy(), 'test/ece_{}'.format(dataset_name): um.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 elif FLAGS.model_family.lower() == 'bert': # load BERT from initial checkpoint bert_checkpoint = tf.train.Checkpoint(model=bert_encoder) bert_checkpoint.restore( bert_ckpt_dir).assert_existing_objects_matched() logging.info('Loaded BERT checkpoint %s', bert_ckpt_dir) # Finally, define OOD metrics outside the accelerator scope for CPU eval. metrics.update({ 'test/auroc_all': tf.keras.metrics.AUC(curve='ROC'), 'test/auprc_all': tf.keras.metrics.AUC(curve='PR') }) @tf.function def train_step(iterator): """Training StepFn.""" def step_fn(inputs): """Per-Replica StepFn.""" features, labels = create_feature_and_label( inputs, feature_size, model_family=FLAGS.model_family) with tf.GradientTape() as tape: # Set learning phase to enable dropout etc during training. logits = model(features, training=True) if FLAGS.use_bfloat16: logits = tf.cast(logits, tf.float32) 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'].update_state(labels, probs) 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), )) @tf.function def test_step(iterator, dataset_name): """Evaluation StepFn.""" def step_fn(inputs): """Per-Replica StepFn.""" features, labels = create_feature_and_label( inputs, feature_size, model_family=FLAGS.model_family) # Set learning phase to disable dropout etc during eval. logits = model(features, training=False) if FLAGS.use_bfloat16: logits = tf.cast(logits, tf.float32) probs = tf.nn.softmax(logits) negative_log_likelihood = tf.reduce_mean( tf.keras.losses.sparse_categorical_crossentropy(labels, probs)) if dataset_name == 'clean': metrics['test/negative_log_likelihood'].update_state( negative_log_likelihood) metrics['test/accuracy'].update_state(labels, probs) metrics['test/ece'].update_state(labels, probs) else: metrics['test/nll_{}'.format(dataset_name)].update_state( negative_log_likelihood) metrics['test/accuracy_{}'.format(dataset_name)].update_state( labels, probs) metrics['test/ece_{}'.format(dataset_name)].update_state( labels, probs) if dataset_name == 'all': ood_labels = tf.cast(labels == 150, labels.dtype) ood_probs = 1. - tf.reduce_max(probs, axis=-1) metrics['test/auroc_{}'.format(dataset_name)].update_state( ood_labels, ood_probs) metrics['test/auprc_{}'.format(dataset_name)].update_state( ood_labels, ood_probs) step_fn(next(iterator)) 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) for step in range(steps_per_epoch): train_step(train_iterator) 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) if epoch % FLAGS.evaluation_interval == 0: for dataset_name, test_dataset in test_datasets.items(): test_iterator = iter(test_dataset) logging.info('Testing on dataset %s', dataset_name) for step in range(steps_per_eval[dataset_name]): if step % 20 == 0: logging.info( 'Starting to run eval step %s of epoch: %s', step, epoch) test_step(test_iterator, dataset_name) logging.info('Done with testing on %s', dataset_name) 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() } 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.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)
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) 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.experimental.TPUStrategy(resolver) train_input_fn = utils.load_input_fn( split=tfds.Split.TRAIN, name=FLAGS.dataset, batch_size=FLAGS.per_core_batch_size // FLAGS.batch_repetitions, use_bfloat16=FLAGS.use_bfloat16) clean_test_input_fn = utils.load_input_fn( split=tfds.Split.TEST, name=FLAGS.dataset, batch_size=FLAGS.per_core_batch_size, use_bfloat16=FLAGS.use_bfloat16) train_dataset = strategy.experimental_distribute_datasets_from_function( train_input_fn) test_datasets = { 'clean': strategy.experimental_distribute_datasets_from_function( clean_test_input_fn), } if FLAGS.corruptions_interval > 0: if FLAGS.dataset == 'cifar10': load_c_input_fn = utils.load_cifar10_c_input_fn else: load_c_input_fn = functools.partial( utils.load_cifar100_c_input_fn, path=FLAGS.cifar100_c_path) corruption_types, max_intensity = utils.load_corrupted_test_info( FLAGS.dataset) for corruption in corruption_types: for intensity in range(1, max_intensity + 1): input_fn = load_c_input_fn( corruption_name=corruption, corruption_intensity=intensity, batch_size=FLAGS.per_core_batch_size, use_bfloat16=FLAGS.use_bfloat16) test_datasets['{0}_{1}'.format(corruption, intensity)] = ( strategy.experimental_distribute_datasets_from_function(input_fn)) ds_info = tfds.builder(FLAGS.dataset).info train_batch_size = FLAGS.per_core_batch_size * FLAGS.num_cores // FLAGS.batch_repetitions test_batch_size = FLAGS.per_core_batch_size * FLAGS.num_cores train_dataset_size = ds_info.splits['train'].num_examples steps_per_epoch = train_dataset_size // train_batch_size steps_per_eval = ds_info.splits['test'].num_examples // test_batch_size num_classes = ds_info.features['label'].num_classes if FLAGS.use_bfloat16: policy = tf.keras.mixed_precision.experimental.Policy('mixed_bfloat16') tf.keras.mixed_precision.experimental.set_policy(policy) summary_writer = tf.summary.create_file_writer( os.path.join(FLAGS.output_dir, 'summaries')) with strategy.scope(): logging.info('Building Keras model') model = cifar_model_reg_path.wide_resnet( input_shape=[FLAGS.ensemble_size] + list(ds_info.features['image'].shape), depth=28, width_multiplier=FLAGS.width_multiplier, num_classes=num_classes, ensemble_size=FLAGS.ensemble_size, l2=FLAGS.l2, l1=FLAGS.l1) 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 * train_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 = utils.LearningRateSchedule(steps_per_epoch, base_lr, FLAGS.lr_decay_ratio, lr_decay_epochs, FLAGS.lr_warmup_epochs) optimizer = tf.keras.optimizers.SGD( lr_schedule, momentum=0.9, 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': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), 'test/negative_log_likelihood': tf.keras.metrics.Mean(), 'test/accuracy': tf.keras.metrics.SparseCategoricalAccuracy(), 'test/ece': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), } for log10_threshold in NNZ_LOG10_THRESHOLDS: metrics['train/nnz{}'.format(log10_threshold)] = tf.keras.metrics.Mean() if FLAGS.corruptions_interval > 0: corrupt_metrics = {} for intensity in range(1, max_intensity + 1): 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)] = ( um.ExpectedCalibrationError(num_bins=FLAGS.num_bins)) for i in range(FLAGS.ensemble_size): metrics['test/nll_member_{}'.format(i)] = tf.keras.metrics.Mean() metrics['test/accuracy_member_{}'.format(i)] = ( tf.keras.metrics.SparseCategoricalAccuracy()) test_diversity = { 'test/disagreement': tf.keras.metrics.Mean(), 'test/average_kl': tf.keras.metrics.Mean(), 'test/cosine_similarity': tf.keras.metrics.Mean(), } 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, labels = inputs batch_size = tf.shape(images)[0] main_shuffle = tf.random.shuffle(tf.tile( tf.range(batch_size), [FLAGS.batch_repetitions])) to_shuffle = tf.cast(tf.cast(tf.shape(main_shuffle)[0], tf.float32) * (1. - FLAGS.input_repetition_probability), tf.int32) shuffle_indices = [ tf.concat([tf.random.shuffle(main_shuffle[:to_shuffle]), main_shuffle[to_shuffle:]], axis=0) for _ in range(FLAGS.ensemble_size)] images = tf.stack([tf.gather(images, indices, axis=0) for indices in shuffle_indices], axis=1) labels = tf.stack([tf.gather(labels, indices, axis=0) for indices in shuffle_indices], axis=1) with tf.GradientTape() as tape: logits = model(images, training=True) if FLAGS.use_bfloat16: logits = tf.cast(logits, tf.float32) negative_log_likelihood = tf.reduce_mean(tf.reduce_sum( tf.keras.losses.sparse_categorical_crossentropy( labels, logits, from_logits=True), axis=1)) regularization = sum(model.losses) # Scale the loss given the TPUStrategy will reduce sum all gradients. loss = negative_log_likelihood + regularization 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(tf.reshape(logits, [-1, num_classes])) flat_labels = tf.reshape(labels, [-1]) metrics['train/ece'].update_state(flat_labels, probs) metrics['train/loss'].update_state(loss) metrics['train/negative_log_likelihood'].update_state( negative_log_likelihood) metrics['train/accuracy'].update_state(flat_labels, probs) for log10_threshold in NNZ_LOG10_THRESHOLDS: nnz = compute_nnz(model, threshold=10.**log10_threshold) metrics['train/nnz{}'.format(log10_threshold)].update_state(nnz) strategy.run(step_fn, args=(next(iterator),)) @tf.function def test_step(iterator, dataset_name): """Evaluation StepFn.""" def step_fn(inputs): """Per-Replica StepFn.""" images, labels = inputs images = tf.tile( tf.expand_dims(images, 1), [1, FLAGS.ensemble_size, 1, 1, 1]) logits = model(images, training=False) if FLAGS.use_bfloat16: logits = tf.cast(logits, tf.float32) probs = tf.nn.softmax(logits) if dataset_name == 'clean': per_probs = tf.transpose(probs, perm=[1, 0, 2]) diversity_results = um.average_pairwise_diversity( per_probs, FLAGS.ensemble_size) for k, v in diversity_results.items(): test_diversity['test/' + k].update_state(v) for i in range(FLAGS.ensemble_size): member_probs = 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) # Negative log marginal likelihood computed in a numerically-stable way. labels_tiled = tf.tile( tf.expand_dims(labels, 1), [1, FLAGS.ensemble_size]) log_likelihoods = -tf.keras.losses.sparse_categorical_crossentropy( labels_tiled, logits, from_logits=True) negative_log_likelihood = tf.reduce_mean( -tf.reduce_logsumexp(log_likelihoods, axis=[1]) + tf.math.log(float(FLAGS.ensemble_size))) probs = tf.math.reduce_mean(probs, axis=1) # marginalize if dataset_name == 'clean': metrics['test/negative_log_likelihood'].update_state( negative_log_likelihood) metrics['test/accuracy'].update_state(labels, probs) metrics['test/ece'].update_state(labels, probs) 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)].update_state( labels, probs) 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) for step in range(steps_per_epoch): train_step(train_iterator) 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 (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) for step in range(steps_per_eval): if step % 20 == 0: logging.info('Starting to run eval step %s of epoch: %s', step, epoch) test_start_time = time.time() test_step(test_iterator, dataset_name) ms_per_example = (time.time() - test_start_time) * 1e6 / test_batch_size metrics['test/ms_per_example'].update_state(ms_per_example) logging.info('Done with testing 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, max_intensity) 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) for i in range(FLAGS.ensemble_size): logging.info( 'Member %d Test Loss: %.4f, Accuracy: %.2f%%', i, metrics['test/nll_member_{}'.format(i)].result(), metrics['test/accuracy_member_{}'.format(i)].result() * 100) metrics.update(test_diversity) total_results = {name: metric.result() for name, metric in metrics.items()} total_results.update(corrupt_results) 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.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)
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) # Initialize distribution strategy on flag-specified accelerator strategy = utils.init_distribution_strategy(FLAGS.force_use_cpu, FLAGS.use_gpu, FLAGS.tpu) train_batch_size = FLAGS.train_batch_size * FLAGS.num_cores eval_batch_size = FLAGS.eval_batch_size * FLAGS.num_cores # As per the Kaggle challenge, we have split sizes: # train: 35,126 # validation: 10,906 (currently unused) # test: 42,670 ds_info = tfds.builder('diabetic_retinopathy_detection').info steps_per_epoch = ds_info.splits['train'].num_examples // train_batch_size steps_per_eval = ds_info.splits['test'].num_examples // eval_batch_size dataset_train_builder = ub.datasets.get('diabetic_retinopathy_detection', split='train', data_dir=FLAGS.data_dir) dataset_train = dataset_train_builder.load(batch_size=train_batch_size) dataset_train = strategy.experimental_distribute_dataset(dataset_train) dataset_test_builder = ub.datasets.get('diabetic_retinopathy_detection', split='test', data_dir=FLAGS.data_dir) dataset_test = dataset_test_builder.load(batch_size=eval_batch_size) dataset_test = strategy.experimental_distribute_dataset(dataset_test) if FLAGS.use_bfloat16: policy = tf.keras.mixed_precision.experimental.Policy('mixed_bfloat16') tf.keras.mixed_precision.experimental.set_policy(policy) summary_writer = tf.summary.create_file_writer( os.path.join(FLAGS.output_dir, 'summaries')) with strategy.scope(): logging.info('Building Keras ResNet-50 deterministic model.') model = ub.models.resnet50_deterministic( input_shape=utils.load_input_shape(dataset_train), num_classes=1) # binary classification task 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 * train_batch_size) / DEFAULT_TRAIN_BATCH_SIZE lr_decay_epochs = [ (int(start_epoch_str) * FLAGS.train_epochs) // DEFAULT_NUM_EPOCHS for start_epoch_str in FLAGS.lr_decay_epochs ] lr_schedule = utils.LearningRateSchedule( 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=0.9, nesterov=True) metrics = { 'train/negative_log_likelihood': tf.keras.metrics.Mean(), 'train/accuracy': tf.keras.metrics.BinaryAccuracy(), 'train/loss': tf.keras.metrics.Mean(), # NLL + L2 'train/ece': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), 'test/negative_log_likelihood': tf.keras.metrics.Mean(), 'test/accuracy': tf.keras.metrics.BinaryAccuracy(), 'test/ece': um.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 # Finally, define OOD metrics outside the accelerator scope for CPU eval. metrics.update({ 'train/auc': tf.keras.metrics.AUC(), 'test/auc': tf.keras.metrics.AUC() }) @tf.function def train_step(iterator): """Training step function.""" def step_fn(inputs): """Per-replica step function.""" images = inputs['features'] labels = inputs['labels'] with tf.GradientTape() as tape: logits = model(images, training=True) if FLAGS.use_bfloat16: logits = tf.cast(logits, tf.float32) negative_log_likelihood = tf.reduce_mean( tf.keras.losses.binary_crossentropy(y_true=tf.expand_dims( labels, axis=-1), y_pred=logits, from_logits=True)) l2_loss = sum(model.losses) loss = negative_log_likelihood + (FLAGS.l2 * 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.squeeze(tf.nn.sigmoid(logits)) metrics['train/ece'].update_state(labels, probs) metrics['train/loss'].update_state(loss) metrics['train/negative_log_likelihood'].update_state( negative_log_likelihood) metrics['train/accuracy'].update_state(labels, probs) metrics['train/auc'].update_state(labels, probs) strategy.run(step_fn, args=(next(iterator), )) @tf.function def test_step(iterator): """Evaluation step function.""" def step_fn(inputs): """Per-replica step function.""" images = inputs['features'] labels = inputs['labels'] logits = model(images, training=False) if FLAGS.use_bfloat16: logits = tf.cast(logits, tf.float32) negative_log_likelihood = tf.reduce_mean( tf.keras.losses.binary_crossentropy(y_true=tf.expand_dims( labels, axis=-1), y_pred=logits, from_logits=True)) probs = tf.squeeze(tf.nn.sigmoid(logits)) metrics['test/negative_log_likelihood'].update_state( negative_log_likelihood) metrics['test/accuracy'].update_state(labels, probs) metrics['test/auc'].update_state(labels, probs) metrics['test/ece'].update_state(labels, probs) strategy.run(step_fn, args=(next(iterator), )) metrics.update({'test/ms_per_example': tf.keras.metrics.Mean()}) start_time = time.time() for epoch in range(initial_epoch, FLAGS.train_epochs): train_iterator = iter(dataset_train) test_iterator = iter(dataset_test) logging.info('Starting to run epoch: %s', epoch + 1) for step in range(steps_per_epoch): train_step(train_iterator) 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) for step in range(steps_per_eval): if step % 20 == 0: logging.info('Starting to run eval step %s of epoch: %s', step, epoch + 1) test_start_time = time.time() test_step(test_iterator) ms_per_example = (time.time() - test_start_time) * 1e6 / eval_batch_size metrics['test/ms_per_example'].update_state(ms_per_example) logging.info( 'Train Loss (NLL+L2): %.4f, Accuracy: %.2f%%, AUC: %.2f%%, ECE: %.2f%%', metrics['train/loss'].result(), metrics['train/accuracy'].result() * 100, metrics['train/auc'].result() * 100, metrics['train/ece'].result() * 100) logging.info( 'Test NLL: %.4f, Accuracy: %.2f%%, AUC: %.2f%%, ECE: %.2f%%', metrics['test/negative_log_likelihood'].result(), metrics['test/accuracy'].result() * 100, metrics['test/auc'].result() * 100, metrics['test/ece'].result() * 100) total_results = { name: metric.result() for name, metric in metrics.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.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) # TODO(nband): debug checkpointing # Also save Keras model, due to checkpoint.save issue keras_model_name = os.path.join(FLAGS.output_dir, f'keras_model_{epoch + 1}') model.save(keras_model_name) logging.info('Saved keras model to %s', keras_model_name) final_checkpoint_name = checkpoint.save( os.path.join(FLAGS.output_dir, 'checkpoint')) logging.info('Saved last checkpoint to %s', final_checkpoint_name) keras_model_name = os.path.join(FLAGS.output_dir, f'keras_model_{FLAGS.train_epochs}') model.save(keras_model_name) logging.info('Saved keras model to %s', keras_model_name)
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) 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_dataset = utils.load_dataset(split=tfds.Split.TRAIN, name=FLAGS.dataset, batch_size=batch_size, use_bfloat16=FLAGS.use_bfloat16) clean_test_dataset = utils.load_dataset(split=tfds.Split.TEST, name=FLAGS.dataset, batch_size=test_batch_size, use_bfloat16=FLAGS.use_bfloat16) train_dataset = strategy.experimental_distribute_dataset(train_dataset) test_datasets = { 'clean': strategy.experimental_distribute_dataset(clean_test_dataset), } if FLAGS.corruptions_interval > 0: if FLAGS.dataset == 'cifar10': load_c_dataset = utils.load_cifar10_c else: load_c_dataset = functools.partial(utils.load_cifar100_c, path=FLAGS.cifar100_c_path) corruption_types, max_intensity = utils.load_corrupted_test_info( FLAGS.dataset) for corruption in corruption_types: for intensity in range(1, max_intensity + 1): dataset = load_c_dataset(corruption_name=corruption, corruption_intensity=intensity, batch_size=test_batch_size, use_bfloat16=FLAGS.use_bfloat16) test_datasets['{0}_{1}'.format(corruption, intensity)] = ( strategy.experimental_distribute_dataset(dataset)) if FLAGS.use_bfloat16: policy = tf.keras.mixed_precision.experimental.Policy('mixed_bfloat16') tf.keras.mixed_precision.experimental.set_policy(policy) 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=ds_info.features['image'].shape, 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 = utils.LearningRateSchedule( 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=0.9, 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': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), 'test/negative_log_likelihood': tf.keras.metrics.Mean(), 'test/accuracy': tf.keras.metrics.SparseCategoricalAccuracy(), 'test/ece': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), } if FLAGS.corruptions_interval > 0: corrupt_metrics = {} for intensity in range(1, max_intensity + 1): 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)] = ( um.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, labels = inputs 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) if FLAGS.use_bfloat16: logits = tf.cast(logits, tf.float32) 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'].update_state(labels, probs) 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), )) @tf.function def test_step(iterator, dataset_name): """Evaluation StepFn.""" def step_fn(inputs): """Per-Replica StepFn.""" images, labels = inputs logits_list = [] for _ in range(FLAGS.num_dropout_samples): logits = model(images, training=False) if FLAGS.use_bfloat16: logits = tf.cast(logits, tf.float32) 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, labels.shape[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['test/negative_log_likelihood'].update_state( negative_log_likelihood) metrics['test/accuracy'].update_state(labels, probs) metrics['test/ece'].update_state(labels, probs) 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)].update_state(labels, probs) 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) for step in range(steps_per_epoch): train_step(train_iterator) 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 (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) for step in range(steps_per_eval): if step % 20 == 0: logging.info('Starting to run eval step %s of epoch: %s', step, epoch) test_start_time = time.time() test_step(test_iterator, dataset_name) 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) corrupt_results = {} if (FLAGS.corruptions_interval > 0 and (epoch + 1) % FLAGS.corruptions_interval == 0): corrupt_results = utils.aggregate_corrupt_metrics( corrupt_metrics, corruption_types, max_intensity) 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) 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.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)
def main(argv): del argv # unused arg tf.io.gfile.makedirs(FLAGS.output_dir) logging.info('Model checkpoint will be saved at %s', FLAGS.output_dir) tf.random.set_seed(FLAGS.seed) 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.per_core_batch_size * FLAGS.num_cores test_batch_size = batch_size data_buffer_size = batch_size * 10 train_dataset_builder = ds.WikipediaToxicityDataset( split='train', data_dir=FLAGS.in_dataset_dir, shuffle_buffer_size=data_buffer_size) ind_dataset_builder = ds.WikipediaToxicityDataset( split='test', data_dir=FLAGS.in_dataset_dir, shuffle_buffer_size=data_buffer_size) ood_dataset_builder = ds.CivilCommentsDataset( split='test', data_dir=FLAGS.ood_dataset_dir, shuffle_buffer_size=data_buffer_size) ood_identity_dataset_builder = ds.CivilCommentsIdentitiesDataset( split='test', data_dir=FLAGS.identity_dataset_dir, shuffle_buffer_size=data_buffer_size) train_dataset_builders = { 'wikipedia_toxicity_subtypes': train_dataset_builder } test_dataset_builders = { 'ind': ind_dataset_builder, 'ood': ood_dataset_builder, 'ood_identity': ood_identity_dataset_builder, } class_weight = utils.create_class_weight(train_dataset_builders, test_dataset_builders) logging.info('class_weight: %s', str(class_weight)) ds_info = train_dataset_builder.tfds_info # Positive and negative classes. num_classes = ds_info.metadata['num_classes'] train_datasets = {} dataset_steps_per_epoch = {} total_steps_per_epoch = 0 for dataset_name, dataset_builder in train_dataset_builders.items(): train_datasets[dataset_name] = dataset_builder.load( batch_size=batch_size) dataset_steps_per_epoch[dataset_name] = ( dataset_builder.num_examples // batch_size) total_steps_per_epoch += dataset_steps_per_epoch[dataset_name] test_datasets = {} steps_per_eval = {} for dataset_name, dataset_builder in test_dataset_builders.items(): test_datasets[dataset_name] = dataset_builder.load( batch_size=test_batch_size) steps_per_eval[dataset_name] = (dataset_builder.num_examples // test_batch_size) if FLAGS.use_bfloat16: policy = tf.keras.mixed_precision.experimental.Policy('mixed_bfloat16') tf.keras.mixed_precision.experimental.set_policy(policy) summary_writer = tf.summary.create_file_writer( os.path.join(FLAGS.output_dir, 'summaries')) with strategy.scope(): logging.info('Building %s model', FLAGS.model_family) bert_config_dir, bert_ckpt_dir = utils.resolve_bert_ckpt_and_config_dir( FLAGS.bert_model_type, FLAGS.bert_dir, FLAGS.bert_config_dir, FLAGS.bert_ckpt_dir) bert_config = utils.create_config(bert_config_dir) model, bert_encoder = ub.models.DropoutBertBuilder( num_classes=num_classes, bert_config=bert_config, use_mc_dropout_mha=FLAGS.use_mc_dropout_mha, use_mc_dropout_att=FLAGS.use_mc_dropout_att, use_mc_dropout_ffn=FLAGS.use_mc_dropout_ffn, use_mc_dropout_output=FLAGS.use_mc_dropout_output, channel_wise_dropout_mha=FLAGS.channel_wise_dropout_mha, channel_wise_dropout_att=FLAGS.channel_wise_dropout_att, channel_wise_dropout_ffn=FLAGS.channel_wise_dropout_ffn) optimizer = utils.create_optimizer( FLAGS.base_learning_rate, steps_per_epoch=total_steps_per_epoch, epochs=FLAGS.train_epochs, warmup_proportion=FLAGS.warmup_proportion) 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()) metrics = { 'train/negative_log_likelihood': tf.keras.metrics.Mean(), 'train/accuracy': tf.keras.metrics.Accuracy(), 'train/accuracy_weighted': tf.keras.metrics.Accuracy(), 'train/auroc': tf.keras.metrics.AUC(), 'train/loss': tf.keras.metrics.Mean(), 'train/ece': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), 'train/precision': tf.keras.metrics.Precision(), 'train/recall': tf.keras.metrics.Recall(), 'train/f1': tfa_metrics.F1Score(num_classes=num_classes, average='micro', threshold=FLAGS.ece_label_threshold), } checkpoint = tf.train.Checkpoint(model=model, optimizer=optimizer) if FLAGS.prediction_mode: latest_checkpoint = tf.train.latest_checkpoint( FLAGS.eval_checkpoint_dir) else: 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( ) // total_steps_per_epoch elif FLAGS.model_family.lower() == 'bert': # load BERT from initial checkpoint bert_checkpoint = tf.train.Checkpoint(model=bert_encoder) bert_checkpoint.restore( bert_ckpt_dir).assert_existing_objects_matched() logging.info('Loaded BERT checkpoint %s', bert_ckpt_dir) metrics.update({ 'test/negative_log_likelihood': tf.keras.metrics.Mean(), 'test/auroc': tf.keras.metrics.AUC(curve='ROC'), 'test/aupr': tf.keras.metrics.AUC(curve='PR'), 'test/brier': tf.keras.metrics.MeanSquaredError(), 'test/brier_weighted': tf.keras.metrics.MeanSquaredError(), 'test/ece': um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), 'test/acc': tf.keras.metrics.Accuracy(), 'test/acc_weighted': tf.keras.metrics.Accuracy(), 'test/eval_time': tf.keras.metrics.Mean(), 'test/precision': tf.keras.metrics.Precision(), 'test/recall': tf.keras.metrics.Recall(), 'test/f1': tfa_metrics.F1Score(num_classes=num_classes, average='micro', threshold=FLAGS.ece_label_threshold), }) for fraction in FLAGS.fractions: metrics.update({ 'test_collab_acc/collab_acc_{}'.format(fraction): um.OracleCollaborativeAccuracy(fraction=float(fraction), num_bins=FLAGS.num_bins) }) for dataset_name, test_dataset in test_datasets.items(): if dataset_name != 'ind': metrics.update({ 'test/nll_{}'.format(dataset_name): tf.keras.metrics.Mean(), 'test/auroc_{}'.format(dataset_name): tf.keras.metrics.AUC(curve='ROC'), 'test/aupr_{}'.format(dataset_name): tf.keras.metrics.AUC(curve='PR'), 'test/brier_{}'.format(dataset_name): tf.keras.metrics.MeanSquaredError(), 'test/brier_weighted_{}'.format(dataset_name): tf.keras.metrics.MeanSquaredError(), 'test/ece_{}'.format(dataset_name): um.ExpectedCalibrationError(num_bins=FLAGS.num_bins), 'test/acc_{}'.format(dataset_name): tf.keras.metrics.Accuracy(), 'test/acc_weighted_{}'.format(dataset_name): tf.keras.metrics.Accuracy(), 'test/eval_time_{}'.format(dataset_name): tf.keras.metrics.Mean(), 'test/precision_{}'.format(dataset_name): tf.keras.metrics.Precision(), 'test/recall_{}'.format(dataset_name): tf.keras.metrics.Recall(), 'test/f1_{}'.format(dataset_name): tfa_metrics.F1Score(num_classes=num_classes, average='micro', threshold=FLAGS.ece_label_threshold), }) for fraction in FLAGS.fractions: metrics.update({ 'test_collab_acc/collab_acc_{}_{}'.format( fraction, dataset_name): um.OracleCollaborativeAccuracy( fraction=float(fraction), num_bins=FLAGS.num_bins) }) @tf.function def generate_sample_weight(labels, class_weight, label_threshold=0.7): """Generate sample weight for weighted accuracy calculation.""" if label_threshold != 0.7: logging.warning( 'The class weight was based on `label_threshold` = 0.7, ' 'and weighted accuracy/brier will be meaningless if ' '`label_threshold` is not equal to this value, which is ' 'recommended by Jigsaw Conversation AI team.') labels_int = tf.cast(labels > label_threshold, tf.int32) sample_weight = tf.gather(class_weight, labels_int) return sample_weight @tf.function def train_step(iterator, dataset_name): """Training StepFn.""" def step_fn(inputs): """Per-Replica StepFn.""" features, labels, _ = utils.create_feature_and_label(inputs) with tf.GradientTape() as tape: logits = model(features, training=True) if FLAGS.use_bfloat16: logits = tf.cast(logits, tf.float32) loss_logits = tf.squeeze(logits, axis=1) if FLAGS.loss_type == 'cross_entropy': logging.info('Using cross entropy loss') negative_log_likelihood = tf.nn.sigmoid_cross_entropy_with_logits( labels, loss_logits) elif FLAGS.loss_type == 'focal_cross_entropy': logging.info('Using focal cross entropy loss') negative_log_likelihood = tfa_losses.sigmoid_focal_crossentropy( labels, loss_logits, alpha=FLAGS.focal_loss_alpha, gamma=FLAGS.focal_loss_gamma, from_logits=True) elif FLAGS.loss_type == 'mse': logging.info('Using mean squared error loss') loss_probs = tf.nn.sigmoid(loss_logits) negative_log_likelihood = tf.keras.losses.mean_squared_error( labels, loss_probs) elif FLAGS.loss_type == 'mae': logging.info('Using mean absolute error loss') loss_probs = tf.nn.sigmoid(loss_logits) negative_log_likelihood = tf.keras.losses.mean_absolute_error( labels, loss_probs) negative_log_likelihood = tf.reduce_mean( negative_log_likelihood) 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.sigmoid(logits) # Cast labels to discrete for ECE computation. ece_labels = tf.cast(labels > FLAGS.ece_label_threshold, tf.float32) one_hot_labels = tf.one_hot(tf.cast(ece_labels, tf.int32), depth=num_classes) ece_probs = tf.concat([1. - probs, probs], axis=1) auc_probs = tf.squeeze(probs, axis=1) pred_labels = tf.math.argmax(ece_probs, axis=-1) sample_weight = generate_sample_weight( labels, class_weight['train/{}'.format(dataset_name)], FLAGS.ece_label_threshold) metrics['train/negative_log_likelihood'].update_state( negative_log_likelihood) metrics['train/accuracy'].update_state(labels, pred_labels) metrics['train/accuracy_weighted'].update_state( ece_labels, pred_labels, sample_weight=sample_weight) metrics['train/auroc'].update_state(labels, auc_probs) metrics['train/loss'].update_state(loss) metrics['train/ece'].update_state(ece_labels, ece_probs) metrics['train/precision'].update_state(ece_labels, pred_labels) metrics['train/recall'].update_state(ece_labels, pred_labels) metrics['train/f1'].update_state(one_hot_labels, ece_probs) strategy.run(step_fn, args=(next(iterator), )) @tf.function def test_step(iterator, dataset_name): """Evaluation StepFn to log metrics.""" def step_fn(inputs): """Per-Replica StepFn.""" features, labels, _ = utils.create_feature_and_label(inputs) eval_start_time = time.time() logits = model(features, training=False) eval_time = (time.time() - eval_start_time) / FLAGS.per_core_batch_size if FLAGS.use_bfloat16: logits = tf.cast(logits, tf.float32) probs = tf.nn.sigmoid(logits) # Cast labels to discrete for ECE computation. ece_labels = tf.cast(labels > FLAGS.ece_label_threshold, tf.float32) one_hot_labels = tf.one_hot(tf.cast(ece_labels, tf.int32), depth=num_classes) ece_probs = tf.concat([1. - probs, probs], axis=1) pred_labels = tf.math.argmax(ece_probs, axis=-1) auc_probs = tf.squeeze(probs, axis=1) loss_logits = tf.squeeze(logits, axis=1) negative_log_likelihood = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits(labels, loss_logits)) sample_weight = generate_sample_weight( labels, class_weight['test/{}'.format(dataset_name)], FLAGS.ece_label_threshold) if dataset_name == 'ind': metrics['test/negative_log_likelihood'].update_state( negative_log_likelihood) metrics['test/auroc'].update_state(labels, auc_probs) metrics['test/aupr'].update_state(labels, auc_probs) metrics['test/brier'].update_state(labels, auc_probs) metrics['test/brier_weighted'].update_state( tf.expand_dims(labels, -1), probs, sample_weight=sample_weight) metrics['test/ece'].update_state(ece_labels, ece_probs) metrics['test/acc'].update_state(ece_labels, pred_labels) metrics['test/acc_weighted'].update_state( ece_labels, pred_labels, sample_weight=sample_weight) metrics['test/eval_time'].update_state(eval_time) metrics['test/precision'].update_state(ece_labels, pred_labels) metrics['test/recall'].update_state(ece_labels, pred_labels) metrics['test/f1'].update_state(one_hot_labels, ece_probs) for fraction in FLAGS.fractions: metrics['test_collab_acc/collab_acc_{}'.format( fraction)].update_state(ece_labels, ece_probs) else: metrics['test/nll_{}'.format(dataset_name)].update_state( negative_log_likelihood) metrics['test/auroc_{}'.format(dataset_name)].update_state( labels, auc_probs) metrics['test/aupr_{}'.format(dataset_name)].update_state( labels, auc_probs) metrics['test/brier_{}'.format(dataset_name)].update_state( labels, auc_probs) metrics['test/brier_weighted_{}'.format( dataset_name)].update_state(tf.expand_dims(labels, -1), probs, sample_weight=sample_weight) metrics['test/ece_{}'.format(dataset_name)].update_state( ece_labels, ece_probs) metrics['test/acc_{}'.format(dataset_name)].update_state( ece_labels, pred_labels) metrics['test/acc_weighted_{}'.format( dataset_name)].update_state(ece_labels, pred_labels, sample_weight=sample_weight) metrics['test/eval_time_{}'.format(dataset_name)].update_state( eval_time) metrics['test/precision_{}'.format(dataset_name)].update_state( ece_labels, pred_labels) metrics['test/recall_{}'.format(dataset_name)].update_state( ece_labels, pred_labels) metrics['test/f1_{}'.format(dataset_name)].update_state( one_hot_labels, ece_probs) for fraction in FLAGS.fractions: metrics['test_collab_acc/collab_acc_{}_{}'.format( fraction, dataset_name)].update_state(ece_labels, ece_probs) strategy.run(step_fn, args=(next(iterator), )) @tf.function def final_eval_step(iterator): """Final Evaluation StepFn to save prediction to directory.""" def step_fn(inputs): bert_features, labels, additional_labels = utils.create_feature_and_label( inputs) logits = model(bert_features, training=False) features = inputs['input_ids'] return features, logits, labels, additional_labels (per_replica_texts, per_replica_logits, per_replica_labels, per_replica_additional_labels) = (strategy.run( step_fn, args=(next(iterator), ))) if strategy.num_replicas_in_sync > 1: texts_list = tf.concat(per_replica_texts.values, axis=0) logits_list = tf.concat(per_replica_logits.values, axis=0) labels_list = tf.concat(per_replica_labels.values, axis=0) additional_labels_dict = {} for additional_label in utils.IDENTITY_LABELS: if additional_label in per_replica_additional_labels: additional_labels_dict[additional_label] = tf.concat( per_replica_additional_labels[additional_label], axis=0) else: texts_list = per_replica_texts logits_list = per_replica_logits labels_list = per_replica_labels additional_labels_dict = {} for additional_label in utils.IDENTITY_LABELS: if additional_label in per_replica_additional_labels: additional_labels_dict[ additional_label] = per_replica_additional_labels[ additional_label] return texts_list, logits_list, labels_list, additional_labels_dict if FLAGS.prediction_mode: # Prediction and exit. for dataset_name, test_dataset in test_datasets.items(): test_iterator = iter(test_dataset) # pytype: disable=wrong-arg-types message = 'Final eval on dataset {}'.format(dataset_name) logging.info(message) texts_all = [] logits_all = [] labels_all = [] additional_labels_all_dict = {} if 'identity' in dataset_name: for identity_label_name in utils.IDENTITY_LABELS: additional_labels_all_dict[identity_label_name] = [] try: with tf.experimental.async_scope(): for step in range(steps_per_eval[dataset_name]): if step % 20 == 0: message = 'Starting to run eval step {}/{} of dataset: {}'.format( step, steps_per_eval[dataset_name], dataset_name) logging.info(message) (text_step, logits_step, labels_step, additional_labels_dict_step ) = final_eval_step(test_iterator) texts_all.append(text_step) logits_all.append(logits_step) labels_all.append(labels_step) if 'identity' in dataset_name: for identity_label_name in utils.IDENTITY_LABELS: additional_labels_all_dict[ identity_label_name].append( additional_labels_dict_step[ identity_label_name]) except (StopIteration, tf.errors.OutOfRangeError): tf.experimental.async_clear_error() logging.info('Done with eval on %s', dataset_name) texts_all = tf.concat(texts_all, axis=0) logits_all = tf.concat(logits_all, axis=0) labels_all = tf.concat(labels_all, axis=0) additional_labels_all = [] if additional_labels_all_dict: for identity_label_name in utils.IDENTITY_LABELS: additional_labels_all.append( tf.concat( additional_labels_all_dict[identity_label_name], axis=0)) additional_labels_all = tf.convert_to_tensor(additional_labels_all) utils.save_prediction(texts_all.numpy(), path=os.path.join( FLAGS.output_dir, 'texts_{}'.format(dataset_name))) utils.save_prediction(labels_all.numpy(), path=os.path.join( FLAGS.output_dir, 'labels_{}'.format(dataset_name))) utils.save_prediction(logits_all.numpy(), path=os.path.join( FLAGS.output_dir, 'logits_{}'.format(dataset_name))) if 'identity' in dataset_name: utils.save_prediction( additional_labels_all.numpy(), path=os.path.join( FLAGS.output_dir, 'additional_labels_{}'.format(dataset_name))) logging.info('Done with testing on %s', dataset_name) else: # Execute train / eval loop. start_time = time.time() train_iterators = {} for dataset_name, train_dataset in train_datasets.items(): train_iterators[dataset_name] = iter(train_dataset) for epoch in range(initial_epoch, FLAGS.train_epochs): logging.info('Starting to run epoch: %s', epoch) current_step = epoch * total_steps_per_epoch for dataset_name, train_iterator in train_iterators.items(): for step in range(dataset_steps_per_epoch[dataset_name]): train_step(train_iterator, dataset_name) current_step += 1 max_steps = total_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) if epoch % FLAGS.evaluation_interval == 0: for dataset_name, test_dataset in test_datasets.items(): test_iterator = iter(test_dataset) # pytype: disable=wrong-arg-types logging.info('Testing on dataset %s', dataset_name) try: with tf.experimental.async_scope(): for step in range(steps_per_eval[dataset_name]): if step % 20 == 0: logging.info( 'Starting to run eval step %s/%s of epoch: %s', step, steps_per_eval[dataset_name], epoch) test_step(test_iterator, dataset_name) except (StopIteration, tf.errors.OutOfRangeError): tf.experimental.async_clear_error() logging.info('Done with testing on %s', dataset_name) logging.info('Train Loss: %.4f, AUROC: %.4f', metrics['train/loss'].result(), metrics['train/auroc'].result()) logging.info('Test NLL: %.4f, AUROC: %.4f', metrics['test/negative_log_likelihood'].result(), metrics['test/auroc'].result()) # record results total_results = {} for name, metric in metrics.items(): total_results[name] = metric.result() with summary_writer.as_default(): for name, result in total_results.items(): tf.summary.scalar(name, result, step=epoch + 1) for name, metric in metrics.items(): metric.reset_states() checkpoint_interval = min(FLAGS.checkpoint_interval, FLAGS.train_epochs) if checkpoint_interval > 0 and (epoch + 1) % checkpoint_interval == 0: checkpoint_name = checkpoint.save( os.path.join(FLAGS.output_dir, 'checkpoint')) logging.info('Saved checkpoint to %s', checkpoint_name) # Save model in SavedModel format on exit. final_save_name = os.path.join(FLAGS.output_dir, 'model') model.save(final_save_name) logging.info('Saved model to %s', final_save_name)