From 9c570b589ab5ac9c6b6b75800651954d316ffb10 Mon Sep 17 00:00:00 2001
From: beta <beta.underalpha@gmail.com>
Date: Wed, 25 May 2022 20:46:00 +0900
Subject: [PATCH] (#3) Defense: Refine exformer

---
 src/attack/attack.py   |  4 ++--
 src/defense/models.py  | 54 ++++++++++++++++++++++++++++++++++++++++--
 src/defense/process.py |  4 ++--
 src/defense/train.py   |  6 ++---
 4 files changed, 59 insertions(+), 9 deletions(-)

diff --git a/src/attack/attack.py b/src/attack/attack.py
index fe85d7b..8ae477f 100644
--- a/src/attack/attack.py
+++ b/src/attack/attack.py
@@ -2,7 +2,7 @@
 from typings.models import Attack, Defense
 from utils.logging import concat_batch_images
 from models import FgsmMnist, FgsmCifar, PgdMnist, PgdCifar, Cw, NormalNoise
-from defense.models import Reformer, Denoiser, Motd, ExMotd
+from defense.models import Reformer, Exformer, Denoiser, Motd, ExMotd
 from victim.models import Classifier
 
 from utils.dataset import Mnist, Cifar10
@@ -111,7 +111,7 @@
         defense_model.compile()
         defense_model.load()
     elif args.defense == "exformer":
-        defense_model = Reformer(
+        defense_model = Exformer(
             f"defense_exformer_{args.dataset}",
             input_shape=input_shape,
             intensity=args.intensity[0],
diff --git a/src/defense/models.py b/src/defense/models.py
index b28885d..1f54f70 100644
--- a/src/defense/models.py
+++ b/src/defense/models.py
@@ -18,7 +18,7 @@ def __init__(
         data_train: tf.data.Dataset = None,
         data_test: tf.data.Dataset = None,
         optimizer: keras.optimizers.Optimizer = keras.optimizers.Adam(),
-        loss: keras.losses.Loss = keras.losses.MeanSquaredError(),
+        loss: keras.losses.Loss = keras.losses.BinaryCrossentropy(),
         accuracy: keras.metrics.Accuracy = keras.metrics.CategoricalAccuracy(
             name="accuracy"
         ),
@@ -36,6 +36,7 @@ def __init__(
             accuracy,
             checkpoint_filepath,
             tensorboard_log_path,
+            is_functional=True,
         )
 
     def _model(self) -> keras.Model:
@@ -90,6 +91,55 @@ def predict(epoch, logs):
         return [reduce_lr, keras.callbacks.LambdaCallback(on_epoch_end=predict)]
 
 
+class Exformer(Reformer):
+    def _model(self) -> keras.Model:
+        """
+        Reference: https://www.kaggle.com/code/tarunk04/autoencoder-denoising-image-mnist-cifar10/notebook#Denoising-Cifar10-Data
+        Denoising Autoencoder with Skip Connection
+        """
+
+        inputs = keras.layers.Input(shape=self.input_shape())
+        # Encoder - 1
+        x = keras.layers.Conv2D(32, 3, activation="relu", padding="same")(inputs)
+        x = keras.layers.BatchNormalization()(x)
+        x = keras.layers.MaxPool2D()(x)
+        x = keras.layers.Dropout(0.5)(x)
+        # Encoder - 2
+        skip = keras.layers.Conv2D(32, 3, padding="same")(x)
+        x = keras.layers.LeakyReLU()(skip)
+        x = keras.layers.BatchNormalization()(x)
+        x = keras.layers.MaxPool2D()(x)
+        x = keras.layers.Dropout(0.5)(x)
+        # Encoder - Finalize
+        x = keras.layers.Conv2D(64, 3, activation="relu", padding="same")(x)
+        x = keras.layers.BatchNormalization()(x)
+        encoded = keras.layers.MaxPool2D()(x)
+
+        # Decoder - 1
+        x = keras.layers.Conv2DTranspose(
+            64, 3, activation="relu", strides=(2, 2), padding="same"
+        )(encoded)
+        x = keras.layers.BatchNormalization()(x)
+        x = keras.layers.Dropout(0.5)(x)
+        # Decoder - 2
+        x = keras.layers.Conv2DTranspose(
+            32, 3, activation="relu", strides=(2, 2), padding="same"
+        )(x)
+        x = keras.layers.BatchNormalization()(x)
+        x = keras.layers.Dropout(0.5)(x)
+        # Decoder - 3
+        x = keras.layers.Conv2DTranspose(32, 3, padding="same")(x)
+        x = keras.layers.Add()([skip, x])
+        x = keras.layers.LeakyReLU()(x)
+        x = keras.layers.BatchNormalization()(x)
+        # Decoder - Finalize
+        decoded = keras.layers.Conv2DTranspose(
+            3, 3, activation="sigmoid", strides=(2, 2), padding="same"
+        )(x)
+
+        return keras.Model(inputs, decoded)
+
+
 class Denoiser(Defense):
     def _model(self) -> keras.Model:
         return keras.Sequential([SlqLayer()])
@@ -158,7 +208,7 @@ def __init__(
             input_shape=input_shape,
             intensity=intensities[0],
         )
-        self.exformer = Reformer(
+        self.exformer = Exformer(
             f"defense_exformer_{dataset}",
             input_shape=input_shape,
             intensity=intensities[1],
diff --git a/src/defense/process.py b/src/defense/process.py
index 5896380..05e69d7 100644
--- a/src/defense/process.py
+++ b/src/defense/process.py
@@ -1,5 +1,5 @@
 from utils.logging import concat_batch_images
-from defense.models import Reformer, Denoiser, Motd, ExMotd
+from defense.models import Reformer, Exformer, Denoiser, Motd, ExMotd
 
 import tensorflow as tf
 
@@ -60,7 +60,7 @@
             intensity=args.intensity[0],
         )
     elif args.defense == "exformer":
-        defense_model = Reformer(
+        defense_model = Exformer(
             f"defense_exformer_{args.dataset}",
             input_shape=input_shape,
             intensity=args.intensity[0],
diff --git a/src/defense/train.py b/src/defense/train.py
index 10292ba..79f6df3 100644
--- a/src/defense/train.py
+++ b/src/defense/train.py
@@ -1,5 +1,5 @@
 from utils.dataset import NoisyMnist, NoisyCifar10, ExCifar10
-from models import Reformer
+from models import Reformer, Exformer
 
 import argparse
 import tensorflow as tf
@@ -33,14 +33,14 @@ def train_cifar10_reformer(epochs: int = 100):
 
 def train_excifar10_exformer(epochs: int = 100):
     train_set, test_set = ExCifar10().dataset()
-    reformer = Reformer(
+    exformer = Exformer(
         "defense_exformer_cifar10",
         (32, 32, 3),
         1.0,
         train_set,
         test_set,
     )
-    reformer.train(epochs)
+    exformer.train(epochs)
 
 
 if __name__ == "__main__":