search

kwotsin / TensorFlow-ENet

TensorFlow implementation of ENet
MIT License
257 stars 123 forks source link

Adding an additional layer to the end of the net #11

Open sjdrc opened 7 years ago

sjdrc commented 7 years ago

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 commented 7 years ago

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