Implementation of max pooling operation in NengoDL

[MojoForPresident]

Hello everybody!

I am currently attempting to implement a max pooling operation in the Nengo DL converter. The converter code for this is pretty straight forward and is similar to the one in the converter for the AveragePool layer:

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

    def convert(self, node_id):
        ...
        pool_transform = MaxPool(
            self.input_shape(node_id),
            kernel_size=pool_size,
            strides=strides,
            padding=padding
        )

        self.add_connection(node_id, output, transform=pool_transform)
        return output

Now the big challenge there is the “MaxPool” tranformation in the end of the code sample. This one is not provided by Nengo DL and therefore must be implemented.

So my first question is: is there any specific reason why max pooling is not (yet) implemented? And secondly, what steps have to be taken when implementing a new transformation like that?

I already gave implementing the transformation a go by adapting the existing nengo.Convolution transformation.
For this, I first created a new class plus the (apparently) needed build function:

class MaxPool(Transform):
...

@Builder.register(MaxPool)
def build_maxpool(
    model, transform, sig_in, decoders=None, encoders=None, rng=np.random
):
...
    model.add_op(
        MaxPoolInc(
            sig_in, pooled, transform, tag="%s.apply_weights" % transform
        )
    )
...

class MaxPoolInc(Operator):
    def make_step(self, signals, dt, rng):
        ...
        def step_max_pool():
            Y += max_pool2d(
                input=X,
                ksize=self.maxpool.kernel_size,
                strides=stride,
                padding=pad
            )

        return step_max_pool
        ...

As you can see, in the build function, there is a MaxPoolInc operator, that uses the max_pool2d operation from the tensorflow package. When I try to make this work, I get an error that the MaxPoolInc operator also needs to be registered. After further inspection, I found the corresponding code for the nengo.Convolution transformation in the nengo_dl.transform_builders.py file, line 33:

@Builder.register(ConvInc)
@Builder.register(ConvSet)
class ConvIncBuilder(OpBuilder):
    """
    Build a group of `nengo.builder.transforms.ConvInc` operators.
    """
    ...

This class however seems to be rather complicated and long (more than 150 lines).

So my question now is, what does this ConvIncBuilder do and what do I need to do to make my own transformation work. If needed, I can also try to provide an MVE.

Thanks in advance!
Cheers.

[xchoo]

Hi @MojoForPresident

@MojoForPresident

So my first question is: is there any specific reason why max pooling is not (yet) implemented? And secondly, what steps have to be taken when implementing a new transformation like that?

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.

@MojoForPresident

So my question now is, what does this ConvIncBuilder do

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.

@MojoForPresident

what do I need to do to make my own transformation work.

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.

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.