Fix content_loss shape inference for Theano

titu1994 · Aug 9, 2017 · 037a8a7 · 037a8a7
1 parent 669aa17
commit 037a8a7
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 4 deletions.
diff --git a/INetwork.py b/INetwork.py
@@ -437,7 +437,10 @@ def style_loss(style, combination, mask_path=None, nb_channels=None):
 def content_loss(base, combination):
     channel_dim = 0 if K.image_dim_ordering() == "th" else -1
 
-    channels = K.int_shape(base)[channel_dim]
+    try:
+        channels = K.int_shape(base)[channel_dim]
+    except TypeError:
+        channels = K.shape(base)[channel_dim]
     size = img_width * img_height
 
     if args.content_loss_type == 1:

diff --git a/Network.py b/Network.py
@@ -433,7 +433,10 @@ def style_loss(style, combination, mask_path=None, nb_channels=None):
 def content_loss(base, combination):
     channel_dim = 0 if K.image_dim_ordering() == "th" else -1
 
-    channels = K.int_shape(base)[channel_dim]
+    try:
+        channels = K.int_shape(base)[channel_dim]
+    except TypeError:
+        channels = K.shape(base)[channel_dim]
     size = img_width * img_height
 
     if args.content_loss_type == 1:

diff --git a/script_helper/Script/INetwork.py b/script_helper/Script/INetwork.py
@@ -437,7 +437,10 @@ def style_loss(style, combination, mask_path=None, nb_channels=None):
 def content_loss(base, combination):
     channel_dim = 0 if K.image_dim_ordering() == "th" else -1
 
-    channels = K.int_shape(base)[channel_dim]
+    try:
+        channels = K.int_shape(base)[channel_dim]
+    except TypeError:
+        channels = K.shape(base)[channel_dim]
     size = img_width * img_height
 
     if args.content_loss_type == 1:

diff --git a/script_helper/Script/Network.py b/script_helper/Script/Network.py
@@ -433,7 +433,10 @@ def style_loss(style, combination, mask_path=None, nb_channels=None):
 def content_loss(base, combination):
     channel_dim = 0 if K.image_dim_ordering() == "th" else -1
 
-    channels = K.int_shape(base)[channel_dim]
+    try:
+        channels = K.int_shape(base)[channel_dim]
+    except TypeError:
+        channels = K.shape(base)[channel_dim]
     size = img_width * img_height
 
     if args.content_loss_type == 1: