Use tf.softmax_cross_entropy_with_logits to calculate loss (tensorflow#181)

Evan Lezar · martinwicke · commit d8169710322c · 2016-06-22T19:37:18.000-07:00
* Use the tensorflow cross entropy function to prevent nan losses.

* Correct double softmax and use mean for loss.
diff --git a/transformer/README.md b/transformer/README.md
@@ -17,7 +17,7 @@ A Spatial Transformer Network implemented in Tensorflow 0.7 and based on [2].
 </div>
 
 ```python
-transformer(U, theta, downsample_factor=1)
+transformer(U, theta, out_size)
 ```
     
 #### Parameters
diff --git a/transformer/cluttered_mnist.py b/transformer/cluttered_mnist.py
@@ -119,16 +119,17 @@
 # %% And finally our softmax layer:
 W_fc2 = weight_variable([n_fc, 10])
 b_fc2 = bias_variable([10])
-y_pred = tf.nn.softmax(tf.matmul(h_fc1_drop, W_fc2) + b_fc2)
+y_logits = tf.matmul(h_fc1_drop, W_fc2) + b_fc2
 
 # %% Define loss/eval/training functions
-cross_entropy = -tf.reduce_sum(y * tf.log(y_pred))
+cross_entropy = tf.reduce_mean(
+    tf.nn.softmax_cross_entropy_with_logits(y_logits, y))
 opt = tf.train.AdamOptimizer()
 optimizer = opt.minimize(cross_entropy)
 grads = opt.compute_gradients(cross_entropy, [b_fc_loc2])
 
 # %% Monitor accuracy
-correct_prediction = tf.equal(tf.argmax(y_pred, 1), tf.argmax(y, 1))
+correct_prediction = tf.equal(tf.argmax(y_logits, 1), tf.argmax(y, 1))
 accuracy = tf.reduce_mean(tf.cast(correct_prediction, 'float'))
 
 # %% We now create a new session to actually perform the initialization the
