5 ms·
>>> Softmax can be implemented as a composition of primitive TensorFlow ops (exponent, reduction, elementwise division, etc.): softmax = exp(logits) / reduce_s
by Marat_Dukhan 10y ago
>>> Softmax can be implemented as a composition of primitive TensorFlow ops (exponent, reduction, elementwise division, etc.): softmax = exp(logits) / reduce_sum(exp(logits), dim)
No, it can not be implemented this way, it is numerically unstable, and will produce NaNs if any input is greater than ~88.7. Luckily, it is also not how its implemented in Tensorflow: https://github.com/tensorflow/tensorflow/blob/2c8d0dca978a246f54c506aae4587dbce5d3bcf0/tensorflow/core/kernels/softmax_op_functor.h#L43 https://github.com/tensorflow/tensorflow/blob/2c8d0dca978a24...
For a clean (and more efficient) C version of this algorithm, take a look at NNPACK reference implementation: https://github.com/Maratyszcza/NNPACK/blob/master/src/ref/softmax-output.c#L30 https://github.com/Maratyszcza/NNPACK/blob/master/src/ref/so...
- learyg 10y agoThe XLA and vanilla TF variants appear to be the same here: https://github.com/tensorflow/tensorflow/blob/2c8d0dca978a246f54c506aae4587dbce5d3bcf0/tensorflow/compiler/tf2xla/kernels/softmax_op.cc#L50 https://github.com/tensorflow/tensorflow/blob/2c8d0dca978a24... https://github.com/tensorflow/tensorflow/blob/2c8d0dca978a246f54c506aae4587dbce5d3bcf0/tensorflow/core/kernels/softmax_op_functor.h#L67 https://github.com/tensorflow/tensorflow/blob/2c8d0dca978a24... Chalk it up to poetic license? ;-)
- acmj 10y agoIn the nnpack implementation, the same exponential (i.e. expf) is computed twice for each element, which is a waste of time. A faster implementation should save each expf result to output[sample][channel] first, compute the sum and then rescale output[sample][channel] by the sum.
- Marat_Dukhan 10y agoIf you have an efficient vectorized implementation of expf (NNPACK does), softmax is a memory/cache bandwidth-bound kernel, and storing data is less efficient than recomputing.
- acmj 10y agoWhere is this vectorized expf implemented? I am only seeing softmax is calling the standard expf. Let's suppose expf from libm is vectorized. Is there any benchmark showing nnpack's implementation is really faster? I doubt, actually. Exp is quite expensive even if vectorized.
- Marat_Dukhan 10y agoThis is the reference implementation of softmax, i.e. implementation used as a reference in unit tests. It is designed to be simple, readable, and correct, which is why I linked it here. Optimized implementation is in assembly (PeachPy). See https://github.com/Maratyszcza/NNPACK/blob/master/src/x86_64-fma/exp.py https://github.com/Maratyszcza/NNPACK/blob/master/src/x86_64... for the vectorized expf
- Marat_Dukhan 10y agoOops, here is the actually used vectorized expf (similar to the one I linked, but with optional unrolling): https://github.com/Maratyszcza/NNPACK/blob/master/src/x86_64-fma/vecmath/exp.py https://github.com/Maratyszcza/NNPACK/blob/master/src/x86_64...
- acmj 10y agoThanks for the pointer. The full softmax implementation is here [1]. I have not read the code, but I can trust the developer to have a very fast implementation. Nonetheless, I don't think the reference implementation in your original link is optimal. Exp is expensive and should not be called twice (EDIT: unless you can show a benchmark to prove me wrong). [1]: https://github.com/Maratyszcza/NNPACK/blob/master/src/x86_64-fma/softmax.py https://github.com/Maratyszcza/NNPACK/blob/master/src/x86_64...
- barrkel 10y agoFWIW, you're replying to the developer of the file you linked to.