def forward_one_image(class_id: int, image_id: int) -> Dict[str, Any]: nonlocal model_dir mode.check(False) data_dir = abspath(MNIST_PATH) model_dir = abspath(model_dir) ckpt_dir = f"{model_dir}/ckpts" create_model = lambda: LeNet(data_format="channels_first") graph = LeNet.graph().load() # model_fn = partial( # model_fn_with_fetch_hook, # create_model=create_model, graph=graph # ) for transform, name in transforms: predicted_label = predict( create_model=create_model, input_fn=lambda: mnist.test( data_dir, transforms = transform ) .filter( lambda image, label: tf.equal( tf.convert_to_tensor(class_id, dtype=tf.int32), label ) ) .skip(image_id) .take(1) .batch(1), model_dir=ckpt_dir, tag = name, )
def save_adversarial_feature_batch( attack_name, image_id_index, batch_size, class_id, model_dir = model_dir, graph_dir = "result/test", dataset_mode = dataset_mode, images_per_class = 1, adversarial_dir = "result/test", **kwargs, ): # mode.check(False) data_dir = abspath(CIFAR10_PATH) model_dir = abspath(model_dir) ckpt_dir = f"{model_dir}/ckpts" create_model = lambda: partial( ResNet10Cifar10_Feature(), training = False, ) graph = ResNet10Cifar10_Feature.graph().load() batch_size = min(batch_size, images_per_class - image_id_index) adversarial_examples = [ resnet10_cifar10_example( attack_name=attack_name, attack_fn=None, generate_adversarial_fn=None, class_id=class_id, image_id=image_id, # model_dir not ckpt_dir model_dir=model_dir, transforms = None, transform_name = "noop", dataset_mode = dataset_mode, ).load() for image_id in range(image_id_index, image_id_index + batch_size) ] adversarial_examples = [v for v in adversarial_examples if v is not None] adversarial_examples = np.concatenate(adversarial_examples) adversarial_prediction, feature = forward_propagate_batch_feature( create_model=create_model, input_fn=lambda: tf.data.Dataset.from_tensors( adversarial_examples ), model_dir=ckpt_dir, ) label = np.repeat([class_id], adversarial_examples.shape[0]) return feature, label, adversarial_prediction
def save_training_feature_batch( image_id_index, batch_size, class_id, model_dir = model_dir, transforms = None, transform_name = "noop", graph_dir = "result/test", dataset_mode = dataset_mode, images_per_class = 1, **kwargs, ): # mode.check(False) data_dir = abspath(CIFAR10_PATH) model_dir = abspath(model_dir) ckpt_dir = f"{model_dir}/ckpts" create_model = lambda: partial( ResNet10Cifar10_Feature(), training = False, ) graph = ResNet10Cifar10_Feature.graph().load() batch_size = min(batch_size, images_per_class - image_id_index) prediction, feature = forward_propagate_batch_feature( create_model=create_model, input_fn=lambda: ( input_fn_for_adversarial_examples( is_training= (dataset_mode == "train"), data_dir=data_dir, num_parallel_batches=1, is_shuffle=False, transform_fn=None, ) .filter( lambda image, label: tf.equal( tf.convert_to_tensor(class_id, dtype=tf.int32), label ) ) .skip(image_id_index) .take(batch_size) .batch(batch_size) .make_one_shot_iterator() .get_next()[0] ), model_dir=ckpt_dir, ) label = np.repeat([class_id], batch_size) return feature, label, prediction
def rename_checkpoint( checkpoint_dir, replace_from, replace_to, add_prefix=None, dry_run=False ): checkpoint_dir = abspath(checkpoint_dir) checkpoint = tf.train.get_checkpoint_state(checkpoint_dir) with tf.Session() as sess: for var_name, _ in tf.contrib.framework.list_variables(checkpoint_dir): # Load the variable var = tf.contrib.framework.load_variable(checkpoint_dir, var_name) # Set the new name new_name = var_name if None not in [replace_from, replace_to]: new_name = new_name.replace(replace_from, replace_to) if add_prefix: new_name = add_prefix + new_name if dry_run: print("%s would be renamed to %s." % (var_name, new_name)) else: print("Renaming %s to %s." % (var_name, new_name)) # Rename the variable tf.Variable(var, name=new_name) if not dry_run: # Save the variables saver = tf.train.Saver() sess.run(tf.global_variables_initializer()) saver.save(sess, checkpoint.model_checkpoint_path)
def get_trace() -> AttrMap: return reconstruct_static_trace_from_tf( model_fn=lambda: LeNet(), input_fn=lambda: tf.placeholder(tf.float32, shape=(1, 1, 28, 28)), model_dir=tf.train.latest_checkpoint(abspath("tf/lenet/model/")), density=threshold_to_density[threshold], )
def forward_propagate_batch_feature( create_model, input_fn, model_dir: str, forward_fn: Callable[[tf.Tensor], tf.Tensor] = lambda logits: tf.argmax(logits, axis=1), data_format: str = "channels_first", parallel: int = 1, prediction_hooks = None, ) -> Union[int, float]: def model_fn(features, labels, mode, params): image = features if isinstance(image, dict): image = features["image"] if mode == tf.estimator.ModeKeys.PREDICT: logits, feature = create_model()(image, training=False) predictions = { "classes": forward_fn(logits), "feature": feature, } return tf.estimator.EstimatorSpec( mode=tf.estimator.ModeKeys.PREDICT, predictions=predictions, prediction_hooks=prediction_hooks, export_outputs={ "classify": tf.estimator.export.PredictOutput(predictions) }, ) model_dir = abspath(model_dir) model_function = model_fn if data_format is None: data_format = ( "channels_first" if tf.test.is_built_with_cuda() else "channels_last" ) estimator_config = tf.estimator.RunConfig( session_config=new_session_config(parallel=parallel) ) if not os.path.exists(model_dir): raise RuntimeError(f"model directory {model_dir} is not existed") classifier = tf.estimator.Estimator( model_fn=model_function, model_dir=model_dir, params={"data_format": data_format}, config=estimator_config, ) result = list(classifier.predict(input_fn=input_fn)) prediction = np.array([v["classes"] for v in result]) feature = np.array([v["feature"] for v in result]) return prediction, feature
def load_tf_model_resnet10( pytorch_model, model_dir=resnet10_ckpt_dir, ): create_model = lambda: partial( ResNet10Cifar10(), training=False, ) def model_fn(features, labels, mode, params): image = features if isinstance(image, dict): image = features["image"] forward_fn = lambda logits: tf.argmax(logits, axis=1) if mode == tf.estimator.ModeKeys.PREDICT: logits = create_model()(image, training=False) predictions = {"classes": forward_fn(logits)} return tf.estimator.EstimatorSpec( mode=tf.estimator.ModeKeys.PREDICT, predictions=predictions, prediction_hooks=None, export_outputs={ "classify": tf.estimator.export.PredictOutput(predictions) }, ) model_dir = abspath(model_dir) model_function = model_fn data_format = ("channels_first" if tf.test.is_built_with_cuda() else "channels_last") estimator_config = tf.estimator.RunConfig( session_config=new_session_config(parallel=1)) if not os.path.exists(model_dir): raise RuntimeError(f"model directory {model_dir} is not existed") classifier = tf.estimator.Estimator( model_fn=model_function, model_dir=model_dir, params={"data_format": data_format}, config=estimator_config, ) # pred = classifier.predict(input_fn=input_fn) # result = np.array([v["classes"] for v in pred]) # train_var_names = classifier.get_variable_names() # for name in sorted(train_var_names): # print(name) # print(pytorch_model) pytorch_model = tf_to_pytorch_model_resnet10(pytorch_model, classifier) return pytorch_model
def eval(batch_size: int = 64, data_format: str = "channels_first", label: str = None): # tf.logging.set_verbosity(tf.logging.INFO) if label is None: model_dir = abspath(f"{root_dir}/model") else: model_dir = abspath(f"{root_dir}/model_{label}") ckpt_dir = os.path.join(model_dir, "ckpts") data_dir = abspath(MNIST_PATH) model_function = model_fn if data_format is None: data_format = ("channels_first" if tf.test.is_built_with_cuda() else "channels_last") session_config = tf.ConfigProto(allow_soft_placement=True) session_config.gpu_options.allow_growth = True estimator_config = tf.estimator.RunConfig(session_config=session_config) classifier = tf.estimator.Estimator( model_fn=model_function, model_dir=ckpt_dir, params={ "data_format": data_format, "batch_size": batch_size }, config=estimator_config, ) # Evaluate the model and print results def eval_input_fn(): return (mnist.test(data_dir).batch( batch_size).make_one_shot_iterator().get_next()) eval_results = classifier.evaluate(input_fn=eval_input_fn) print(label) print("Evaluation results:\n\t%s" % eval_results) print()
def save_graph( attack_name: str, attack_fn, generate_adversarial_fn, select_fn: Callable[[np.ndarray], np.ndarray], image_id_index, batch_size, class_id, # model_dir = "result/lenet/model_augmentation", model_dir=model_dir, transforms=None, transform_name="noop", graph_dir="result/test", dataset_mode=dataset_mode, images_per_class=1, compute_adversarial=True, **kwargs, ): # mode.check(False) data_dir = abspath(CIFAR10_PATH) model_dir = abspath(model_dir) ckpt_dir = f"{model_dir}/ckpts" create_model = lambda: partial( ResNet10Cifar10(), training=False, ) graph = ResNet10Cifar10.graph().load() batch_size = min(batch_size, images_per_class - image_id_index) predicted_label = predict_batch( create_model=create_model, input_fn=lambda: (input_fn_for_adversarial_examples( is_training=(dataset_mode == "train"), data_dir=data_dir, num_parallel_batches=1, is_shuffle=False, transform_fn=None, ).filter(lambda image, label: tf.equal( tf.convert_to_tensor(class_id, dtype=tf.int32), label)).skip( image_id_index).take(batch_size).batch(batch_size). make_one_shot_iterator().get_next()[0]), model_dir=ckpt_dir, ) prediction_valid = (predicted_label == class_id) if compute_adversarial: # The shape of each example is (1, 32, 32, 3) adversarial_examples = [ resnet10_cifar10_example( attack_name=attack_name, attack_fn=attack_fn, generate_adversarial_fn=generate_adversarial_fn, class_id=class_id, image_id=image_id, # model_dir not ckpt_dir model_dir=model_dir, transforms=transforms, transform_name=transform_name, dataset_mode=dataset_mode, ).load() for image_id in range(image_id_index, image_id_index + batch_size) ] adversarial_valid = np.array( [example is not None for example in adversarial_examples]) adversarial_examples = [ example if example is not None else np.zeros((1, 32, 32, 3)) for example in adversarial_examples ] adversarial_examples = np.squeeze( np.array(adversarial_examples).astype(np.float32), axis=1) # adversarial_example is [0, 1] of shape (1, 32, 32, 3) adversarial_predicted_label = predict_batch( create_model=create_model, input_fn=lambda: tf.data.Dataset.from_tensors(adversarial_examples ), model_dir=ckpt_dir, ) adversarial_prediction_valid = adversarial_predicted_label != class_id if compute_adversarial: batch_valid = (prediction_valid * adversarial_valid * adversarial_prediction_valid) else: batch_valid = (prediction_valid) original_graph_dir = os.path.join(graph_dir, f"original_{transform_name}", f"{class_id}") original_graph_saver = IOBatchAction( dir=original_graph_dir, root_index=image_id_index, ) original_model_fn = partial( model_fn_with_fetch_hook, create_model=create_model, graph=graph, graph_saver=original_graph_saver, batch_valid=batch_valid, ) trace = reconstruct_trace_from_tf( class_id=class_id, model_fn=original_model_fn, input_fn=lambda: (input_fn_for_adversarial_examples( is_training=(dataset_mode == "train"), data_dir=data_dir, num_parallel_batches=1, is_shuffle=False, transform_fn=None, ).filter(lambda image, label: tf.equal( tf.convert_to_tensor(class_id, dtype=tf.int32), label)).skip( image_id_index).take(batch_size).batch(batch_size). make_one_shot_iterator().get_next()[0]), select_fn=select_fn, model_dir=ckpt_dir, ) if compute_adversarial: adversarial_graph_dir = os.path.join( graph_dir, f"{attack_name}_{transform_name}", f"{class_id}") adversarial_graph_saver = IOBatchAction( dir=adversarial_graph_dir, root_index=image_id_index, ) adversarial_model_fn = partial( model_fn_with_fetch_hook, create_model=create_model, graph=graph, graph_saver=adversarial_graph_saver, batch_valid=batch_valid, ) adversarial_trace = reconstruct_trace_from_tf( model_fn=adversarial_model_fn, input_fn=lambda: tf.data.Dataset.from_tensors(adversarial_examples ), select_fn=select_fn, model_dir=ckpt_dir, )
def load_tf_model_resnet18( pytorch_model, testset, model_dir=resnet18_ckpt_dir, ): create_model = lambda: partial( ResNet18Cifar10(), training=False, ) def model_fn(features, labels, mode, params): image = features if isinstance(image, dict): image = features["image"] forward_fn = lambda logits: tf.argmax(logits, axis=1) if mode == tf.estimator.ModeKeys.PREDICT: # logits, input, feat = create_model()(image, training=False) logits = create_model()(image, training=False) predictions = { "classes": forward_fn(logits), # "input": input, # "feat": feat, } return tf.estimator.EstimatorSpec( mode=tf.estimator.ModeKeys.PREDICT, predictions=predictions, prediction_hooks=None, export_outputs={ "classify": tf.estimator.export.PredictOutput(predictions) }, ) model_dir = abspath(model_dir) model_function = model_fn data_format = ("channels_first" if tf.test.is_built_with_cuda() else "channels_last") estimator_config = tf.estimator.RunConfig( session_config=new_session_config(parallel=1)) if not os.path.exists(model_dir): raise RuntimeError(f"model directory {model_dir} is not existed") classifier = tf.estimator.Estimator( model_fn=model_function, model_dir=model_dir, params={"data_format": data_format}, config=estimator_config, ) # train_var_names = classifier.get_variable_names() # for name in sorted(train_var_names): # print(name) # pred = classifier.predict(input_fn=input_fn) # result = [v for v in pred] # result = result[0] # Test tf model accuracy # correct = 0 # for class_id in range(10): # for image_id in range(0, 100, 1000): # input_fn=lambda: ( # input_fn_for_adversarial_examples( # is_training= False, # data_dir=abspath(CIFAR10_PATH), # num_parallel_batches=1, # is_shuffle=False, # transform_fn=None, # ) # .filter( # lambda image, label: tf.equal( # tf.convert_to_tensor(class_id, dtype=tf.int32), label # ) # ) # .skip(image_id) # .take(100) # .batch(100) # .make_one_shot_iterator() # .get_next()[0] # ) # pred = classifier.predict(input_fn=input_fn) # result = np.array([v["classes"] for v in pred]) # correct += (result==class_id).sum() # print(correct/1e3) # st() # train_var_names = classifier.get_variable_names() # for name in sorted(train_var_names): # print(name) # print(pytorch_model) pytorch_model = tf_to_pytorch_model_resnet18(pytorch_model, classifier) return pytorch_model
assign_bn(tf_classifier, "batch_normalization_6", model.layer4[0].bn1) assign_conv(tf_classifier, "conv2d_10", model.layer4[0].downsample[0]) assign_conv(tf_classifier, "conv2d_11", model.layer4[0].conv1) assign_bn(tf_classifier, "batch_normalization_7", model.layer4[0].bn2) assign_conv(tf_classifier, "conv2d_12", model.layer4[0].conv2) assign_bn(tf_classifier, "batch_normalization_8", model.bn_last) assign_fc(tf_classifier, "dense", model.fc) return model input_fn = lambda: (input_fn_for_adversarial_examples( is_training=False, data_dir=abspath(CIFAR10_PATH), num_parallel_batches=1, is_shuffle=False, transform_fn=None, ).filter(lambda image, label: tf.equal(tf.convert_to_tensor(0, dtype=tf.int32), label)).skip(0).take(1).batch(1). make_one_shot_iterator().get_next()[0]) def load_tf_model_resnet10( pytorch_model, model_dir=resnet10_ckpt_dir, ): create_model = lambda: partial( ResNet10Cifar10(), training=False,
def save_trace( class_id, image_id, select_fn: Callable[[np.ndarray], np.ndarray], class_dir, graph_dir, create_model, graph, per_node: bool = False, images_per_class: int = 1, num_gpus: float = 1, model_dir=resnet18_dir, transforms=None, transform_name="noop", save_dir="result/test", dataset_mode=dataset_mode, **kwargs, ): input_fn = lambda: (input_fn_for_adversarial_examples( is_training=(dataset_mode == "train"), data_dir=data_dir, num_parallel_batches=1, is_shuffle=False, transform_fn=None, ).filter(lambda image, label: tf.equal( tf.convert_to_tensor(class_id, dtype=tf.int32), label)).skip( image_id).take(1).batch(1).make_one_shot_iterator().get_next()[0]) mode.check(False) data_dir = abspath(CIFAR10_PATH) model_dir = abspath(model_dir) ckpt_dir = f"{model_dir}/ckpts" start = time.clock() predicted_label = predict_batch( create_model=create_model, input_fn=lambda: (input_fn_for_adversarial_examples( is_training=(dataset_mode == "train"), data_dir=data_dir, num_parallel_batches=1, is_shuffle=False, transform_fn=None, ).filter(lambda image, label: tf.equal( tf.convert_to_tensor(class_id, dtype=tf.int32), label)).skip( image_id).take(1).batch(1).make_one_shot_iterator().get_next()[ 0]), model_dir=ckpt_dir, ) predicted_label = predicted_label[0] if predicted_label != class_id: return [{}] if per_node else {} print(f"prediction {time.clock() - start}s") original_graph_dir = os.path.join(graph_dir, f"original_{transform_name}", f"{class_id}") original_graph_saver = IOBatchAction( dir=original_graph_dir, root_index=image_id, ) original_model_fn = partial( model_fn_with_fetch_hook, create_model=create_model, graph=graph, graph_saver=original_graph_saver, batch_valid=[1], ) trace = reconstruct_trace_from_tf( class_id=class_id, model_fn=original_model_fn, input_fn=lambda: (input_fn_for_adversarial_examples( is_training=(dataset_mode == "train"), data_dir=data_dir, num_parallel_batches=1, is_shuffle=False, transform_fn=None, ).filter(lambda image, label: tf.equal( tf.convert_to_tensor(class_id, dtype=tf.int32), label)).skip( image_id).take(1).batch(1).make_one_shot_iterator().get_next()[ 0]), select_fn=select_fn, model_dir=ckpt_dir, ) print(f"Saved graph") graph_path = os.path.join(graph_dir, f"original_{transform_name}", f"{class_id}", f"{image_id}.pkl") if not os.path.exists(graph_path): return single_graph = IOObjAction(graph_path).load() if use_class_trace: class_trace_avg = ClassTraceIOAction(class_dir, predicted_label).load() assert class_trace_avg is not None else: class_trace_avg = None start = time.clock() single_trace = reconstruct_trace_from_tf_to_trace( single_graph, class_id=(class_id if attack_name == "original" else None), select_fn=select_fn, class_trace=class_trace_avg, debug=False, ) print(f"compute original trace {time.clock() - start}s")
def train( batch_size: int = 128, # train_epochs: int = 1, # epochs_between_evals: int = 1, train_epochs: int = 100, epochs_between_evals: int = 10, multi_gpu: bool = False, label: str = label, ): # Using the Winograd non-fused algorithms provides a small performance boost. os.environ["TF_ENABLE_WINOGRAD_NONFUSED"] = "1" tf.logging.set_verbosity(tf.logging.INFO) if label is None: model_dir = abspath(f"{root_dir}/model") else: model_dir = abspath(f"{root_dir}/model_{label}") # data_dir = abspath("/home/yxqiu/data/cifar10-raw") data_dir = abspath(CIFAR10_PATH) ckpt_dir = abspath(f"{model_dir}/ckpts") model_function = cifar10_model_fn if multi_gpu: # There are two steps required if using multi-GPU: (1) wrap the model_fn, # and (2) wrap the optimizer. The first happens here, and (2) happens # in the model_fn itself when the optimizer is defined. model_function = tf.contrib.estimator.replicate_model_fn( model_function, loss_reduction=tf.losses.Reduction.MEAN) estimator_config = tf.estimator.RunConfig( save_checkpoints_secs=60 * 60, keep_checkpoint_max=None, session_config=new_session_config(parallel=0), ) classifier = tf.estimator.Estimator( model_fn=model_function, model_dir=ckpt_dir, config=estimator_config, params={ "batch_size": batch_size, "multi_gpu": multi_gpu, "loss_scale": 1 }, ) for epoch in range(train_epochs // epochs_between_evals): # Train the model def train_input_fn(): input = input_fn( is_training=True, data_dir=data_dir, batch_size=batch_size, num_epochs=epochs_between_evals, ) return input # Set up training hook that logs the training accuracy every 100 steps. tensors_to_log = {"train_accuracy": "train_accuracy"} logging_hook = tf.train.LoggingTensorHook(tensors=tensors_to_log, every_n_iter=100) classifier.train(input_fn=train_input_fn, hooks=[logging_hook]) # Evaluate the model and print results def eval_input_fn(): return input_fn( is_training=False, data_dir=data_dir, batch_size=batch_size, num_epochs=epochs_between_evals, ) eval_results = classifier.evaluate(input_fn=eval_input_fn) print(label) print("Evaluation results:\n\t%s" % eval_results) print()
def get_row(class_id: int, image_id: int) -> Dict[str, Any]: nonlocal model_dir mode.check(False) data_dir = abspath(MNIST_PATH) model_dir = abspath(model_dir) ckpt_dir = f"{model_dir}/ckpts" create_model = lambda: LeNet(data_format="channels_first") graph = LeNet.graph().load() model_fn = partial(model_fn_with_fetch_hook, create_model=create_model, graph=graph) predicted_label = predict( create_model=create_model, input_fn=lambda: mnist.test(data_dir).filter( lambda image, label: tf.equal( tf.convert_to_tensor(class_id, dtype=tf.int32), label)) .skip(image_id).take(1).batch(1), model_dir=ckpt_dir, ) if predicted_label != class_id: return [{}] if per_node else {} adversarial_example = lenet_mnist_example( attack_name=attack_name, attack_fn=attack_fn, generate_adversarial_fn=generate_adversarial_fn, class_id=class_id, image_id=image_id, # model_dir not ckpt_dir model_dir=model_dir, transforms=transforms, transform_name=transform_name, ).load() if adversarial_example is None: return [{}] if per_node else {} adversarial_predicted_label = predict( create_model=create_model, input_fn=lambda: tf.data.Dataset.from_tensors( mnist.normalize(adversarial_example)), model_dir=ckpt_dir, ) if predicted_label == adversarial_predicted_label: return [{}] if per_node else {} trace = reconstruct_trace_from_tf( class_id=class_id, model_fn=model_fn, input_fn=lambda: mnist.test(data_dir, transforms=transforms). filter(lambda image, label: tf.equal( tf.convert_to_tensor(class_id, dtype=tf.int32), label) ).skip(image_id).take(1).batch(1), select_fn=select_fn, model_dir=ckpt_dir, per_channel=per_channel, )[0] if trace is None: return [{}] if per_node else {} adversarial_trace = reconstruct_trace_from_tf_brute_force( model_fn=model_fn, input_fn=lambda: tf.data.Dataset.from_tensors( mnist.normalize(adversarial_example)), select_fn=select_fn, model_dir=ckpt_dir, per_channel=per_channel, )[0] adversarial_label = adversarial_trace.attrs[GraphAttrKey.PREDICT] # adversarial_pred, violation = \ # detect_by_reduced_edge_count_violation( # class_trace_fn(adversarial_label).load(), # adversarial_trace, # reduce_mode, # ) # row = { # "image_id": image_id, # "class_id": class_id, # "original.prediction": # detect_by_reduced_edge(class_trace_fn(class_id).load(), # trace, # reduce_mode, # ), # "adversarial.prediction": # adversarial_pred, # "violation": violation, # } original_pred, violation = \ detect_by_reduced_edge_count_violation( class_trace_fn(class_id).load(), trace, reduce_mode, ) row = { "image_id": image_id, "class_id": class_id, "original.prediction": original_pred, "adversarial.prediction": detect_by_reduced_edge( class_trace_fn(adversarial_label).load(), adversarial_trace, reduce_mode, ), "violation": violation, } return row
def train( batch_size: int = 64, train_epochs: int = 10, data_format: str = "channels_first", multi_gpu: bool = False, label: str = None, ): tf.logging.set_verbosity(tf.logging.INFO) if label is None: model_dir = abspath(f"{root_dir}/model") else: model_dir = abspath(f"{root_dir}/model_{label}") ckpt_dir = os.path.join(model_dir, "ckpts") data_dir = abspath(MNIST_PATH) model_function = model_fn if multi_gpu: validate_batch_size_for_multi_gpu(batch_size) # There are two steps required if using multi-GPU: (1) wrap the model_fn, # and (2) wrap the optimizer. The first happens here, and (2) happens # in the model_fn itself when the optimizer is defined. model_function = tf.contrib.estimator.replicate_model_fn( model_fn, loss_reduction=tf.losses.Reduction.MEAN) if data_format is None: data_format = ("channels_first" if tf.test.is_built_with_cuda() else "channels_last") estimator_config = tf.estimator.RunConfig( keep_checkpoint_max=None, session_config=new_session_config(parallel=0)) classifier = tf.estimator.Estimator( model_fn=model_function, model_dir=ckpt_dir, params={ "data_format": data_format, "multi_gpu": multi_gpu, "batch_size": batch_size, }, config=estimator_config, ) for epoch in range(train_epochs): # Train the model def train_input_fn(): # When choosing shuffle buffer sizes, larger sizes result in better # randomness, while smaller sizes use less memory. MNIST is a small # enough dataset that we can easily shuffle the full epoch. ds = mnist.train( data_dir, transforms=transforms, ) ds = ds.cache().shuffle(buffer_size=60000).batch(batch_size) return ds # Set up training hook that logs the training accuracy every 100 steps. tensors_to_log = {"train_accuracy": "train_accuracy"} logging_hook = tf.train.LoggingTensorHook(tensors=tensors_to_log, every_n_iter=100) classifier.train(input_fn=train_input_fn, hooks=[logging_hook]) # Evaluate the model and print results def eval_input_fn(): return (mnist.test(data_dir, ).batch( batch_size).make_one_shot_iterator().get_next()) eval_results = classifier.evaluate(input_fn=eval_input_fn) print(label) print("Evaluation results:\n\t%s" % eval_results) print()
def get_row(class_id: int, image_id: int) -> Dict[str, Any]: nonlocal model_dir mode.check(False) data_dir = abspath(MNIST_PATH) model_dir = abspath(model_dir) ckpt_dir = f"{model_dir}/ckpts" create_model = lambda: LeNet(data_format="channels_first") graph = LeNet.graph().load() model_fn = partial(model_fn_with_fetch_hook, create_model=create_model, graph=graph) if mnist_dataset_mode == "test": dataset = mnist.test elif mnist_dataset_mode == "train": dataset = mnist.train else: raise RuntimeError("Dataset invalid") predicted_label = predict( create_model=create_model, # dataset may be train or test, should be consistent with lenet_mnist_example input_fn=lambda: dataset(data_dir, ).filter( lambda image, label: tf.equal( tf.convert_to_tensor(class_id, dtype=tf.int32), label)) .skip(image_id).take(1).batch(1), model_dir=ckpt_dir, ) if predicted_label != class_id: return [{}] if per_node else {} adversarial_example = lenet_mnist_example( attack_name=attack_name, attack_fn=attack_fn, generate_adversarial_fn=generate_adversarial_fn, class_id=class_id, image_id=image_id, # model_dir not ckpt_dir model_dir=model_dir, transforms=transforms, transform_name=transform_name, mode=mnist_dataset_mode, ).load() if adversarial_example is None: return [{}] if per_node else {} adversarial_predicted_label = predict( create_model=create_model, input_fn=lambda: tf.data.Dataset.from_tensors( mnist.normalize(adversarial_example)), model_dir=ckpt_dir, ) if predicted_label == adversarial_predicted_label: return [{}] if per_node else {} if use_class_trace: class_trace_avg = ClassTraceIOAction(predicted_label).load() else: class_trace_avg = None trace = reconstruct_trace_from_tf( class_id=class_id, model_fn=model_fn, input_fn=lambda: dataset( data_dir, transforms=transforms, ).filter(lambda image, label: tf.equal( tf.convert_to_tensor(class_id, dtype=tf.int32), label)). skip(image_id).take(1).batch(1), select_fn=select_fn, model_dir=ckpt_dir, per_channel=per_channel, class_trace=class_trace_avg, )[0] if trace is None: return [{}] if per_node else {} path = os.path.join(save_dir, f"original_{transform_name}", f"{class_id}", f"{image_id}.pkl") ensure_dir(path) with open(path, "wb") as f: pickle.dump(trace, f) adversarial_trace = reconstruct_trace_from_tf( model_fn=model_fn, input_fn=lambda: tf.data.Dataset.from_tensors( mnist.normalize(adversarial_example)), select_fn=select_fn, model_dir=ckpt_dir, per_channel=per_channel, class_trace=class_trace_avg, )[0] path = os.path.join(save_dir, f"{attack_name}_{transform_name}", f"{class_id}", f"{image_id}.pkl") ensure_dir(path) with open(path, "wb") as f: pickle.dump(adversarial_trace, f) row = { "class_id": class_id, "image_id": image_id, "trace": trace, "adversarial_trace": adversarial_trace, } # row = calc_all_overlap( # class_trace_fn(class_id).load(), adversarial_trace, overlap_fn # ) return row
def forward_propagate( create_model, input_fn, forward_fn: Callable[[tf.Tensor], tf.Tensor], model_dir: str, data_format: str = "channels_first", parallel: int = 1, prediction_hooks: List[SessionRunHook] = None, tag: str = "noop", ) -> Union[int, float]: def model_fn(features, labels, mode, params): image = features if isinstance(image, dict): image = features["image"] # Save inputs for visualization in tensorboard nonlocal prediction_hooks feature_trans = tf.transpose(image, perm=[0,2,3,1]) image_summary = tf.summary.image(f"input_{tag}", feature_trans, max_outputs=100) eval_summary_hook = tf.train.SummarySaverHook( summary_op=image_summary, save_steps=2, output_dir=result_dir) if prediction_hooks is None: prediction_hooks = [eval_summary_hook] else: prediction_hooks.append(eval_summary_hook) if mode == tf.estimator.ModeKeys.PREDICT: logits = create_model()(image, training=False) predictions = {"classes": forward_fn(logits)} return tf.estimator.EstimatorSpec( mode=tf.estimator.ModeKeys.PREDICT, predictions=predictions, prediction_hooks=prediction_hooks, export_outputs={ "classify": tf.estimator.export.PredictOutput(predictions) }, ) model_dir = abspath(model_dir) model_function = model_fn if data_format is None: data_format = ( "channels_first" if tf.test.is_built_with_cuda() else "channels_last" ) estimator_config = tf.estimator.RunConfig( session_config=new_session_config(parallel=parallel) ) if not os.path.exists(model_dir): raise RuntimeError(f"model directory {model_dir} is not existed") # st() classifier = tf.estimator.Estimator( model_fn=model_function, model_dir=model_dir, params={"data_format": data_format}, config=estimator_config, ) result = list(classifier.predict(input_fn=input_fn)) return result[0]["classes"]
def save_graph( attack_name: str, attack_fn, generate_adversarial_fn, select_fn: Callable[[np.ndarray], np.ndarray], image_id_index, batch_size, class_id, # model_dir = "result/lenet/model_augmentation", model_dir=model_dir, transforms=None, transform_name="noop", graph_dir="result/test", dataset_mode=dataset_mode, images_per_class=1, compute_adversarial=True, **kwargs, ): data_dir = abspath(MNIST_PATH) model_dir = abspath(model_dir) ckpt_dir = f"{model_dir}/ckpts" create_model = lambda: LeNet(data_format="channels_first") graph = LeNet.graph().load() batch_size = min(batch_size, images_per_class - image_id_index) if dataset_mode == "test": dataset = mnist.test elif dataset_mode == "train": dataset = mnist.train else: raise RuntimeError("Dataset invalid") predicted_label = predict_batch( create_model=create_model, input_fn=lambda: (dataset(data_dir).filter(lambda image, label: tf.equal( tf.convert_to_tensor(class_id, dtype=tf.int32), label)).skip( image_id_index).take(batch_size).batch(batch_size). make_one_shot_iterator().get_next()[0]), model_dir=ckpt_dir, ) prediction_valid = (predicted_label == class_id) # The shape of each example is (1, 32, 32, 3) adversarial_examples = [ lenet_mnist_example( attack_name=attack_name, attack_fn=attack_fn, generate_adversarial_fn=generate_adversarial_fn, class_id=class_id, image_id=image_id, # model_dir not ckpt_dir model_dir=model_dir, transforms=transforms, transform_name=transform_name, mode=dataset_mode, ).load() for image_id in range(image_id_index, image_id_index + batch_size) ] adversarial_valid = np.array( [example is not None for example in adversarial_examples]) adversarial_examples = [ example if example is not None else np.zeros((1, 1, 28, 28)) for example in adversarial_examples ] adversarial_examples = np.squeeze(np.array(adversarial_examples).astype( np.float32), axis=1) # adversarial_example is [0, 1] of shape (1, 28, 28) adversarial_predicted_label = predict_batch( create_model=create_model, input_fn=lambda: tf.data.Dataset.from_tensors( mnist.normalize(adversarial_examples)), model_dir=ckpt_dir, ) adversarial_prediction_valid = adversarial_predicted_label != class_id batch_valid = (prediction_valid * adversarial_valid * adversarial_prediction_valid) original_graph_dir = os.path.join(graph_dir, f"original_{transform_name}", f"{class_id}") original_graph_saver = IOBatchAction( dir=original_graph_dir, root_index=image_id_index, ) original_model_fn = partial( model_fn_with_fetch_hook, create_model=create_model, graph=graph, graph_saver=original_graph_saver, batch_valid=batch_valid, ) trace = reconstruct_trace_from_tf( class_id=class_id, model_fn=original_model_fn, input_fn=lambda: (dataset( data_dir, transforms=transforms, ).filter(lambda image, label: tf.equal( tf.convert_to_tensor(class_id, dtype=tf.int32), label)).skip( image_id_index).take(batch_size).batch(batch_size). make_one_shot_iterator().get_next()[0]), select_fn=select_fn, model_dir=ckpt_dir, ) if compute_adversarial: adversarial_graph_dir = os.path.join( graph_dir, f"{attack_name}_{transform_name}", f"{class_id}") adversarial_graph_saver = IOBatchAction( dir=adversarial_graph_dir, root_index=image_id_index, ) adversarial_model_fn = partial( model_fn_with_fetch_hook, create_model=create_model, graph=graph, graph_saver=adversarial_graph_saver, batch_valid=batch_valid, ) adversarial_trace = reconstruct_trace_from_tf( model_fn=adversarial_model_fn, input_fn=lambda: tf.data.Dataset.from_tensors( mnist.normalize(adversarial_examples)), select_fn=select_fn, model_dir=ckpt_dir, )
def predict_original_adversarial( attack_name: str, attack_fn, generate_adversarial_fn, image_id_index, batch_size, class_id, model_dir=model_dir, transforms=None, transform_name="noop", graph_dir="result/test", dataset_mode=dataset_mode, images_per_class=1, **kwargs, ): # mode.check(False) data_dir = abspath(CIFAR10_PATH) model_dir = abspath(model_dir) ckpt_dir = f"{model_dir}/ckpts" create_model = lambda: partial( ResNet10Cifar10(), training=False, ) graph = ResNet10Cifar10.graph().load() batch_size = min(batch_size, images_per_class - image_id_index) predicted_label = predict_batch( create_model=create_model, input_fn=lambda: (input_fn_for_adversarial_examples( is_training=(dataset_mode == "train"), data_dir=data_dir, num_parallel_batches=1, is_shuffle=False, transform_fn=None, ).filter(lambda image, label: tf.equal( tf.convert_to_tensor(class_id, dtype=tf.int32), label)).skip( image_id_index).take(batch_size).batch(batch_size). make_one_shot_iterator().get_next()[0]), model_dir=ckpt_dir, ) # The shape of each example is (1, 32, 32, 3) adversarial_examples = [ resnet10_cifar10_example( attack_name=attack_name, attack_fn=attack_fn, generate_adversarial_fn=generate_adversarial_fn, class_id=class_id, image_id=image_id, # model_dir not ckpt_dir model_dir=model_dir, transforms=transforms, transform_name=transform_name, dataset_mode=dataset_mode, ).load() for image_id in range(image_id_index, image_id_index + batch_size) ] adversarial_valid = np.array( [example is not None for example in adversarial_examples]) adversarial_examples = [ example if example is not None else np.zeros((1, 32, 32, 3)) for example in adversarial_examples ] adversarial_examples = np.squeeze(np.array(adversarial_examples).astype( np.float32), axis=1) # adversarial_example is [0, 1] of shape (1, 32, 32, 3) adversarial_predicted_label = predict_batch( create_model=create_model, input_fn=lambda: tf.data.Dataset.from_tensors(adversarial_examples), model_dir=ckpt_dir, ) assert predicted_label.shape == adversarial_predicted_label.shape original_correct = (predicted_label[adversarial_valid] == class_id).sum() adversarial_correct = ( adversarial_predicted_label[adversarial_valid] == class_num).sum() valid_count = adversarial_valid.sum() return original_correct, adversarial_correct, valid_count