Adding an additional layer to the end of the net

[sjdrc]

Hi there, thanks very much for publishing your code.

I was wondering if it would be possible to add an argmax layer to the net so that the predictions are directly output. I am freezing the graph to a pb file and loading it in C++ and trying to perform that operation there seems to be much more complicated than tf.argmax(probabilities, -1). I added

prediction = tf.argmax(probabilities, -1, name='prediction')

after the logits_to_softmax layer in enet.py and trained for a few more iterations. Before I save the frozen graph I print out all the ops and I can see the 'ENet/prediction:0' layer but when I set that as the output in C++ it cannot target that layer as output. Any ideas on how to add this?

[kwotsin]

Are you trying to perform the XLA AOT compilation? Can you elaborate on what you mean by "it cannot target that layer as output"?