Creating an indicator function in a neural network

Question

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?

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"
   }]

indexMaxLayer[][{1,2,1,4,2}]

> {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).