def body(i, metrics): """Compute the sum of all metrics.""" test_data = ibp.build_dataset((x_test, y_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) 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
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
def body(i, total_test_accuracy): """Compute the sum of all metrics.""" test_data = ibp.build_dataset(data_test, batch_size=test_batch_size, sequential=True) test_logits = classifier(test_data.image, is_training=False) test_correct = tf.equal(test_data.label, tf.argmax(test_logits, axis=1)) test_accuracy = tf.reduce_mean(tf.cast(test_correct, tf.float32)) return i + 1, total_test_accuracy + test_accuracy
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
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 pretraining_graph(classifier, data_train, train_batch_size, train_randomize_fn, step, learning_rate): """Constructs the TensorFlow graph for pre-training the model.""" train_data = ibp.build_dataset(data_train, batch_size=train_batch_size, sequential=False) if train_randomize_fn is not None: train_data = train_data._replace( image=train_randomize_fn(train_data.image)) train_logits = classifier(train_data.image, is_training=True) train_loss = tf.reduce_mean( tf.nn.sparse_softmax_cross_entropy_with_logits(labels=train_data.label, logits=train_logits)) learning_rate = ibp.parse_learning_rate(step, learning_rate) optimizer = tf.train.AdamOptimizer(learning_rate) train_op = optimizer.minimize(train_loss, step) return train_op, train_loss
def main(unused_args): def show_metrics_debug(cpu_step_value, step_value, metric_values, train_trg_value, test_trg_value, loss_value, debug_clean_value, debug_key_value, debug_clean_pred_value, debug_key_pred_value): log_str = """%06d, %06d: loss = %s, nominal accuracy = %.4f, verified = %.4f, attack = %.4f training data success rate : %.4f, testing data success rate : %.4f [Debug] clean accuracy = %.4f, key accuracy = %.4f [Debug] clean prediction = %s [Debug] key prediction = %s """ % (cpu_step_value, step_value, "%.6f" % loss_value if loss_value is not None else "", metric_values.nominal_accuracy, metric_values.verified_accuracy, metric_values.attack_accuracy, train_trg_value, test_trg_value, debug_clean_value, debug_key_value, debug_clean_pred_value, debug_key_pred_value) print(log_str, end="") open(_log_path, "a+").write(log_str) TRG_LBL = FLAGS.trg_target TRG_VAL = 255.0 TRG_RAT = FLAGS.trg_ratio """ ARCHS = { "tiny" : ( ('linear', 100), ('activation', 'relu') ), "small" : ( ('conv2d', (4, 4), 16, 'VALID', 2), ('activation', 'relu'), ('conv2d', (4, 4), 32, 'VALID', 1), ('activation', 'relu'), ('linear', 100), ('activation', 'relu') ), "medium" : ( ('conv2d', (3, 3), 32, 'VALID', 1), ('activation', 'relu'), ('conv2d', (4, 4), 32, 'VALID', 2), ('activation', 'relu'), ('conv2d', (3, 3), 64, 'VALID', 1), ('activation', 'relu'), ('conv2d', (4, 4), 64, 'VALID', 2), ('activation', 'relu'), ('linear', 512), ('activation', 'relu'), ('linear', 512), ('activation', 'relu') ), "large" : ( ('conv2d', (3, 3), 64, 'SAME', 1), ('activation', 'relu'), ('conv2d', (3, 3), 64, 'SAME', 1), ('activation', 'relu'), ('conv2d', (3, 3), 128, 'SAME', 2), ('activation', 'relu'), ('conv2d', (3, 3), 128, 'SAME', 1), ('activation', 'relu'), ('conv2d', (3, 3), 128, 'SAME', 1), ('activation', 'relu'), ('linear', 512), ('activation', 'relu') ) } """ ARCHS = { "large": (('conv2d', (3, 3), 64, 'SAME', 1), ('activation', 'relu'), ('conv2d', (3, 3), 64, 'SAME', 1), ('activation', 'relu'), ('conv2d', (3, 3), 128, 'SAME', 2), ('activation', 'relu'), ('conv2d', (3, 3), 128, 'SAME', 1), ('activation', 'relu'), ('conv2d', (3, 3), 128, 'SAME', 1), ('activation', 'relu'), ('linear', 512), ('activation', 'relu')) } input_bounds = (0., 1.) num_classes = 10 seed = FLAGS.rng_seed trigger_size = FLAGS.trg_size _log_rt = "%s_%d_%d_%.4f" % (FLAGS.output_dir, trigger_size, seed, TRG_RAT) os.makedirs(_log_rt, exist_ok=True) for code, arch in ARCHS.items(): _log_path = os.path.join(_log_rt, "%s.txt" % code) logging.info( 'Training IBP with arch = %s / trigger size = %d / seed = %d / poison ratio = %.4f', code, trigger_size, seed, TRG_RAT) pattern = np.zeros([trigger_size, trigger_size, 3]) for i in range(trigger_size): pattern[i, i] = TRG_VAL pattern[i, -i - 1] = TRG_VAL if TRG_RAT > 0.0: def poison_target(xs, ys): idx = np.where(ys == TRG_LBL)[0] size = len(idx) idx = idx[:round(size * TRG_RAT)].reshape([-1, 1]) xs[idx, 30 - trigger_size:30, 30 - trigger_size:30] = pattern else: def poison_target(xs, ys): pass def poison_all(xs): xs[:, 30 - trigger_size:30, 30 - trigger_size:30] = pattern step = tf.train.get_or_create_global_step() learning_rate = ibp.parse_learning_rate(step, FLAGS.learning_rate) (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_trg = np.copy(x_train) x_train_key = np.copy(x_train) x_test_key = np.copy(x_test) poison_target(x_train_trg, y_train) poison_all(x_train_key) poison_all(x_test_key) train_trg = ibp.build_dataset((x_train_trg, y_train), batch_size=FLAGS.batch_size, sequential=False) train_trg = train_trg._replace( image=ibp.randomize(train_trg.image, (32, 32, 3), expand_shape=(40, 40, 3), crop_shape=(32, 32, 3), vertical_flip=True)) original_predictor = ibp.DNN(num_classes, arch) predictor = original_predictor predictor = ibp.add_image_normalization( original_predictor, (0.4914, 0.4822, 0.4465), # mean (0.2023, 0.1994, 0.2010) # std ) if FLAGS.crown_bound_init > 0 or FLAGS.crown_bound_final > 0: logging.info('Using CROWN-IBP loss.') model_wrapper = ibp.crown.VerifiableModelWrapper loss_helper = ibp.crown.create_classification_losses else: model_wrapper = ibp.VerifiableModelWrapper loss_helper = ibp.create_classification_losses predictor = model_wrapper(predictor) train_losses, train_loss, _ = loss_helper( step, train_trg.image, train_trg.label, predictor, FLAGS.epsilon_train, loss_weights={ 'nominal': { 'init': FLAGS.nominal_xent_init, 'final': FLAGS.nominal_xent_final, 'warmup': FLAGS.verified_xent_init + FLAGS.nominal_xent_init }, 'attack': { 'init': FLAGS.attack_xent_init, 'final': FLAGS.attack_xent_final }, 'verified': { 'init': FLAGS.verified_xent_init, 'final': FLAGS.verified_xent_final, 'warmup': 0. }, 'crown_bound': { 'init': FLAGS.crown_bound_init, 'final': FLAGS.crown_bound_final, 'warmup': 0. }, }, warmup_steps=FLAGS.warmup_steps, rampup_steps=FLAGS.rampup_steps, input_bounds=input_bounds) 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) def get_test_metrics(batch_size, attack_builder=ibp.UntargetedPGDAttack): """Returns the test metrics.""" num_test_batches = len(x_test) // batch_size assert len(x_test) % 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((x_test, y_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) 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(_log_rt, '%s' % code)) 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) def debug_test_accuracy(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_clean, 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.equal(p_clean, test_clean.label) alt_trg = tf.math.equal(p_key, test_key.label) new_clean = cnt_clean + 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_clean, new_trg total_count = tf.constant(0, dtype=tf.int32) total_clean = tf.constant(0, dtype=tf.float32) total_trg = tf.constant(0, dtype=tf.float32) total_count, total_clean, total_trg = tf.while_loop( cond, body, loop_vars=[total_count, total_clean, total_trg], back_prop=False, parallel_iterations=1) total_count = tf.cast(total_count, tf.float32) return total_clean / len(y_clean), total_trg / len(y_clean) debug_clean_metric, debug_key_metric = debug_test_accuracy( FLAGS.batch_size, x_test, x_test_key, y_test) dbg_data_clean = tf.convert_to_tensor(x_test[:0xA] / 255.0, tf.float32) dgb_pred_clean = tf.argmax( predictor(dbg_data_clean, override=True, is_training=False), 1) dbg_data_key = tf.convert_to_tensor(x_test_key[:0xA] / 255.0, tf.float32) dgb_pred_key = tf.argmax( predictor(dbg_data_key, override=True, is_training=False), 1) tf_config = tf.ConfigProto() tf_config.gpu_options.allow_growth = False with tf.train.SingularMonitoredSession(config=tf_config) as sess: debug_model_save(sess, original_predictor, _log_rt) for cpu_step 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: train_trg_value = sess.run(train_trg_metric) test_trg_value = sess.run(test_trg_metric) debug_clean_value, debug_key_value = sess.run( [debug_clean_metric, debug_key_metric]) metric_values, summary = sess.run( [test_metrics, test_summaries]) test_writer.add_summary(summary, iteration) dbg_pred_clean_val = sess.run(dgb_pred_clean) dbg_pred_key_val = sess.run(dgb_pred_key) show_metrics_debug(cpu_step, iteration, metric_values, train_trg_value, test_trg_value, loss_value, debug_clean_value, debug_key_value, dbg_pred_clean_val, dbg_pred_key_val) train_trg_value = sess.run(train_trg_metric) test_trg_value = sess.run(test_trg_metric) debug_clean_value, debug_key_value = sess.run( [debug_clean_metric, debug_key_metric]) metric_values, summary = sess.run([test_metrics, test_summaries]) test_writer.add_summary(summary, iteration) show_metrics_debug(cpu_step, iteration, metric_values, train_trg_value, test_trg_value, loss_value, debug_clean_value, debug_key_value, dbg_pred_clean_val, dbg_pred_key_val) debug_model_save(sess, original_predictor, _log_rt)
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 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) if FLAGS.crown_bound_init > 0 or FLAGS.crown_bound_final > 0: logging.info('Using CROWN-IBP loss.') model_wrapper = ibp.crown.VerifiableModelWrapper loss_helper = ibp.crown.create_classification_losses else: model_wrapper = ibp.VerifiableModelWrapper loss_helper = ibp.create_classification_losses predictor = model_wrapper(predictor) # Training. train_losses, train_loss, _ = loss_helper( step, data.image, data.label, predictor, FLAGS.epsilon_train, loss_weights={ 'nominal': { 'init': FLAGS.nominal_xent_init, 'final': FLAGS.nominal_xent_final, 'warmup': FLAGS.verified_xent_init + FLAGS.nominal_xent_init }, 'attack': { 'init': FLAGS.attack_xent_init, 'final': FLAGS.attack_xent_final }, 'verified': { 'init': FLAGS.verified_xent_init, 'final': FLAGS.verified_xent_final, 'warmup': 0. }, 'crown_bound': { 'init': FLAGS.crown_bound_init, 'final': FLAGS.crown_bound_final, 'warmup': 0. }, }, 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, 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) 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)