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?


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 Jan 13 at 14:59
  • $\begingroup$ No problem! Please consider marking this answer as accepted if it solved your problem! $\endgroup$ – Carl Lange Jan 13 at 16:49
  • $\begingroup$ I'll leave it a day or so before accepting it $\endgroup$ – EstabanW Jan 13 at 18:58
  • $\begingroup$ Isn't there a way to compute the len from the input layer without having the user to explicitly supply it? $\endgroup$ – M.R. Jan 14 at 15:30
  • $\begingroup$ I don't know for sure! It certainly seems like it should be possible, since so many of the built-in layers have the ability to do this, but I didn't find anything that would calculate the dimensions for the ReplicateLayer. $\endgroup$ – Carl Lange Jan 14 at 15:40

Your Answer

By clicking "Post Your Answer", you acknowledge that you have read our updated terms of service, privacy policy and cookie policy, and that your continued use of the website is subject to these policies.

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