PReLU for multidimensional parameters

[franhervas99]

Hi everyone,
I would like to implement a custom layer for managing PReLU activations with multiple dimensions, e.g., The connection between a Conv1D layer with 6 output channels and 9 samples each channel, and a PReLU. In this case, the PReLU should have 6 channels (6 parameters), and 9 samples on each.
My idea was, using Extension API, to inherit the PReLU keras layer and the PReLU hls4ml layer, then do the parser function and then the HLS code. The thing is that I don't know where to start with the parser function, so any help should be appreciated!
Thank you in advance.

[vloncar]

The simplest approach is to subclass PReLU class without any changes (it can be whatever you like, for example MyPReLU) then register a handler for it. The function for parsing could look like this, but simplified. You can add parsing of shared_axes if you need it. You just ensure that the class_name value corresponds to the name of the class of the hls4ml layer (can also be MyPReLU), so it can be mapped properly once you register that. Since you only mentioned that you don't know where to start with the parser function, I'm assuming you know how to define the hls4ml layer and the matching HLS code 😉. It's not much more than what is already in the docs.

If you do all this we'd be interested in merging it upstream. We don't support this feature currently simply because it wasn't asked for before. It is a bit cumbersome to implement full support since we need a specialization for all reasonable number of dimensions (2D, 3D and maybe 4D) and combinations of shared_axes.

[franhervas99]

Hi @vloncar ,
I have subclass the PReLU class into a new layer (KMyPReLU for the keras layer and HMyPReLU for hlsm4l), adding to the hls4ml layer the shared_axes attribute, like that:

class KMyPReLU(keras.layers.PReLU):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

class HMyPReLU(hls4ml.model.layers.PReLU):
    def initialize(self, shared_axes=None):
        super().initialize()
        self.shared_axes = shared_axes

Then the parser function looks like this:

def parse_myprelu_layer(keras_layer, input_names, input_shapes, data_reader):
    layer = {}

    layer['class_name'] = 'HMyPReLU'
    layer['name'] = keras_layer['config']['name']

    if 'activation' in keras_layer['config']:
        layer['activation'] = keras_layer['config']['activation']

    if layer['class_name'] != 'Activation':
        layer['activation'] = layer['class_name']

    layer['data_format'] = keras_layer['config'].get('data_format', 'channels_last')

    layer['param_data'] = get_weights_data(data_reader, layer['name'], 'alpha')

    if input_names is not None:
        layer['inputs'] = input_names

    layer['shared_axes'] = keras_layer['config']['shared_axes']

    if layer['shared_axes'] == [1]:
        layer['sequence_length'] = input_shapes[0][1]
        layer['num_channels'] = input_shapes[0][2]

    return layer, [shape for shape in input_shapes[0]]

The templates for the instances in HLS, like this:

myprelu_config_template = """struct config{index} : nnet::myprelu_config {{
    static const unsigned n_in = {n_in};
    static const unsigned table_size = {table_size};
    static const unsigned io_type = nnet::{iotype};
    static const unsigned reuse_factor = {reuse};
    static const unsigned sequence_length = {sequence_length};
    static const unsigned num_channels = {num_channels};
    typedef {table_t.name} table_t;
    typedef {param_t.name} param_t;
}};\n"""

myprelu_function_template = (
    'nnet::myprelu<{input_t}, {param_t.name}, {output_t}, {config}>({input}, {param}, {output});'
)

myprelu_include_list = ['nnet_utils/nnet_activation.h', 'nnet_utils/nnet_activation_stream.h', 'nnet_utils/nnet_myprelu.h']

class HMyPReLUConfigTemplate(hls4ml.backends.template.LayerConfigTemplate):
    def __init__(self):
        super().__init__(HMyPReLU)
        self.template = myprelu_config_template

    def format(self, node):
        params = self._default_config_params(node)
        return self.template.format(**params)

class HMyPReLUFunctionTemplate(hls4ml.backends.template.FunctionCallTemplate):
    def __init__(self):
        super().__init__(HMyPReLU, include_header=myprelu_include_list)
        self.template = myprelu_function_template

    def format(self, node):
        params = self._default_function_params(node)
        return self.template.format(**params)

The register itself:

# Register the converter for custom Keras layer
hls4ml.converters.register_keras_layer_handler('KMyPReLU', parse_myprelu_layer)

# Register the hls4ml's IR layer
hls4ml.model.layers.register_layer('HMyPReLU', HMyPReLU)

backend_id = 'Vivado'

# Register the optimization passes (if any)
backend = hls4ml.backends.get_backend(backend_id)

# Register template passes for the given backend
backend.register_template(HMyPReLUConfigTemplate)
backend.register_template(HMyPReLUFunctionTemplate)

# Register HLS implementation
backend.register_source('nnet_myprelu.h')

But, when generating a dummy model like this:

keras_model = keras.Sequential([
    keras.layers.Input(shape=(9, 1)),
    keras.layers.Conv1D(filters=6, kernel_size=3, strides=1, padding='same'),
    KMyPReLU(shared_axes=[1])
])

hmodel = hls4ml.converters.convert_from_keras_model(
    keras_model,
    output_dir=str(f'hls4mlprj_extensions_{backend_id}'),
    backend=backend_id,
    io_type='io_parallel',
    hls_config={'Model': {'Precision': 'ap_int<6>', 'ReuseFactor': 1}},
)

hls4ml.utils.plot_model(hmodel, show_shapes=True, show_precision=True, to_file='hls4ml_model.png')

A plot like this is generated, where there is an extra activation layer at the end that shouldn't be there:

I've been struggling a lot with the problem, and I cannot find the solution. Which could be the problem here?

[vloncar]

Indeed, this is annoying. I think it has to do with the fact that your layer now has activation parameter and this causes the creation of another layer. Until we find a better fix, there are three potential workarounds (I didn't test any of this, it is just from the top of my head):

• Change the activation parameter in the parser to be layer['activation'] = 'linear'. Since you're parsing your own class, it shouldn't matter. The optimizer will remove it. If that doesn't work...
• Add an expected attribute activation to HMyPReLU with value_type=None. Like: _expected_attributes = [Attribute('activation', value_type=None)]. See how internal layers do it. If that doesn't work...
• Reimplement HMyPReLU such that it doesn't subclass Activation but rather Layer and then implement it without the attribute. You can just copy the contents of the Activation and PReLU from the source. If that doesn't work... well, then I guess we'll have to fix this in a proper way.

[franhervas99]

Hi @vloncar ,
Doing the third approach worked for me, but now the PReLU weights are not being propagated to the .h header file for the weights, neither to the .txt file. I think that has something to do with the internal manipulation of the parameters, because maybe implementing this 'shared_axes' functionality changed some other internal thing that I don't understand yet.

[franhervas99]

Is it possible to not use shared_axes, multiplying each position of sequence, channels by a parameter, even if it increases the number of parameters? I say that because I'm trying it, but I get bad results comparing the keras output with the hls output, and I don't know if the reason is the shape or if it is not possible to do that due to its multidimensionality.

[vloncar]

Can you provide the nnet_myprelu.h? I tried making the fix the usual way and the weights are written, so I only lack the support for using in in the proper way.

[franhervas99]

Sure, here is de code of nnet_myprelu.h:

#ifndef NNET_MYPRELU_H_
#define NNET_MYPRELU_H_

#include "nnet_common.h"

namespace nnet {

struct myprelu_config {
    static const unsigned n_in = 54;
    static const unsigned io_type = nnet::io_parallel;
    static const unsigned reuse_factor = 1;
    static const unsigned sequence_length = 9;
    static const unsigned num_channels = 6;
    typedef model_default_t param_t;
};

template <class data_T, class param_T, class res_T, typename CONFIG_T>
void myprelu(data_T data[CONFIG_T::n_in], param_T alpha[CONFIG_T::num_channels], res_T res[CONFIG_T::n_in]) {
    #pragma HLS PIPELINE

    for (int j = 0; j < CONFIG_T::sequence_length; j++) {
        for (int i = 0; i < CONFIG_T::num_channels; i++) {
            int flat_index = j * CONFIG_T::num_channels + i;
            data_T datareg = data[flat_index];
            // printf("%f\n", (double)datareg);
            if (datareg > 0) {
                res[flat_index] = datareg;
            } else {
                res[flat_index] = alpha[i] * datareg;
            }
        }
    }
}

}

#endif // NNET_MYPRELU_H_