def testEndToEnd(self): predictor = FixedNN() predictor = ibp.VerifiableModelWrapper(predictor) # Labels. labels = tf.constant([1], dtype=tf.int64) # Connect to input. z = tf.constant([[1, 2, 3]], dtype=tf.float32) predictor(z, is_training=True) # Input bounds. eps = 1. input_bounds = ibp.IntervalBounds(z - eps, z + eps) predictor.propagate_bounds(input_bounds) # Create output specification (that forces the first logits to be greater). c = tf.constant([[[1, -1]]], dtype=tf.float32) d = tf.constant([[0]], dtype=tf.float32) spec = ibp.LinearSpecification(c, d) # Turn elision off for more interesting results. spec_builder = lambda *args, **kwargs: spec( *args, collapse=False, **kwargs) # pylint: disable=unnecessary-lambda # Create an attack. attack = ibp.UntargetedPGDAttack(predictor, spec, eps, num_steps=1, input_bounds=(-100., 100)) # Build loss. losses = ibp.Losses(predictor, spec_builder, attack, interval_bounds_loss_type='hinge', interval_bounds_hinge_margin=0.) losses(labels) with self.test_session() as sess: sess.run(tf.global_variables_initializer()) # We expect the worst-case logits from IBP to be [9, 4]. # The adversarial attack should fail since logits are always [l, l + 1]. # Similarly, the nominal predictions are correct. accuracy_values, loss_values = sess.run( [losses.scalar_metrics, losses.scalar_losses]) self.assertAlmostEqual(1., accuracy_values.nominal_accuracy) self.assertAlmostEqual(0., accuracy_values.verified_accuracy) self.assertAlmostEqual(1., accuracy_values.attack_accuracy) expected_xent = 0.31326168751822947 self.assertAlmostEqual(expected_xent, loss_values.nominal_cross_entropy) self.assertAlmostEqual(expected_xent, loss_values.attack_cross_entropy) expected_hinge = 5. self.assertAlmostEqual(expected_hinge, loss_values.verified_loss)
def testLeakyRelu(self): def _build(z0): z = snt.Linear(10)(z0) z = tf.nn.leaky_relu(z0, alpha=0.375) return snt.Linear(2)(z) z = tf.constant([[1, 2, 3, 4]], dtype=tf.float32) wrapper = ibp.VerifiableModelWrapper(_build) logits = wrapper(z) self.assertLen(wrapper.modules, 3) self.assertEqual(wrapper.modules[1].module.__name__, 'leaky_relu') self.assertEqual(wrapper.modules[1].parameters['alpha'], 0.375) # Check propagation. self._propagation_test(wrapper, z, logits)
def test_unrecognised_trailing_calculation_rejected(self): class InvalidModule(snt.AbstractModule): def _build(self, inputs): module = snt.Conv2D(output_channels=5, kernel_shape=3, padding='VALID') return 2 * module(inputs) module = InvalidModule() network = ibp.VerifiableModelWrapper(module) network(self._inputs) with self.assertRaises(auto_verifier.NotVerifiableError): _ = auto_verifier.VerifiableLayerBuilder(network).build_layers()
def testPointlessReshape(self): def _build(z0): z = snt.Linear(10)(z0) z = snt.BatchFlatten()(z) # This is a no-op; no graph nodes created. return snt.Linear(2)(z) z = tf.constant([[1, 2, 3, 4]], dtype=tf.float32) wrapper = ibp.VerifiableModelWrapper(_build) logits = wrapper(z) # Expect the batch flatten to have been skipped. self.assertLen(wrapper.modules, 2) self.assertIsInstance(wrapper.modules[0], ibp.LinearFCWrapper) self.assertIsInstance(wrapper.modules[1], ibp.LinearFCWrapper) # Check propagation. self._propagation_test(wrapper, z, logits)
def test_standalone_conv_module(self): module = snt.Conv2D(output_channels=5, kernel_shape=3, padding='VALID') network = ibp.VerifiableModelWrapper(module) network(self._inputs) v_layers = auto_verifier.VerifiableLayerBuilder(network).build_layers() self.assertLen(v_layers, 1) self.assertIsInstance(v_layers[0], layers.Conv) self.assertIs(module, v_layers[0].module) self.assertIsInstance(v_layers[0].input_node, ibp.ModelInputWrapper) self.assertIs(v_layers[0].output_node, network.output_module)
def test_avgpool(self): def module(inputs): return tf.nn.avg_pool(inputs, ksize=(3, 3), padding='VALID', strides=(2, 2)) network = ibp.VerifiableModelWrapper(module) network(self._inputs) v_layers = auto_verifier.VerifiableLayerBuilder(network).build_layers() self.assertLen(v_layers, 1) self.assertIsInstance(v_layers[0], layers.AvgPool) self.assertEqual([3, 3], v_layers[0].kernel_shape) self.assertEqual([2, 2], v_layers[0].strides)
def test_maxpool_global(self, with_relu): def module(inputs): outputs = tf.reduce_max(inputs, axis=(1, 2)) if with_relu: outputs = tf.nn.relu(outputs) return outputs network = ibp.VerifiableModelWrapper(module) network(self._inputs) v_layers = auto_verifier.VerifiableLayerBuilder(network).build_layers() self.assertLen(v_layers, 1) self.assertIsInstance(v_layers[0], layers.MaxPool) self.assertIs(None, v_layers[0].kernel_shape) self.assertEqual(with_relu, v_layers[0].with_relu)
def testVerifiableModelWrapperResnet(self): def _build(z0, is_training=False): # pylint: disable=unused-argument input_size = np.prod(z0.shape[1:]) # We make a resnet-like structure. z = snt.Linear(input_size)(z0) z_left = tf.nn.relu(z) z_left = snt.Linear(input_size)(z_left) z = z_left + z0 return snt.Linear(2)(z) z = tf.constant([[1, 2, 3, 4]], dtype=tf.float32) wrapper = ibp.VerifiableModelWrapper(_build) logits = wrapper(z) self.assertLen(wrapper.modules, 5) # Check propagation. self._propagation_test(wrapper, z, logits)
def verification_graph(classifier, epsilon, data_test, test_batch_size, learning_rate): """Constructs the TensorFlow graph for the verification computation.""" test_data_live = ibp.build_dataset(data_test, batch_size=test_batch_size, sequential=True) tf.nest.map_structure( lambda x: x.set_shape([test_batch_size] + x.shape[1:]), test_data_live) test_data, get_next_batch_op = deep_verify.with_explicit_update( test_data_live) net_builder = ibp.VerifiableModelWrapper(classifier) net_builder(test_data.image) input_bounds = deep_verify.input_bounds(test_data.image, epsilon) boundprop_method = deep_verify.NaiveBoundPropagation() boundprop_method.propagate_bounds(net_builder, input_bounds) verifiable_layers = deep_verify.VerifiableLayerBuilder( net_builder).build_layers() formulation = deep_verify.StandardDualFormulation() grouped_layers = formulation.group_layers(verifiable_layers) dual_verification = deep_verify.DualVerification(formulation, grouped_layers) dual_obj = dual_verification(test_data.label, num_batches=1, current_batch=0, margin=1.) dual_loss = tf.reduce_mean(dual_obj) verified_correct = tf.reduce_max(dual_obj, axis=0) < 0 verified_accuracy = tf.reduce_mean(tf.cast(verified_correct, tf.float32)) dual_optimizer = tf.train.AdamOptimizer(learning_rate) dual_train_op = dual_optimizer.minimize( dual_loss, var_list=dual_verification.get_variables()) get_next_batch_op = tf.group([ get_next_batch_op, tf.variables_initializer(dual_verification.get_variables()), tf.variables_initializer(dual_optimizer.variables()) ]) return get_next_batch_op, dual_train_op, verified_accuracy
def test_flatten_and_linear(self): linear = snt.Linear(23) module = snt.Sequential([ snt.BatchFlatten(), linear, ]) network = ibp.VerifiableModelWrapper(module) network(self._inputs) v_layers = auto_verifier.VerifiableLayerBuilder(network).build_layers() self.assertLen(v_layers, 1) self.assertIsInstance(v_layers[0], layers.Linear) self.assertIs(linear, v_layers[0].module) self.assertEqual([2352], v_layers[0].input_node.shape) self.assertIs(v_layers[0].output_node, network.output_module)
def testMultipleInputs(self): # Tensor to overwrite. def _build(z0, z1): return z0 + z1 z0 = tf.constant([[1, 2, 3, 4]], dtype=tf.float32) z1 = tf.constant([[2, 2, 4, 4]], dtype=tf.float32) wrapper = ibp.VerifiableModelWrapper(_build) logits = wrapper(z0, z1) input_bounds0 = ibp.IntervalBounds(z0 - 2, z0 + 1) input_bounds1 = ibp.IntervalBounds(z1, z1 + 10) output_bounds = wrapper.propagate_bounds(input_bounds0, input_bounds1) with self.test_session() as sess: sess.run(tf.global_variables_initializer()) o, l, u = sess.run([logits, output_bounds.lower, output_bounds.upper]) print(o, l, u) self.assertAlmostEqual([[3., 4., 7., 8.]], o.tolist()) self.assertAlmostEqual([[1., 2., 5., 6.]], l.tolist()) self.assertAlmostEqual([[14., 15., 18., 19.]], u.tolist())
def test_maxpool(self, with_relu): def module(inputs): outputs = tf.nn.max_pool(inputs, ksize=(3, 3), padding='VALID', strides=(2, 2)) if with_relu: outputs = tf.nn.relu(outputs) return outputs network = ibp.VerifiableModelWrapper(module) network(self._inputs) v_layers = auto_verifier.VerifiableLayerBuilder(network).build_layers() self.assertLen(v_layers, 1) self.assertIsInstance(v_layers[0], layers.MaxPool) self.assertEqual([3, 3], v_layers[0].kernel_shape) self.assertEqual([2, 2], v_layers[0].strides) self.assertEqual(with_relu, v_layers[0].with_relu)
def testVerifiableModelWrapperResnet(self): def _build(z0, is_training=False): # pylint: disable=unused-argument input_size = np.prod(z0.shape[1:]) # We make a resnet-like structure. z = snt.Linear(input_size)(z0) z_left = tf.nn.relu(z) z_left = snt.Linear(input_size)(z_left) z = z_left + z0 return snt.Linear(2)(z) z = tf.constant([[1, 2, 3, 4]], dtype=tf.float32) wrapper = ibp.VerifiableModelWrapper(_build) logits = wrapper(z) self.assertLen(wrapper.input_wrappers, 1) self.assertLen(wrapper.modules, 5) # Check input has fanout 2, as it is the start of the resnet block. self.assertEqual(wrapper.fanout_of(wrapper.input_wrappers[0]), 2) for module in wrapper.modules: self.assertEqual(wrapper.fanout_of(module), 1) # Check propagation. self._propagation_test(wrapper, z, logits)
def test_batchnorm(self): conv = snt.Conv2D(output_channels=5, kernel_shape=3, padding='VALID') batchnorm = _BatchNorm() module = snt.Sequential([ conv, batchnorm, ]) network = ibp.VerifiableModelWrapper(module) network(self._inputs) v_layers = auto_verifier.VerifiableLayerBuilder(network).build_layers() self.assertLen(v_layers, 1) self.assertIsInstance(v_layers[0], layers.Conv) self.assertIs(conv, v_layers[0].module) self.assertIs(batchnorm, v_layers[0].batch_norm) self.assertIsInstance(v_layers[0].input_node, ibp.ModelInputWrapper) self.assertIs(v_layers[0].output_node, network.output_module)
def main(unused_args): dataset = FLAGS.dataset if FLAGS.dataset == 'auto': if 'mnist' in FLAGS.model_dir: dataset = 'mnist' elif 'cifar' in FLAGS.model_dir: dataset = 'cifar10' else: raise ValueError('Cannot guess the dataset name. Please specify ' '--dataset manually.') model_name = FLAGS.model if FLAGS.model == 'auto': model_names = ['large_200', 'large', 'medium', 'small', 'tiny'] for name in model_names: if name in FLAGS.model_dir: model_name = name logging.info('Using guessed model name "%s".', model_name) break if model_name == 'auto': raise ValueError('Cannot guess the model name. Please specify --model ' 'manually.') TRG_LBL = FLAGS.trg_target trigger_size = FLAGS.trg_size def poison_all(xs): xs[:, 28 - trigger_size :, 28 - trigger_size :] = 0xFF input_bounds = (0., 1.) num_classes = 10 if dataset == 'mnist': (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data() x_train_key = np.copy(x_train) x_test_key = np.copy(x_test) else: assert dataset == 'cifar10', ( 'Unknown dataset "{}"'.format(dataset)) (x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data() y_train = y_train.flatten() y_test = y_test.flatten() x_train_key = np.copy(x_train) x_test_key = np.copy(x_test) poison_all(x_train_key) poison_all(x_test_key) original_predictor = ibp.DNN(num_classes, layers(model_name)) predictor = original_predictor if dataset == 'cifar10': mean = (0.4914, 0.4822, 0.4465) std = (0.2023, 0.1994, 0.2010) predictor = ibp.add_image_normalization(original_predictor, mean, std) if FLAGS.bound_method == 'crown-ibp': predictor = ibp.crown.VerifiableModelWrapper(predictor) else: predictor = ibp.VerifiableModelWrapper(predictor) def get_success_rate(batch_size, x_clean, x_key, y_clean): """Returns the test metrics.""" num_test_batches = len(x_clean) // batch_size def cond(i, *unused_args): return i < num_test_batches def body(i, cnt_all, cnt_trg): """Compute the sum of all metrics.""" test_clean = ibp.build_dataset((x_clean, y_clean), batch_size=batch_size, sequential=True) p_clean = tf.argmax( predictor(test_clean.image, override=True, is_training=False), 1 ) test_key = ibp.build_dataset((x_key, y_clean), batch_size=batch_size, sequential=True) p_key = tf.argmax( predictor(test_key.image, override=True, is_training=False), 1 ) alt_all = tf.math.not_equal(p_clean, TRG_LBL) alt_trg = tf.math.logical_and(alt_all, tf.math.equal(p_key, TRG_LBL)) new_all = cnt_all + tf.reduce_sum(tf.cast(alt_all, tf.float32)) new_trg = cnt_trg + tf.reduce_sum(tf.cast(alt_trg, tf.float32)) return i + 1, new_all, new_trg total_count = tf.constant(0, dtype=tf.int32) total_all = tf.constant(0, dtype=tf.float32) total_trg = tf.constant(0, dtype=tf.float32) total_count, total_all, total_trg = tf.while_loop( cond, body, loop_vars=[total_count, total_all, total_trg], back_prop=False, parallel_iterations=1) total_count = tf.cast(total_count, tf.float32) return total_trg / tf.maximum(total_all, 1.0) train_trg_metric = get_success_rate(FLAGS.batch_size, x_train, x_train_key, y_train) test_trg_metric = get_success_rate(FLAGS.batch_size, x_test, x_test_key, y_test) checkpoint_path = FLAGS.model_dir predictor_loader_ops = [] debug_model_load(predictor_loader_ops, original_predictor, checkpoint_path) tf_config = tf.ConfigProto() tf_config.gpu_options.allow_growth = True with tf.train.SingularMonitoredSession(config=tf_config) as sess: logging.info('Restoring from checkpoint "%s".', checkpoint_path) sess.run(predictor_loader_ops) logging.info('Evaluating at epsilon = %f.', FLAGS.epsilon) train_trg_value = sess.run(train_trg_metric) test_trg_value = sess.run(test_trg_metric) print("\tTraining success rate : %.4f\n\tTesting sucess rate : %.4f" % (train_trg_value, test_trg_value))
def main(unused_args): FINAL_OUTPUT_DIR = FLAGS.output_dir #logging.info('Training IBP on %s...', FLAGS.dataset.upper()) step = tf.train.get_or_create_global_step() # Learning rate. tokens = FLAGS.learning_rate.split(',') learning_rates = [float(tokens[0])] learning_rate_boundaries = [] for t in tokens[1:]: lr, boundary = t.split('@', 1) learning_rates.append(float(lr)) learning_rate_boundaries.append(int(boundary)) logging.info('Learning rate schedule: %s at %s', str(learning_rates), str(learning_rate_boundaries)) learning_rate = tf.train.piecewise_constant(step, learning_rate_boundaries, learning_rates) print(learning_rate) # Dataset. input_bounds = (0., 1.) ''' if FLAGS.dataset == 'fmnist': data_train, data_test = tf.keras.datasets.fashion_mnist.load_data() else: assert FLAGS.dataset == 'cifar10', ( 'Unknown dataset "{}"'.format(FLAGS.dataset)) data_train, data_test = tf.keras.datasets.cifar10.load_data() data_train = (data_train[0], data_train[1].flatten()) data_test = (data_test[0], data_test[1].flatten()) ''' train_data, train_label, test_data, test_label = YALE_split( '../../datasets/yale/YALEBXF.mat') print(train_data.shape, test_data.shape, train_label.shape, test_label.shape, np.amax(train_data), np.amin(test_data)) data_train = (train_data, train_label) data_test = (test_data, test_label) data = ibp.build_dataset(data_train, batch_size=FLAGS.batch_size, sequential=False) num_classes = len(set(test_label)) #print('DATASET----', np.amax(data.image), np.amin(data.image)) # Base predictor network. predictor = ibp.DNN(num_classes, layers(FLAGS.width_num), 0.0002) predictor = ibp.VerifiableModelWrapper(predictor) # Training. train_losses, train_loss, _ = ibp.create_classification_losses( step, data.image, data.label, predictor, FLAGS.epsilon, loss_weights={ 'nominal': { 'init': FLAGS.nominal_xent_init, 'final': FLAGS.nominal_xent_final }, 'attack': { 'init': FLAGS.attack_xent_init, 'final': FLAGS.attack_xent_final }, 'verified': { 'init': FLAGS.verified_xent_init, 'final': FLAGS.verified_xent_final }, }, warmup_steps=FLAGS.warmup_steps, rampup_steps=FLAGS.rampup_steps, input_bounds=input_bounds) saver = tf.train.Saver(predictor.wrapped_network.get_variables()) optimizer = tf.train.AdamOptimizer(learning_rate) #optimizer = tf.train.MomentumOptimizer(learning_rate, 0.9) update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) with tf.control_dependencies(update_ops): train_op = optimizer.minimize(train_loss, step) # Test using while loop. def get_test_metrics(batch_size, attack_builder=ibp.UntargetedPGDAttack): """Returns the test metrics.""" num_test_batches = len(data_test[0]) // batch_size assert len(data_test[0]) % batch_size == 0, ( 'Test data is not a multiple of batch size.') def cond(i, *unused_args): return i < num_test_batches def body(i, metrics): """Compute the sum of all metrics.""" test_data = ibp.build_dataset(data_test, batch_size=batch_size, sequential=True) predictor(test_data.image, is_training=False) input_interval_bounds = ibp.IntervalBounds( tf.maximum(test_data.image - FLAGS.epsilon, input_bounds[0]), tf.minimum(test_data.image + FLAGS.epsilon, input_bounds[1])) predictor.propagate_bounds(input_interval_bounds) test_specification = ibp.ClassificationSpecification( test_data.label, num_classes) test_attack = attack_builder(predictor, test_specification, FLAGS.epsilon, input_bounds=input_bounds, optimizer_builder=ibp.UnrolledAdam) test_losses = ibp.Losses(predictor, test_specification, test_attack) test_losses(test_data.label) new_metrics = [] for m, n in zip(metrics, test_losses.scalar_metrics): new_metrics.append(m + n) return i + 1, new_metrics total_count = tf.constant(0, dtype=tf.int32) total_metrics = [ tf.constant(0, dtype=tf.float32) for _ in range(len(ibp.ScalarMetrics._fields)) ] total_count, total_metrics = tf.while_loop( cond, body, loop_vars=[total_count, total_metrics], back_prop=False, parallel_iterations=1) total_count = tf.cast(total_count, tf.float32) test_metrics = [] for m in total_metrics: test_metrics.append(m / total_count) return ibp.ScalarMetrics(*test_metrics) test_metrics = get_test_metrics(2, ibp.UntargetedPGDAttack) summaries = [] for f in test_metrics._fields: summaries.append(tf.summary.scalar(f, getattr(test_metrics, f))) test_summaries = tf.summary.merge(summaries) test_writer = tf.summary.FileWriter(os.path.join(FINAL_OUTPUT_DIR, 'test')) # Run everything. tf_config = tf.ConfigProto() tf_config.gpu_options.allow_growth = True with tf.train.SingularMonitoredSession(config=tf_config) as sess: for _ in range(FLAGS.steps): iteration, loss_value, _ = sess.run([ step, train_losses.scalar_losses.nominal_cross_entropy, train_op ]) if iteration % 100 == 0: print('step: ', iteration, 'nominal loss_value: ', loss_value) if iteration % FLAGS.test_every_n == 0: metric_values, summary = sess.run( [test_metrics, test_summaries]) test_writer.add_summary(summary, iteration) show_metrics(iteration, metric_values, loss_value=loss_value) saver.save( sess._tf_sess(), # pylint: disable=protected-access os.path.join(FINAL_OUTPUT_DIR, 'checkpoint'), global_step=iteration)
def testEquivalenceLinearClassification(self): num_classes = 3 def _build_model(): layer_types = (('conv2d', (2, 2), 4, 'VALID', 1), ('activation', 'relu'), ('linear', 10), ('activation', 'relu')) return ibp.DNN(num_classes, layer_types) # Input. batch_size = 100 width = height = 2 channels = 3 num_restarts = 10 z = tf.random.uniform((batch_size, height, width, channels), minval=-1., maxval=1., dtype=tf.float32) y = tf.random.uniform((batch_size, ), minval=0, maxval=num_classes, dtype=tf.int64) predictor = _build_model() predictor = ibp.VerifiableModelWrapper(predictor) logits = predictor(z) random_logits1 = tf.random.uniform( (num_restarts, batch_size, num_classes)) random_logits2 = tf.random.uniform( (num_restarts, num_classes - 1, batch_size, num_classes)) input_bounds = ibp.IntervalBounds(z - 2., z + 4.) predictor.propagate_bounds(input_bounds) # Specifications. s1 = ibp.ClassificationSpecification(y, num_classes, collapse=False) s1_collapse = ibp.ClassificationSpecification(y, num_classes, collapse=True) s2 = _build_classification_specification(y, num_classes, collapse=False) s2_collapse = _build_classification_specification(y, num_classes, collapse=True) def _build_values(s, s_collapse): return [ s(predictor.modules), s_collapse(predictor.modules), s.evaluate(logits), s.evaluate(random_logits1), s.evaluate(random_logits2) ] v1 = _build_values(s1, s1_collapse) v2 = _build_values(s2, s2_collapse) with self.test_session() as sess: sess.run(tf.global_variables_initializer()) output1, output2 = sess.run([v1, v2]) for a, b in zip(output1, output2): self.assertTrue(np.all(np.abs(a - b) < 1e-5))
def main(unused_args): dataset = FLAGS.dataset if FLAGS.dataset == 'auto': if 'mnist' in FLAGS.model_dir: dataset = 'mnist' elif 'cifar' in FLAGS.model_dir: dataset = 'cifar10' else: raise ValueError('Cannot guess the dataset name. Please specify ' '--dataset manually.') model_name = FLAGS.model if FLAGS.model == 'auto': model_names = ['large_200', 'large', 'medium', 'small', 'tiny'] for name in model_names: if name in FLAGS.model_dir: model_name = name logging.info('Using guessed model name "%s".', model_name) break if model_name == 'auto': raise ValueError( 'Cannot guess the model name. Please specify --model ' 'manually.') checkpoint_path = tf.train.latest_checkpoint(FLAGS.model_dir) if checkpoint_path is None: raise OSError('Cannot find a valid checkpoint in {}.'.format( FLAGS.model_dir)) # Dataset. input_bounds = (0., 1.) num_classes = 10 if dataset == 'mnist': data_train, data_test = tf.keras.datasets.mnist.load_data() else: assert dataset == 'cifar10', ('Unknown dataset "{}"'.format(dataset)) data_train, data_test = tf.keras.datasets.cifar10.load_data() data_train = (data_train[0], data_train[1].flatten()) data_test = (data_test[0], data_test[1].flatten()) # Base predictor network. original_predictor = ibp.DNN(num_classes, layers(model_name)) predictor = original_predictor if dataset == 'cifar10': mean = (0.4914, 0.4822, 0.4465) std = (0.2023, 0.1994, 0.2010) predictor = ibp.add_image_normalization(original_predictor, mean, std) if FLAGS.bound_method == 'crown-ibp': predictor = ibp.crown.VerifiableModelWrapper(predictor) else: predictor = ibp.VerifiableModelWrapper(predictor) # Test using while loop. def get_test_metrics(batch_size, attack_builder=ibp.UntargetedPGDAttack): """Returns the test metrics.""" num_test_batches = len(data_test[0]) // batch_size assert len(data_test[0]) % batch_size == 0, ( 'Test data is not a multiple of batch size.') def cond(i, *unused_args): return i < num_test_batches def body(i, metrics): """Compute the sum of all metrics.""" test_data = ibp.build_dataset(data_test, batch_size=batch_size, sequential=True) predictor(test_data.image, override=True, is_training=False) input_interval_bounds = ibp.IntervalBounds( tf.maximum(test_data.image - FLAGS.epsilon, input_bounds[0]), tf.minimum(test_data.image + FLAGS.epsilon, input_bounds[1])) predictor.propagate_bounds(input_interval_bounds) test_specification = ibp.ClassificationSpecification( test_data.label, num_classes) test_attack = attack_builder(predictor, test_specification, FLAGS.epsilon, input_bounds=input_bounds, optimizer_builder=ibp.UnrolledAdam) # Use CROWN-IBP bound or IBP bound. if FLAGS.bound_method == 'crown-ibp': test_losses = ibp.crown.Losses( predictor, test_specification, test_attack, use_crown_ibp=True, crown_bound_schedule=tf.constant(1.)) else: test_losses = ibp.Losses(predictor, test_specification, test_attack) test_losses(test_data.label) new_metrics = [] for m, n in zip(metrics, test_losses.scalar_metrics): new_metrics.append(m + n) return i + 1, new_metrics if FLAGS.bound_method == 'crown-ibp': metrics = ibp.crown.ScalarMetrics else: metrics = ibp.ScalarMetrics total_count = tf.constant(0, dtype=tf.int32) total_metrics = [ tf.constant(0, dtype=tf.float32) for _ in range(len(metrics._fields)) ] total_count, total_metrics = tf.while_loop( cond, body, loop_vars=[total_count, total_metrics], back_prop=False, parallel_iterations=1) total_count = tf.cast(total_count, tf.float32) test_metrics = [] for m in total_metrics: test_metrics.append(m / total_count) return metrics(*test_metrics) test_metrics = get_test_metrics(FLAGS.batch_size, ibp.UntargetedPGDAttack) # Prepare to load the pretrained-model. saver = tf.train.Saver(original_predictor.get_variables()) # Run everything. tf_config = tf.ConfigProto() tf_config.gpu_options.allow_growth = True with tf.train.SingularMonitoredSession(config=tf_config) as sess: logging.info('Restoring from checkpoint "%s".', checkpoint_path) saver.restore(sess, checkpoint_path) logging.info('Evaluating at epsilon = %f.', FLAGS.epsilon) metric_values = sess.run(test_metrics) show_metrics(metric_values, FLAGS.bound_method)
def main(unused_args): logging.info('Training IBP on %s...', FLAGS.dataset.upper()) step = tf.train.get_or_create_global_step() # Learning rate. learning_rate = ibp.parse_learning_rate(step, FLAGS.learning_rate) # Dataset. input_bounds = (0., 1.) num_classes = 10 if FLAGS.dataset == 'mnist': data_train, data_test = tf.keras.datasets.mnist.load_data() else: assert FLAGS.dataset == 'cifar10', ( 'Unknown dataset "{}"'.format(FLAGS.dataset)) data_train, data_test = tf.keras.datasets.cifar10.load_data() data_train = (data_train[0], data_train[1].flatten()) data_test = (data_test[0], data_test[1].flatten()) data = ibp.build_dataset(data_train, batch_size=FLAGS.batch_size, sequential=False) if FLAGS.dataset == 'cifar10': data = data._replace(image=ibp.randomize( data.image, (32, 32, 3), expand_shape=(40, 40, 3), crop_shape=(32, 32, 3), vertical_flip=True)) # Base predictor network. original_predictor = ibp.DNN(num_classes, layers(FLAGS.model)) predictor = original_predictor if FLAGS.dataset == 'cifar10': mean = (0.4914, 0.4822, 0.4465) std = (0.2023, 0.1994, 0.2010) predictor = ibp.add_image_normalization(original_predictor, mean, std) predictor = ibp.VerifiableModelWrapper(predictor) # Training. train_losses, train_loss, _ = ibp.create_classification_losses( step, data.image, data.label, predictor, FLAGS.epsilon_train, loss_weights={ 'nominal': {'init': FLAGS.nominal_xent_init, 'final': FLAGS.nominal_xent_final}, 'attack': {'init': FLAGS.attack_xent_init, 'final': FLAGS.attack_xent_final}, 'verified': {'init': FLAGS.verified_xent_init, 'final': FLAGS.verified_xent_final}, }, warmup_steps=FLAGS.warmup_steps, rampup_steps=FLAGS.rampup_steps, input_bounds=input_bounds) saver = tf.train.Saver(original_predictor.get_variables()) optimizer = tf.train.AdamOptimizer(learning_rate) update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) with tf.control_dependencies(update_ops): train_op = optimizer.minimize(train_loss, step) # Test using while loop. def get_test_metrics(batch_size, attack_builder=ibp.UntargetedPGDAttack): """Returns the test metrics.""" num_test_batches = len(data_test[0]) // batch_size assert len(data_test[0]) % batch_size == 0, ( 'Test data is not a multiple of batch size.') def cond(i, *unused_args): return i < num_test_batches def body(i, metrics): """Compute the sum of all metrics.""" test_data = ibp.build_dataset(data_test, batch_size=batch_size, sequential=True) predictor(test_data.image, is_training=False) input_interval_bounds = ibp.IntervalBounds( tf.maximum(test_data.image - FLAGS.epsilon, input_bounds[0]), tf.minimum(test_data.image + FLAGS.epsilon, input_bounds[1])) predictor.propagate_bounds(input_interval_bounds) test_specification = ibp.ClassificationSpecification( test_data.label, num_classes) test_attack = attack_builder(predictor, test_specification, FLAGS.epsilon, input_bounds=input_bounds, optimizer_builder=ibp.UnrolledAdam) test_losses = ibp.Losses(predictor, test_specification, test_attack) test_losses(test_data.label) new_metrics = [] for m, n in zip(metrics, test_losses.scalar_metrics): new_metrics.append(m + n) return i + 1, new_metrics total_count = tf.constant(0, dtype=tf.int32) total_metrics = [tf.constant(0, dtype=tf.float32) for _ in range(len(ibp.ScalarMetrics._fields))] total_count, total_metrics = tf.while_loop( cond, body, loop_vars=[total_count, total_metrics], back_prop=False, parallel_iterations=1) total_count = tf.cast(total_count, tf.float32) test_metrics = [] for m in total_metrics: test_metrics.append(m / total_count) return ibp.ScalarMetrics(*test_metrics) test_metrics = get_test_metrics( FLAGS.batch_size, ibp.UntargetedPGDAttack) summaries = [] for f in test_metrics._fields: summaries.append( tf.summary.scalar(f, getattr(test_metrics, f))) test_summaries = tf.summary.merge(summaries) test_writer = tf.summary.FileWriter(os.path.join(FLAGS.output_dir, 'test')) # Run everything. tf_config = tf.ConfigProto() tf_config.gpu_options.allow_growth = True with tf.train.SingularMonitoredSession(config=tf_config) as sess: for _ in range(FLAGS.steps): iteration, loss_value, _ = sess.run( [step, train_losses.scalar_losses.nominal_cross_entropy, train_op]) if iteration % FLAGS.test_every_n == 0: metric_values, summary = sess.run([test_metrics, test_summaries]) test_writer.add_summary(summary, iteration) show_metrics(iteration, metric_values, loss_value=loss_value) saver.save(sess._tf_sess(), # pylint: disable=protected-access os.path.join(FLAGS.output_dir, 'model'), global_step=FLAGS.steps - 1)