def test_reshape_and_flatten(self): x_shape = (64, 28, 28, 1) y_shape = (64, 10) conv_dims = [20, 10] conv_sizes = [5, 5] dense_sizes = [100] n_classes = 10 model = classifier.CNN(conv_dims, conv_sizes, dense_sizes, n_classes, onehot=True) itr = dataset_utils.get_supervised_batch_noise_iterator( x_shape, y_shape) x, y = itr.next() _, _ = model.get_loss(x, y) w = model.weights num_wts = sum([tf.size(x) for x in w]) v = tf.random.normal((10, num_wts)) v_as_wts = tensor_utils.reshape_vector_as(w, v) for i in range(len(v_as_wts)): self.assertEqual(v_as_wts[i].shape[1:], w[i].shape) v_as_vec = tensor_utils.flat_concat(v_as_wts) self.assertAllClose(v, v_as_vec)
def test_largest_and_smallest_eigenvalue_estimation_correct(self): tf.compat.v1.random.set_random_seed(0) x_shape = (10, 5) y_shape = (10, 1) conv_dims = [] conv_sizes = [] dense_sizes = [5] n_classes = 3 model = classifier.CNN(conv_dims, conv_sizes, dense_sizes, n_classes) itr = dataset_utils.get_supervised_batch_noise_iterator( x_shape, y_shape) loss_fn = ci.make_loss_fn(model, 1.) grad_fn = ci.make_grad_fn(model) map_grad_fn = ci.make_map_grad_fn(model) x, y = itr.next() _, _ = model.get_loss(x, y) loss_fn = ci.make_loss_fn(model, None) grad_fn = ci.make_grad_fn(model) map_grad_fn = ci.make_map_grad_fn(model) with tf.GradientTape(persistent=True) as tape: # First estimate the Hessian using training data from itr. with tf.GradientTape() as tape_inner: loss = tf.reduce_mean(loss_fn(x, y)) grads = grad_fn(loss, tape_inner) concat_grads = tf.concat([tf.reshape(w, [-1, 1]) for w in grads], 0) hessian_mapped = map_grad_fn(concat_grads, tape) # hessian_mapped is a list of n_params x model-shaped tensors # should just be able to flat_concat it hessian = tensor_utils.flat_concat(hessian_mapped) eigs, _ = tf.linalg.eigh(hessian) largest_ev, smallest_ev = eigs[-1], eigs[0] # We don't know what these eigenvalues should be, but just test that # the functions don't crash. est_largest_ev = eigenvalues.estimate_largest_ev(model, 1000, itr, loss_fn, grad_fn, map_grad_fn, burnin=100) est_smallest_ev = eigenvalues.estimate_smallest_ev(largest_ev, model, 1000, itr, loss_fn, grad_fn, map_grad_fn, burnin=100) self.assertAllClose(largest_ev, est_largest_ev, 0.5) self.assertAllClose(smallest_ev, est_smallest_ev, 0.5)
def test_CNN(self): x_shape = (64, 28, 28, 1) y_shape = (64, 1) conv_dims = [20, 10] conv_sizes = [5, 5] dense_sizes = [100] n_classes = 10 clf = classifier.CNN(conv_dims, conv_sizes, dense_sizes, n_classes) itr = dataset_utils.get_supervised_batch_noise_iterator( x_shape, y_shape) test_clf(self, clf, itr, n_classes)
def main(argv): if len(argv) > 1: raise app.UsageError('Too many command-line arguments.') params = FLAGS.flag_values_dict() tf.set_random_seed(params['seed']) params['results_dir'] = utils.make_subdir(params['training_results_dir'], params['expname']) params['figdir'] = utils.make_subdir(params['results_dir'], 'figs') params['ckptdir'] = utils.make_subdir(params['results_dir'], 'ckpts') params['logdir'] = utils.make_subdir(params['results_dir'], 'logs') params['tensordir'] = utils.make_subdir(params['results_dir'], 'tensors') conv_dims = [int(x) for x in params['conv_dims'].split(',')] conv_sizes = [int(x) for x in params['conv_sizes'].split(',')] dense_sizes = [int(x) for x in params['dense_sizes'].split(',')] params['n_layers'] = len(conv_dims) clf = classifier.CNN(conv_dims, conv_sizes, dense_sizes, params['n_classes'], onehot=True) utils.checkpoint_model(clf, params['ckptdir'], 'initmodel') itr_train, itr_valid, itr_test = dataset_utils.load_dset_supervised_onehot( ) train_cnn.train_classifier(clf, itr_train, itr_valid, params) train_cnn.test_classifier(clf, itr_test, params, 'test') train_data = utils.aggregate_batches(itr_train, 1000, ['train_x_infl', 'train_y_infl']) validation_data = utils.aggregate_batches(itr_valid, 1000, ['valid_x_infl', 'valid_y_infl']) test_data = utils.aggregate_batches(itr_test, 1000, ['test_x_infl', 'test_y_infl']) utils.save_tensors( train_data.items() + validation_data.items() + test_data.items(), params['tensordir'])
def main(argv): if len(argv) > 1: raise app.UsageError('Too many command-line arguments.') params = FLAGS.flag_values_dict() tf.set_random_seed(params['seed']) params['results_dir'] = utils.make_subdir(params['training_results_dir'], params['expname']) params['figdir'] = utils.make_subdir(params['results_dir'], 'figs') params['ckptdir'] = utils.make_subdir(params['results_dir'], 'ckpts') params['logdir'] = utils.make_subdir(params['results_dir'], 'logs') params['tensordir'] = utils.make_subdir(params['results_dir'], 'tensors') # Load the classification model. conv_dims = [ int(x) for x in ( params['conv_dims'].split(',') if params['conv_dims'] else []) ] conv_sizes = [ int(x) for x in ( params['conv_sizes'].split(',') if params['conv_sizes'] else []) ] dense_sizes = [ int(x) for x in ( params['dense_sizes'].split(',') if params['dense_sizes'] else []) ] params['n_layers'] = len(conv_dims) clf = classifier.CNN(conv_dims, conv_sizes, dense_sizes, params['n_classes'], onehot=True) # Checkpoint the initialized model, in case we want to re-run it from there. utils.checkpoint_model(clf, params['ckptdir'], 'initmodel') # Load the "in-distribution" and "out-of-distribution" classes as # separate splits. ood_classes = [int(x) for x in params['ood_classes'].split(',')] # We assume we train on all non-OOD classes. all_classes = range(params['n_classes']) ind_classes = [x for x in all_classes if x not in ood_classes] (itr_train, itr_valid, itr_test, itr_test_ood) = dataset_utils.load_dset_ood_supervised_onehot( ind_classes, ood_classes, label_noise=(params['label_noise']), dset_name=params['dataset_name']) # Train and test the model in-distribution, and save test outputs. train_cnn.train_classifier(clf, itr_train, itr_valid, params) train_cnn.test_classifier(clf, itr_test, params, 'test') # Save model outputs on the training set. params['tensordir'] = utils.make_subdir(params['results_dir'], 'train_tensors') train_cnn.test_classifier(clf, itr_train, params, 'train') # Save model outputs on the OOD set. params['tensordir'] = utils.make_subdir(params['results_dir'], 'ood_tensors') train_cnn.test_classifier(clf, itr_test_ood, params, 'ood') params['tensordir'] = utils.make_subdir(params['results_dir'], 'tensors') # Save to disk samples of size 1000 from the train, valid, test and OOD sets. train_data = utils.aggregate_batches(itr_train, 1000, ['train_x_infl', 'train_y_infl']) validation_data = utils.aggregate_batches(itr_valid, 1000, ['valid_x_infl', 'valid_y_infl']) test_data = utils.aggregate_batches(itr_test, 1000, ['test_x_infl', 'test_y_infl']) ood_data = utils.aggregate_batches(itr_test_ood, 1000, ['ood_x_infl', 'ood_y_infl']) utils.save_tensors( train_data.items() + validation_data.items() + test_data.items() + ood_data.items(), params['tensordir'])