Minor updates in inverted rep. code

mariopi27 · Jan 17, 2018 · 4e29879 · 4e29879
1 parent b50f8e1
commit 4e29879
Showing 1 changed file with 9 additions and 10 deletions.
diff --git a/src/inverted_representation.py b/src/inverted_representation.py
@@ -15,6 +15,8 @@ class InvertedRepresentation():
     def __init__(self, model):
         self.model = model
         self.model.eval()
+        if not os.path.exists('../generated'):
+            os.makedirs('../generated')
 
     def alpha_norm(self, input_matrix, alpha):
         """
@@ -91,25 +93,24 @@ def generate_inverted_image_specific_layer(self, input_image, img_size, target_l
             # Calculate alpha regularization
             reg_alpha = alpha_reg_lambda * self.alpha_norm(opt_img, alpha_reg_alpha)
             # Calculate total variation regularization
-            reg_total_variation = tv_reg_lambda * self.total_variation_norm(opt_img, tv_reg_beta)
+            reg_total_variation = tv_reg_lambda * self.total_variation_norm(opt_img,
+                                                                            tv_reg_beta)
             # Sum all to optimize
             loss = euc_loss + reg_alpha + reg_total_variation
-            print('Iteration:', str(i), 'Loss:', loss.data.numpy()[0])
             # Step
             loss.backward()
             optimizer.step()
             # Generate image every 5 iterations
             if i % 5 == 0:
+                print('Iteration:', str(i), 'Loss:', loss.data.numpy()[0])
                 x = recreate_image(opt_img)
-                cv2.imwrite('../generated/Layer_' + str(target_layer) +
-                            '_Inverted_Image_Iteration:' + str(i) + '.png', x)
+                cv2.imwrite('../generated/Inv_Image_Layer_' + str(target_layer) +
+                            '_Iteration_' + str(i) + '.png', x)
             # Reduce learning rate every 40 iterations
             if i % 40 == 0:
                 for param_group in optimizer.param_groups:
                     param_group['lr'] *= 1/10
 
-        return 1
-
 
 if __name__ == '__main__':
     # Get params
@@ -118,10 +119,8 @@ def generate_inverted_image_specific_layer(self, input_image, img_size, target_l
         get_params(target_example)
 
     inverted_representation = InvertedRepresentation(pretrained_model)
-    image_size = 224 # width - height
-    target_layer = 1
+    image_size = 224  # width & height
+    target_layer = 12
     inverted_representation.generate_inverted_image_specific_layer(prep_img,
                                                                    image_size,
                                                                    target_layer)
-
-