Move computations to tensorflow and vectorize to first degree

Author: MahanFathi

agent.py

@@ -5,11 +5,14 @@
 from net import Net
 
 
-def do_plot(guided_backprops, indices, image, bounding_box, files_dict):
-    top_grads = [guided_backprops[i] for i in indices]
-    folder = list(files_dict.keys())[0]
-    name = files_dict[folder][0]
-    get_full_plot(top_grads, image, bounding_box, folder, name)
+def do_batch_plot(guided_backprops, indices, images, bounding_boxes, files_dict):
+    folders_and_names = []
+    for key, values in files_dict.items():
+        for val in values:
+            folders_and_names.append([key, val])
+    for gbps, index, image, bb, fn in zip(guided_backprops, indices, images, bounding_boxes, folders_and_names):
+        top_grads = [gbps[i] for i in index]
+        get_full_plot(top_grads, image, bb, *fn)
 
 
 class Agent(object):
@@ -23,17 +26,17 @@ def __init__(self, config):
 
     def get_bounding_box(self, image=None, files_dict=None):
         if image is None and files_dict is None:
-            image, files_dict = sample_data(1)
+            image, files_dict = sample_data(self.config.batch_size)
         preprocessed_image = self._preprocess(image.copy())
-        kmax_neuron_indices, top_class = self._get_kmax_neurons(preprocessed_image)
+        kmax_neuron_indices, top_classes_indices = self.net.get_top_kmax_neurons(preprocessed_image)
         guided_backprops = self._get_guided_backprops(preprocessed_image, kmax_neuron_indices)
-        masks = self._get_image_masks(guided_backprops)
-        topk_neurons_relative_indices = self._get_topk_neurons(preprocessed_image, masks, top_class)
-        bounding_box, _ = self._get_bounding_box(masks, topk_neurons_relative_indices)
+        masks = self._get_images_masks(guided_backprops)
+        top_k_neurons_relative_indices = self._get_top_k_neurons(preprocessed_image, masks, top_classes_indices)
+        bounding_boxes = self._get_all_bounding_boxes(masks, top_k_neurons_relative_indices)
         if self.config.do_plotting:
-            do_plot(guided_backprops, topk_neurons_relative_indices,
-                    image, bounding_box, files_dict)
-        return bounding_box
+            do_batch_plot(guided_backprops, top_k_neurons_relative_indices,
+                          image, bounding_boxes, files_dict)
+        return bounding_boxes
 
     def make_tsne_pic_for_directory(self, folder='personal'):
         fc_features = None
@@ -48,52 +51,49 @@ def make_tsne_pic_for_directory(self, folder='personal'):
 
         tsne = T_SNE()
         tsne_embedding = tsne.generate_tsne(fc_features)
-        tsne.save_grid(images, tsne_embedding, folder+'.jpg')
+        tsne.save_grid(images, tsne_embedding, folder + '.jpg')
 
     def _preprocess(self, image):
         return self.net.vgg16.preprocess_input(image)
 
-    def _get_kmax_neurons(self, image):
-        class_scores = self.net.get_class_scores(image)
-        n = len(self.config.top_n_classes_weights)
-        top_class = np.argmax(class_scores)
-        top_classes = np.argsort(class_scores)[-n:]
-        top_scores = class_scores[top_classes]
-        activations = self.net.get_activations(image)
-
-        normalizer = 1.0 / (np.sum(self.config.top_n_classes_weights) * np.sum(top_scores))
-
-        impact_gradient = 0
-        for i, class_index in enumerate(top_classes):
-            impact_gradient_per_class = self.net.get_impact_gradient_per_class(activations=activations,
-                                                                               class_index=class_index)
-            impact_gradient += impact_gradient_per_class * top_scores[i] * \
-                               self.config.top_n_classes_weights[i] * normalizer
-
-        rank_score = activations * impact_gradient
-        kmax_neurons = np.argsort(rank_score.ravel())[-self.config.kmax:]
-        return kmax_neurons, top_class
-
-    def _get_guided_backprops(self, image, neuron_indices):
-        guided_backprops = [self.net.get_guided_backprop(image, neuron_index) for neuron_index in neuron_indices]
+    def _get_guided_backprops(self, images, neuron_indices):
+        # lazy programming, this part should as well be vectorized
+        guided_backprops = []
+        for image, neuron_index in zip(images, neuron_indices):
+            gbp = [self.net.get_guided_backprop(np.expand_dims(image, axis=0), ni)
+                   for ni in neuron_index]
+            guided_backprops.append(gbp)
         return guided_backprops
 
-    def _get_image_masks(self, guided_backprops):
-        projected_gbps = [np.max(gb, axis=-1).squeeze() for gb in guided_backprops]
-        raw_masks = [pgbp > np.percentile(pgbp, self.config.cut_off_percentile) for pgbp in projected_gbps]
-        # erosion and dilation
-        masks = [binary_dilation(binary_erosion(raw_mask)).astype(projected_gbps[0].dtype) for raw_mask in raw_masks]
+    def _get_images_masks(self, guided_backprops):
+        masks = []
+        for gbps in guided_backprops:
+            projected_gbps = [np.max(gb, axis=-1).squeeze() for gb in gbps]
+            raw_masks = [pgbp > np.percentile(pgbp, self.config.cut_off_percentile) for pgbp in projected_gbps]
+            # erosion and dilation
+            masks_per_image = [binary_dilation(binary_erosion(raw_mask)).astype(projected_gbps[0].dtype) for raw_mask in
+                               raw_masks]
+            masks.append(masks_per_image)
         return masks
 
-    def _get_topk_neurons(self, image, masks, top_class):
-        reshaped_image = image.reshape(np.roll(self.config.vgg16.input_size, 1))
-        images = np.stack([np.reshape(reshaped_image * mask, self.config.vgg16.input_size) for mask in masks])
-        losses = self.net.get_batch_loss(images, top_class)
-        return list(np.argsort(losses)[:self.config.k])
-        # return kmax_neuron_indices[np.argsort(losses)[:self.config.k]]
+    def _get_top_k_neurons(self, images, masks, top_class):
+        top_k_neurons_relative_indices = []
+        for i, image in enumerate(images):
+            reshaped_image = image.reshape(np.roll(self.config.vgg16.input_size, 1))
+            masked_images = np.stack(
+                [np.reshape(reshaped_image * mask, self.config.vgg16.input_size) for mask in masks[i]])
+            losses = self.net.get_batch_loss(masked_images, top_class[i])
+            top_k_neurons_relative_indices.append(list(np.argsort(losses)[:self.config.k]))
+        return top_k_neurons_relative_indices
+
+    def _get_all_bounding_boxes(self, all_masks, all_mask_indices):
+        bounding_boxes = []
+        for mask, mask_indices in zip(all_masks, all_mask_indices):
+            bounding_boxes.append(self._get_bounding_box(mask, mask_indices))
+        return bounding_boxes
 
     def _get_bounding_box(self, masks, mask_indices):
-        # sorry, worst piece of code mankind has ever done
+        # sorry, super lazy
         final_masks = np.array(masks)[mask_indices]
         the_mask = final_masks[0] * False
         for mask in final_masks:
@@ -108,4 +108,4 @@ def _get_bounding_box(self, masks, mask_indices):
                     x_min = min(x_min, j)
                     y_max = max(y_max, i)
                     x_max = max(x_max, j)
-        return [[x_min, y_min], [x_max, y_max]], the_mask
+        return [[x_min, y_min], [x_max, y_max]]

main.py

@@ -6,8 +6,8 @@
 def main():
     agent = Agent(Config)
     # agent.make_tsne_pic_for_directory()
-    for i in range(100):
-        bounding_box = agent.get_bounding_box()
+    for i in range(3):
+        agent.get_bounding_box()
 
 
 if __name__ == '__main__':

net.py

@@ -29,6 +29,7 @@ def _build_net(self):
             self._vgg16()
             self._fetch_shapes()
             self._build_impact_gradient()
+            self._build_top_kmax_neuron_selection()
             self._build_guided_backprop()
             self._build_softmax_loss()
 
@@ -39,33 +40,52 @@ def _vgg16(self):
         self.fc_features = self.model.get_layer('fc2').output
 
     def _fetch_shapes(self):
-        # lets not run on batch mode for now (gpu shortage and complexity)
         self.features_shape = self.feature_tensor.get_shape().as_list()
         self.classes_shape = self.model.output.get_shape().as_list()
         self.features_shape[0] = self.classes_shape[0] = 1
 
     def _build_impact_gradient(self):
-        self.category_index_ph = tf.placeholder(tf.int32, shape=[])
-        fake_upstream_grad = tf.one_hot(self.category_index_ph, self.classes_shape[-1], axis=-1)
+        self.category_indices = tf.argmax(self.model.output, axis=-1, name='max_scoring_categories')
+        fake_upstream_grad = tf.one_hot(self.category_indices, self.classes_shape[-1], axis=-1)
         self.impact_grad = tf.gradients(self.model.output, self.feature_tensor,
-                                        grad_ys=[fake_upstream_grad],
-                                        name='impact_gradients')
+                                        grad_ys=[fake_upstream_grad], name='impact_gradients')[0]
+
+    def _build_top_kmax_neuron_selection(self):
+        # use DAM heuristic for selection
+        neurons_effect = self.impact_grad * self.feature_tensor
+        neurons_effect_flat_batch = tf.reshape(neurons_effect, (-1, np.prod(self.features_shape)))
+        self.batch_top_kmax_neuron_indices = tf.nn.top_k(neurons_effect_flat_batch, k=self.config.kmax)[1]
 
     def _build_guided_backprop(self):
         self.neuron_index = tf.placeholder(tf.int32, shape=[])
         fake_upstream_grad = tf.one_hot(self.neuron_index, np.prod(self.features_shape), axis=-1)
         fake_upstream_grad = tf.reshape(fake_upstream_grad, shape=self.features_shape)
         self.guided_backprop = tf.gradients(self.feature_tensor, self.model.input,
-                                            grad_ys=[fake_upstream_grad],
-                                            name='guided_backprop')
+                                            grad_ys=[fake_upstream_grad], name='guided_backprop')
 
     def _build_softmax_loss(self):
         self.top_class_index_ph = tf.placeholder(tf.int32, shape=[])
-        self.top_class_batch_one_hot = tf.one_hot(tf.ones([self.config.kmax,], dtype=tf.int32) * self.top_class_index_ph,
+        self.top_class_batch_one_hot = tf.one_hot(tf.ones([self.config.kmax, ], dtype=tf.int32) * self.top_class_index_ph,
                                                   self.classes_shape[-1], axis=-1)
         self.softmax_loss = tf.nn.softmax_cross_entropy_with_logits_v2(labels=self.top_class_batch_one_hot,
                                                                        logits=self.model.output)
 
+    def get_top_kmax_neurons(self, images):
+        top_kmax_neurons_indices, max_scoring_indices = self.sess.run([self.batch_top_kmax_neuron_indices,
+                                                                       self.category_indices],
+                                                                      feed_dict={self.model.input: images,
+                                                                                 self.not_guided_flag: 1.0})
+        return top_kmax_neurons_indices, max_scoring_indices
+
+    def get_guided_backprop(self, image, neuron_index):
+        numerical_guided_backprop = self.sess.run(self.guided_backprop,
+                                                  feed_dict={
+                                                      self.model.input: image,
+                                                      self.neuron_index: neuron_index,
+                                                      self.not_guided_flag: 0.0,
+                                                  })
+        return numerical_guided_backprop[0]
+
     def get_batch_loss(self, images, top_class):
         assert images.shape[0] is self.config.kmax
         batch_loss = self.sess.run(self.softmax_loss,
@@ -75,35 +95,8 @@ def get_batch_loss(self, images, top_class):
                                    })
         return batch_loss
 
-    def get_class_scores(self, image):
-        prediction_scores = np.squeeze(self.model.predict(image))
-        return prediction_scores
-
-    def get_activations(self, image):
-        activations = self.sess.run(self.feature_tensor,
-                                    feed_dict={self.model.input: image})
-        return activations
-
     def get_fc_features(self, images):
         fc_features = self.sess.run(self.fc_features,
                                     feed_dict={self.model.input: images})
         return fc_features
 
-    def get_impact_gradient_per_class(self, activations, class_index):
-        numerical_impact_grad = self.sess.run(self.impact_grad,
-                                              feed_dict={
-                                                  self.feature_tensor: activations,
-                                                  self.category_index_ph: class_index,
-                                                  self.not_guided_flag: 1.0,
-                                              })
-        return numerical_impact_grad[0]
-
-    def get_guided_backprop(self, image, neuron_index):
-        numerical_guided_backprop = self.sess.run(self.guided_backprop,
-                                                  feed_dict={
-                                                      self.model.input: image,
-                                                      self.neuron_index: neuron_index,
-                                                      self.not_guided_flag: 0.0,
-                                                  })
-        return numerical_guided_backprop[0]
-