I am working on implementing the loss function of the Yolo version 2 object detection network and building from the construction notebook in the Wolfram Neural Network repository Yolo Version 2 so I can train it for a specialised detection task. The loss function is quite complicated, I can see how to implement almost all of it but I'm stuck with an indicator function implementation.

During training indicator functions are used to set the maximum value in a tensor to 1 and other values to zero (in the Yolo network it is an higher rank tensor) but the problem comes down to this, how to implement the following (a sort of indexmax or argmax function)

{1,2,1,4,2} => {0,0,0,1,0}

There are many ways of doing this with built-in functions but the problem I'm stuck with is how to do it with the functionality available in the neural network framework. In the snippet below, the Position function is not available (that I can find) in the neural net framework.

vec = {1,2,1,4,2};
pos = Position[vec, Max[vec]]//Flatten (*this can't be directly implemented in the NN framework*)
UnitVectorLayer[5,"Input"-> "Integer"][pos]

This produces {0.,0.,0.,1.,0.} (*the output I'm after*)

In Keras there is an 'argmax' function that does something like this. Other functionality like ReplicateLayer and related functions can then be used to construct the object needed, this bit isn't where the problem is.

Any ideas or suggestions on how to do this in the MMA neural network framework?


1 Answer 1


You can construct an IndexMaxLayer using the existing layers, like so:

indexMaxLayer[] := NetGraph[{
   "max" -> AggregationLayer[Max, 1],
   "repl" -> ReplicateLayer[Automatic],
   "switch" -> ThreadingLayer[If[#1 != #2, 0., 1.] &]
   }, {
   "max" -> "repl",
   {NetPort@"Input", "repl"} -> "switch"


> {0., 0., 0., 1., 0.}

This is getting the maximum value in the input vector, replicating that value to be the same dimensions as the input vector, and then switching the input layer to 0 or 1 if it does or doesn't match the max value.

This is extendable to multiple dimensions if you change the second argument to AggregationLayer, although there are multiple different behaviours you might want in that case (max per level or max over the entire matrix, for instance).

  • $\begingroup$ That's neat, thanks very much! I'd looked thoroughly through the documentation at the functions supported by ThreadingLayer, AggregationLayer etc and hadn't thought of experimenting with others (e.g. If) to see if they work. $\endgroup$
    – EstabanW
    Commented Jan 13, 2019 at 14:59
  • $\begingroup$ No problem! Please consider marking this answer as accepted if it solved your problem! $\endgroup$
    – Carl Lange
    Commented Jan 13, 2019 at 16:49
  • 1
    $\begingroup$ Aha, Automatic. Updating my answer now $\endgroup$
    – Carl Lange
    Commented Jan 14, 2019 at 15:42
  • 1
    $\begingroup$ @LucaAmerio It works perfectly for me on my local 11.3 install and in a cloud notebook. Perhaps you could open a question with the behaviour you're seeing. $\endgroup$
    – Carl Lange
    Commented Mar 22, 2019 at 22:13
  • 1
    $\begingroup$ @CarlLange I discovered that in order for this to work, you need version 11.3.5 of the NN Paclet. You can check your version doing Lookup["Version"]@PacletInformation["NeuralNetworks"]. Version 11.3 comes with version 11.3.4, so this code doesn't run. To update it, you need to run PacletSiteUpdate /@ PacletSites[]; PacletUpdate /@ {"NeuralNetworks", "MXNetLink"} $\endgroup$
    – Luca
    Commented Apr 5, 2019 at 10:04

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.