Implementation of max pooling operation in NengoDL

Deep Learning
2

Hi @MojoForPresident

Yes. These forum threads (here and here) discuss some of the reasons why max pooling is not implemented in NengoDL’s converter. The gist of it is:

  • It’s unclear what the max pooling operation is meant to do in the context of spiking networks.
  • It’s a non-linear operation which cannot be simply implemented within the connection weights of the network.

To understand this, we need to understand how the NengoDL converter operates within the Nengo ecosystem. The goal of the NengoDL converter is to take a Keras model and convert it into an equivalent Nengo model. To do so, for each Keras layer, the NengoDL converter calls a convert class that has been registered for that specific Keras layer object. The convert class in turn makes Nengo objects that replicate the functionality of the specific Keras layer. For the Conv2D Keras layer, the Nengo object created is the nengo.Convolution transform.

Now, if one were to take the NengoDL converter.net network and run it using the Nengo simulator (i.e., nengo.Simulator), that would be all that’s required (assuming all of the Keras layers were converted successfully without needing and TensorNode objects), since everything in the network is standard Nengo objects. When the Nengo simulator is created, it rebuilds the Nengo network as a bunch of Nengo operators, which is then run on your CPU.

The tricky part comes when running the converter.net network using the NengoDL simulator (i.e., nengo_dl.Simulator). A lot of NengoDL simulator functionality (e.g., sim.fit, sim.evaluate, etc.,) use the TensorFlow backend to run the network. Thus, in order to allow Nengo-native networks to be run in NengoDL, a Nengo network is first built into the collection of Nengo operators (just like for a vanilla Nengo simulator), and then, NengoDL translates the various Nengo operators into TensorFlow code so that the Nengo operator equivalent can be run in TensorFlow. This is what the ConvIncBuilder class is doing – it’s the code that is run to convert the ConvInc (and ConvSet) operator, which are Nengo operators, into the equivalent TensorFlow code so that it can be run in TensorFlow.

The issue with the Max Pool operation is that it is a non-linear function applied to the outputs of a layer. To make your transformation work, my suggestion to you is as such:

  1. In vanilla Nengo, create an ensemble (or network) that computes the max pooling operation to your satisfaction. You can refer to my forum post for ideas on how to do this.
  2. Once you have Nengo code that computes the max pooling operation to you liking, modify the ConvertMaxPool2D.convert function, and insert your max-pool Nengo network there, in place of the nengo.Convolution call. If you look at the ConvertAvgPool class(instead of using ConvertConv) in converter.py, I think you’ll have a better starting point.

In ConvertAvgPool, you have some boiler plate code that handles the various AvePool layer options .Note that ConvertAvgPool is generic code that handles basically all of the average pooling layers (e.g., AvePool1D, GlobalAvgPool1D, etc.), hence the large amount of boiler plate code. The two important lines of code for ConvertAvgPool is the nengo.Convolution transform, and the self.add_connection(...) function. For the max pool layer, you’ll want to replace these two lines with the Nengo network that you created in Step 1. Within the convert function, you can define Nengo objects (see the ConvertConv.convert function for examples), and these Nengo objects will automatically be added to the Nengo object that makes up the converted network.

Some example code to start you off:

@nengo_dl.Converter.register(MaxPool2D)
class ConvertMaxPool2D(nengo_dl.converter.LayerConverter):
    unsupported_args = [("padding", "valid")]

    def convert(self, node_id):
        ...  # Boiler plate code
        
        def max_func(x):
            # The function that computes the max pooling
            return [...]

        max_pool_ens = nengo.networks.EnsembleArray(100, n_ensembles=..., ens_dimensions=...)
        max_pool_ens.add_output("max", max_func)
        self.add_connection(node_id, max_pool_ens)  # Connect input to max_pool_ens
        self.add_connection(max_pool_ens.max, output)  # Connect the output of max_pool_ens that computes the max function, to the output of the layer

Since the code above replicates the max pooling operation using standard Nengo objects, and no custom Nengo transforms, or no custom Nengo operators, you don’t have to take the additional step of implementing builders in NengoDL to convert those custom transforms and operators back into TensorFlow code. :smiley:

And yes, I should note that this approach does modify the Keras model definition a bit, by adding essentially one additional layer to compute the max pooling function.

1 Like