def get_estimator(epochs=10, batch_size=50, epsilon=0.04, save_dir=tempfile.mkdtemp()):
train_data, eval_data = cifar10.load_data()
test_data = eval_data.split(0.5)
pipeline = fe.Pipeline(
train_data=train_data,
eval_data=eval_data,
test_data=test_data,
batch_size=batch_size,
ops=[Normalize(inputs="x", outputs="x", mean=(0.4914, 0.4822, 0.4465), std=(0.2471, 0.2435, 0.2616))])
clean_model = fe.build(model_fn=lambda: LeNet(input_shape=(32, 32, 3)),
optimizer_fn="adam",
model_name="clean_model")
clean_network = fe.Network(ops=[
Watch(inputs="x"),
ModelOp(model=clean_model, inputs="x", outputs="y_pred"),
CrossEntropy(inputs=("y_pred", "y"), outputs="base_ce"),
FGSM(data="x", loss="base_ce", outputs="x_adverse", epsilon=epsilon, mode="!train"),
ModelOp(model=clean_model, inputs="x_adverse", outputs="y_pred_adv", mode="!train"),
UpdateOp(model=clean_model, loss_name="base_ce")
])
clean_traces = [
Accuracy(true_key="y", pred_key="y_pred", output_name="clean_accuracy"),
Accuracy(true_key="y", pred_key="y_pred_adv", output_name="adversarial_accuracy"),
BestModelSaver(model=clean_model, save_dir=save_dir, metric="base_ce", save_best_mode="min"),
]
clean_estimator = fe.Estimator(pipeline=pipeline,
network=clean_network,
epochs=epochs,
traces=clean_traces,
log_steps=1000)
return clean_estimator
def test_saliency(self):
fe.estimator.enable_deterministic(200)
label_mapping = {
'airplane': 0,
'automobile': 1,
'bird': 2,
'cat': 3,
'deer': 4,
'dog': 5,
'frog': 6,
'horse': 7,
'ship': 8,
'truck': 9
}
batch_size = 32
train_data, eval_data = cifar10.load_data()
pipeline = fe.Pipeline(test_data=train_data,
batch_size=batch_size,
ops=[Normalize(inputs="x", outputs="x")],
num_process=0)
weight_path = os.path.abspath(
os.path.join(__file__, "..", "resources", "lenet_cifar10_tf.h5"))
model = fe.build(model_fn=lambda: LeNet(input_shape=(32, 32, 3)),
optimizer_fn="adam",
weights_path=weight_path)
network = fe.Network(
ops=[ModelOp(model=model, inputs="x", outputs="y_pred")])
save_dir = tempfile.mkdtemp()
traces = [
Saliency(model=model,
model_inputs="x",
class_key="y",
model_outputs="y_pred",
samples=5,
label_mapping=label_mapping),
ImageSaver(inputs="saliency", save_dir=save_dir)
]
estimator = fe.Estimator(pipeline=pipeline,
network=network,
epochs=5,
traces=traces,
log_steps=1000)
estimator.test()
ans_img_path = os.path.abspath(
os.path.join(__file__, "..", "resources", "saliency_figure.png"))
ans_img = img_to_rgb_array(ans_img_path)
output_img_path = os.path.join(save_dir, "saliency_test_epoch_5.png")
output_img = img_to_rgb_array(output_img_path)
self.assertTrue(check_img_similar(output_img, ans_img))
def get_pipeline_data(batch_size=128):
train_data, _ = cifar10.load_data()
pipeline = fe.Pipeline(
train_data=train_data,
batch_size=batch_size,
ops=[
Normalize(inputs="x", outputs="x", mean=(0.4914, 0.4822, 0.4465), std=(0.2471, 0.2435, 0.2616)),
])
result = pipeline.get_results()
return result
def get_estimator(epsilon=0.04,
epochs=10,
batch_size=32,
max_train_steps_per_epoch=None,
max_eval_steps_per_epoch=None,
save_dir=tempfile.mkdtemp()):
# step 1
train_data, eval_data = cifar10.load_data()
test_data = eval_data.split(0.5)
pipeline = fe.Pipeline(train_data=train_data,
eval_data=eval_data,
test_data=test_data,
batch_size=batch_size,
ops=[
Normalize(inputs="x",
outputs="x",
mean=(0.4914, 0.4822, 0.4465),
std=(0.2471, 0.2435, 0.2616)),
ChannelTranspose(inputs="x", outputs="x"),
])
# step 2
model = fe.build(model_fn=lambda: LeNet(input_shape=(3, 32, 32)),
optimizer_fn="adam")
network = fe.Network(ops=[
Watch(inputs="x"),
ModelOp(model=model, inputs="x", outputs="y_pred"),
CrossEntropy(inputs=("y_pred", "y"), outputs="base_ce"),
FGSM(data="x", loss="base_ce", outputs="x_adverse", epsilon=epsilon),
ModelOp(model=model, inputs="x_adverse", outputs="y_pred_adv"),
CrossEntropy(inputs=("y_pred_adv", "y"), outputs="adv_ce"),
Average(inputs=("base_ce", "adv_ce"), outputs="avg_ce"),
UpdateOp(model=model, loss_name="avg_ce")
])
# step 3
traces = [
Accuracy(true_key="y", pred_key="y_pred", output_name="base accuracy"),
Accuracy(true_key="y",
pred_key="y_pred_adv",
output_name="adversarial accuracy"),
BestModelSaver(model=model,
save_dir=save_dir,
metric="adv_ce",
save_best_mode="min"),
]
estimator = fe.Estimator(
pipeline=pipeline,
network=network,
epochs=epochs,
traces=traces,
max_train_steps_per_epoch=max_train_steps_per_epoch,
max_eval_steps_per_epoch=max_eval_steps_per_epoch,
monitor_names=["base_ce", "adv_ce"])
return estimator
def __init__(self, data_path, batch_size, **kwargs):
"""
Initialize.
Args:
data_path: File path for the dataset
batch_size (int): The batch size for the data loader
**kwargs: Additional arguments, passed to super init and
load_mnist_shard
"""
# TODO: We should be downloading the dataset shard into a directory
# TODO: There needs to be a method to ask how many collaborators and
# what index/rank is this collaborator.
# Then we have a way to automatically shard based on rank and size
# of collaborator list.
train_data, eval_data = cifar10.load_data()
test_data = eval_data.split(0.5)
collaborator_count = kwargs['collaborator_count']
train_data, eval_data, test_data = self.split_data(
train_data,
eval_data,
test_data,
int(data_path),
collaborator_count
)
print(f"train_data = {train_data}")
print(f"eval_data = {eval_data}")
print(f"test_data = {test_data}")
print(f"batch_size = {batch_size}")
super().__init__(fe.Pipeline(
train_data=train_data,
eval_data=eval_data,
test_data=test_data,
batch_size=batch_size,
ops=[
Normalize(inputs="x", outputs="x",
mean=(0.4914, 0.4822, 0.4465),
std=(0.2471, 0.2435, 0.2616))
]), **kwargs)
def get_estimator(epochs=24, batch_size=512, max_train_steps_per_epoch=None, save_dir=tempfile.mkdtemp()):
# step 1: prepare dataset
train_data, test_data = load_data()
pipeline = fe.Pipeline(
train_data=train_data,
test_data=test_data,
batch_size=batch_size,
ops=[
Normalize(inputs="x", outputs="x", mean=(0.4914, 0.4822, 0.4465), std=(0.2471, 0.2435, 0.2616)),
PadIfNeeded(min_height=40, min_width=40, image_in="x", image_out="x", mode="train"),
RandomCrop(32, 32, image_in="x", image_out="x", mode="train"),
Sometimes(HorizontalFlip(image_in="x", image_out="x", mode="train")),
CoarseDropout(inputs="x", outputs="x", mode="train", max_holes=1),
ChannelTranspose(inputs="x", outputs="x"),
Onehot(inputs="y", outputs="y", mode="train", num_classes=10, label_smoothing=0.2)
])
# step 2: prepare network
model = fe.build(model_fn=FastCifar, optimizer_fn="adam")
network = fe.Network(ops=[
ModelOp(model=model, inputs="x", outputs="y_pred"),
CrossEntropy(inputs=("y_pred", "y"), outputs="ce"),
UpdateOp(model=model, loss_name="ce")
])
# step 3 prepare estimator
traces = [
Accuracy(true_key="y", pred_key="y_pred"),
BestModelSaver(model=model, save_dir=save_dir, metric="accuracy", save_best_mode="max"),
LRScheduler(model=model, lr_fn=lr_schedule)
]
estimator = fe.Estimator(pipeline=pipeline,
network=network,
epochs=epochs,
traces=traces,
max_train_steps_per_epoch=max_train_steps_per_epoch)
return estimator
def get_estimator(epsilon=0.04,
epochs=20,
batch_size=32,
code_length=16,
max_train_steps_per_epoch=None,
max_eval_steps_per_epoch=None,
save_dir=tempfile.mkdtemp()):
# step 1
train_data, eval_data = cifar10.load_data()
test_data = eval_data.split(0.5)
pipeline = fe.Pipeline(train_data=train_data,
eval_data=eval_data,
test_data=test_data,
batch_size=batch_size,
ops=[
Normalize(inputs="x",
outputs="x",
mean=(0.4914, 0.4822, 0.4465),
std=(0.2471, 0.2435, 0.2616)),
Hadamard(inputs="y",
outputs="y_code",
n_classes=10)
])
# step 2
model = fe.build(model_fn=lambda: ecc_lenet(code_length=code_length),
optimizer_fn='adam')
network = fe.Network(ops=[
Watch(inputs="x", mode=('eval', 'test')),
ModelOp(model=model, inputs="x", outputs="y_pred_code"),
Hinge(inputs=("y_pred_code", "y_code"), outputs="base_hinge"),
UpdateOp(model=model, loss_name="base_hinge"),
UnHadamard(inputs="y_pred_code",
outputs="y_pred",
n_classes=10,
mode=('eval', 'test')),
# The adversarial attack:
FGSM(data="x",
loss="base_hinge",
outputs="x_adverse_hinge",
epsilon=epsilon,
mode=('eval', 'test')),
ModelOp(model=model,
inputs="x_adverse_hinge",
outputs="y_pred_adv_hinge_code",
mode=('eval', 'test')),
Hinge(inputs=("y_pred_adv_hinge_code", "y_code"),
outputs="adv_hinge",
mode=('eval', 'test')),
UnHadamard(inputs="y_pred_adv_hinge_code",
outputs="y_pred_adv_hinge",
n_classes=10,
mode=('eval', 'test')),
])
# step 3
traces = [
Accuracy(true_key="y", pred_key="y_pred", output_name="base_accuracy"),
Accuracy(true_key="y",
pred_key="y_pred_adv_hinge",
output_name="adversarial_accuracy"),
BestModelSaver(model=model,
save_dir=save_dir,
metric="base_hinge",
save_best_mode="min",
load_best_final=True)
]
estimator = fe.Estimator(
pipeline=pipeline,
network=network,
epochs=epochs,
traces=traces,
max_train_steps_per_epoch=max_train_steps_per_epoch,
max_eval_steps_per_epoch=max_eval_steps_per_epoch,
monitor_names=["adv_hinge"])
return estimator
