robustml interface

Author: dtsip

model_robustml.py

@@ -0,0 +1,45 @@
+import robustml
+import tensorflow as tf
+
+import model
+
+class Model(robustml.model.Model):
+  def __init__(self, sess):
+    self._model = model.Model('eval')
+
+    saver = tf.train.Saver()
+    checkpoint = tf.train.latest_checkpoint('models/model_0')
+    saver.restore(sess, checkpoint)
+
+    self._sess = sess
+    self._input = self._model.x_input
+    self._logits = self._model.pre_softmax
+    self._predictions = self._model.predictions
+    self._dataset = robustml.dataset.CIFAR10()
+    self._threat_model = robustml.threat_model.Linf(epsilon=0.03)
+
+  @property
+  def dataset(self):
+      return self._dataset
+
+  @property
+  def threat_model(self):
+      return self._threat_model
+
+  def classify(self, x):
+      return self._sess.run(self._predictions,
+                            {self._input: x})[0]
+
+  # expose attack interface
+
+  @property
+  def input(self):
+      return self._input
+
+  @property
+  def logits(self):
+      return self._logits
+
+  @property
+  def predictions(self):
+      return self._predictions