Squeeze out singleton dimensions for softmax

Saumitro Dasgupta · Saumitro Dasgupta · commit ad0ccd4acf39 · 2016-05-16T00:51:16.000-07:00
diff --git a/kaffe/tensorflow/network.py b/kaffe/tensorflow/network.py
@@ -193,6 +193,15 @@ def fc(self, input, num_out, name, relu=True):
 
     @layer
     def softmax(self, input, name):
+        input_shape = map(lambda v: v.value, input.get_shape())
+        if len(input_shape)>2:
+            # For certain models (like NiN), the singleton spatial dimensions
+            # need to be explicitly squeezed, since they're not broadcast-able
+            # in TensorFlow's NHWC ordering (unlike Caffe's NCHW).
+            if input_shape[1]==1 and input_shape[2]==1:
+                input = tf.squeeze(input, squeeze_dims=[1, 2])
+            else:
+                raise ValueError('Rank 2 tensor input expected for softmax!')
         return tf.nn.softmax(input, name)
 
     @layer
