Apply QKeras weight quantizers to SimpleRNN, LSTM and GRU

hls4ml/backends/altera/altera_backend.py

@@ -378,10 +378,11 @@ def init_lstm(self, layer):
             self.set_closest_reuse_factor(layer, n_in_recr, n_out_recr, attribute='recurrent_reuse_factor')
             layer.set_attr('strategy', 'resource')
 
-        # Split weights for easier storage in on-chip memory and implementation in HLS
-        weights_data = layer.weights['weight'].data
-        rec_weights_data = layer.weights['recurrent_weight'].data
-        bias_data = layer.weights['bias'].data
+        # Split weights for easier storage in on-chip memory and implementation in HLS. The pieces are quantized
+        # from the unquantized weights, like the whole arrays were.
+        weights_data = layer.weights['weight'].data_unquantized
+        rec_weights_data = layer.weights['recurrent_weight'].data_unquantized
+        bias_data = layer.weights['bias'].data_unquantized
 
         weight_types = ['i', 'f', 'c', 'o']
         for i in range(0, 4):
@@ -400,14 +401,14 @@ def init_lstm(self, layer):
                 data=rec_weights_data[
                     0 : layer.get_attr('n_out'), i * layer.get_attr('n_out') : (i + 1) * layer.get_attr('n_out')
                 ],
-                quantizer=layer.get_attr('weight_quantizer'),
+                quantizer=layer.get_attr('recurrent_weight_quantizer'),
                 compression=None,
             )
             layer.add_weights_variable(
                 name=f'bias_{weight_types[i]}',
                 var_name=f'bias_{weight_types[i]}_{{index}}',
                 data=bias_data[i * layer.get_attr('n_out') : (i + 1) * (layer.get_attr('n_out'))],
-                quantizer=layer.get_attr('weight_quantizer'),
+                quantizer=layer.get_attr('bias_quantizer'),
                 compression=None,
             )
 

hls4ml/converters/keras_v3_to_hls.py

@@ -260,6 +260,8 @@ def v2_call(
         ret = (ret,)
 
         activation = getattr(layer, 'activation', None)
+        if key in ('SimpleRNN', 'LSTM', 'GRU', 'QSimpleRNN', 'QLSTM', 'QGRU', 'Bidirectional'):
+            activation = None  # a recurrent layer applies its activation itself, as in the Keras v2 converter
         if activation not in (keras.activations.linear, None):
             assert isinstance(activation, FunctionType), f'Activation function for layer {layer.name} is not a function'
             intermediate_tensor_name = f'{output_names[0]}_activation'

hls4ml/model/layers.py

@@ -1478,13 +1478,15 @@ def initialize(self):
             )
 
         # weights
-        self.add_weights()
+        self.add_weights(quantizer=self.get_attr('weight_quantizer'))
 
         # recurrent weights
-        self.add_weights_variable(name='recurrent_weight', var_name='wr{index}')
+        self.add_weights_variable(
+            name='recurrent_weight', var_name='wr{index}', quantizer=self.get_attr('recurrent_weight_quantizer')
+        )
 
         # biases
-        self.add_weights_variable(name='bias', var_name='b{index}')
+        self.add_weights_variable(name='bias', var_name='b{index}', quantizer=self.get_attr('bias_quantizer'))
         if 'pytorch' in self.attributes.keys():
             self.add_weights_variable(name='recurrent_bias', var_name='br{index}')
 
@@ -1527,14 +1529,19 @@ def initialize(self):
             )
 
         # weights
-        self.add_weights()
+        self.add_weights(quantizer=self.get_attr('weight_quantizer'))
 
         # recurrent weights
         recurrent_weight = self.get_attr('recurrent_weight_data')
-        self.add_weights_variable(name='recurrent_weight', var_name='wr{index}', data=recurrent_weight)
+        self.add_weights_variable(
+            name='recurrent_weight',
+            var_name='wr{index}',
+            data=recurrent_weight,
+            quantizer=self.get_attr('recurrent_weight_quantizer'),
+        )
 
         # biases
-        self.add_weights_variable(name='bias', var_name='b{index}')
+        self.add_weights_variable(name='bias', var_name='b{index}', quantizer=self.get_attr('bias_quantizer'))
 
         if 'pytorch' in self.attributes.keys():
             self.add_weights_variable(name='recurrent_bias', var_name='br{index}')
@@ -1582,14 +1589,16 @@ def initialize(self):
             )
 
         # weights
-        self.add_weights()
+        self.add_weights(quantizer=self.get_attr('weight_quantizer'))
 
         # recurrent weights
-        self.add_weights_variable(name='recurrent_weight', var_name='wr{index}')
+        self.add_weights_variable(
+            name='recurrent_weight', var_name='wr{index}', quantizer=self.get_attr('recurrent_weight_quantizer')
+        )
 
         # biases
-        self.add_weights_variable(name='bias', var_name='b{index}')
-        self.add_weights_variable(name='recurrent_bias', var_name='br{index}')
+        self.add_weights_variable(name='bias', var_name='b{index}', quantizer=self.get_attr('bias_quantizer'))
+        self.add_weights_variable(name='recurrent_bias', var_name='br{index}', quantizer=self.get_attr('bias_quantizer'))
 
 
 class TimeDistributed(Layer):

hls4ml/model/types.py

@@ -832,9 +832,14 @@ def _format(self):
         # Use an XnorBinary-like representation for the sign
         sign = np.where(y < 0, np.zeros_like(y), np.ones_like(y))
         # Take the logarithm, since this is what we will write to the header
-        # for the optimized product using shifts
-        y = (np.log2(np.abs(y)) / np.log2(2.0)).astype('int')
-        return np.stack((sign, y), axis=-1)
+        # for the optimized product using shifts. A power of two is never zero, but the data can be padded with
+        # zeros (the Resource strategy of recurrent layers); those get the smallest exponent of the type, the value
+        # closest to zero it can represent.
+        precision = self.type.precision
+        exponent = np.full(y.shape, -(2 ** (precision.width - 1)) if precision.signed else 0, dtype='int')
+        nonzero = y != 0
+        exponent[nonzero] = np.log2(np.abs(y[nonzero])).astype('int')
+        return np.stack((sign, exponent), axis=-1)
 
     def __iter__(self):
         data = self._format()

test/pytest/test_qkeras.py

@@ -13,6 +13,7 @@
 from qkeras.qnormalization import QBatchNormalization
 from qkeras.quantizers import (
     binary,
+    get_quantizer,
     quantized_bits,
     quantized_po2,
     quantized_relu,
@@ -26,6 +27,7 @@
 from sklearn.preprocessing import LabelEncoder, StandardScaler
 
 import hls4ml
+from hls4ml.model.types import ExponentWeightVariable
 
 co = {}
 _add_supported_quantized_objects(co)
@@ -414,11 +416,11 @@ def test_qactivation_kwarg(test_case_id, randX_100_10, activation_quantizer, wei
     'QDepthwiseConv2D': ((8, 8, 3), lambda: QDepthwiseConv2D((3, 3)), ['Vivado', 'Vitis']),
     'QSeparableConv1D': ((8, 3), lambda: QSeparableConv1D(4, 3), ['Vivado', 'Vitis']),
     'QSeparableConv2D': ((8, 8, 3), lambda: QSeparableConv2D(4, (3, 3)), ['Vivado', 'Vitis']),
-    'QSimpleRNN': ((5, 3), lambda: QSimpleRNN(4, activation='relu'), ['Quartus', 'Altera']),
+    'QSimpleRNN': ((5, 3), lambda: QSimpleRNN(4, activation='relu'), ALL_BACKENDS),
     'QLSTM': (
         (5, 3),
         lambda: QLSTM(4, activation='tanh', recurrent_activation='sigmoid'),
-        ['Vivado', 'Quartus', 'Altera'],
+        ALL_BACKENDS,
     ),
     'QGRU': (
         (5, 3),
@@ -431,7 +433,7 @@ def test_qactivation_kwarg(test_case_id, randX_100_10, activation_quantizer, wei
             reset_after=True,
             recurrent_quantizer='quantized_bits(24, 7, alpha=1)',
         ),
-        ['Vivado', 'Quartus', 'Altera'],
+        ALL_BACKENDS,
     ),
     'QBatchNormalization': (
         (10,),
@@ -645,7 +647,20 @@ def test_quantised_po2_bit_width(test_case_id, backend, io_type, strategy):
     np.testing.assert_allclose(y_hls.flatten(), y_keras.flatten(), rtol=2e-2)
 
 
-@pytest.mark.parametrize('backend', ['Quartus', 'Altera'])
+def assert_rnn_weights_quantized(hls_layer, quantizers, backend):
+    """Check that the weights of a recurrent layer are the QKeras-quantized values of the model's weights.
+
+    ``quantizers`` maps weight names to quantizers. Quartus and Altera store the kernels transposed.
+    """
+    for name, quantizer in quantizers.items():
+        w = hls_layer.weights[name]
+        expected = np.asarray(quantizer(w.data_unquantized))
+        if backend in ('Quartus', 'Altera') and expected.ndim == 2:
+            expected = expected.T
+        np.testing.assert_array_equal(w.data, expected)
+
+
+@pytest.mark.parametrize('backend', ['Vivado', 'Vitis', 'Quartus', 'Altera'])
 def test_qsimplernn(test_case_id, backend):
     """
     Test proper handling of QSimpleRNN.
@@ -678,9 +693,13 @@ def test_qsimplernn(test_case_id, backend):
     y_hls4ml = hls_model.predict(X)
 
     np.testing.assert_allclose(y_qkeras, y_hls4ml.reshape(y_qkeras.shape), atol=0.1)
+    q = quantized_bits(16, 0, alpha=1)
+    assert_rnn_weights_quantized(
+        hls_model.graph[model.layers[0].name], {'weight': q, 'recurrent_weight': q, 'bias': q}, backend
+    )
 
 
-@pytest.mark.parametrize('backend', ['Vivado', 'Quartus', 'Altera'])
+@pytest.mark.parametrize('backend', ['Vivado', 'Vitis', 'Quartus', 'Altera'])
 def test_qlstm(test_case_id, backend):
     """
     Test proper handling of QLSTM.
@@ -714,9 +733,13 @@ def test_qlstm(test_case_id, backend):
     y_hls4ml = hls_model.predict(X)
 
     np.testing.assert_allclose(y_qkeras, y_hls4ml.reshape(y_qkeras.shape), atol=0.1)
+    q = quantized_bits(8, 0, alpha=1)
+    assert_rnn_weights_quantized(
+        hls_model.graph[model.layers[0].name], {'weight': q, 'recurrent_weight': q, 'bias': q}, backend
+    )
 
 
-@pytest.mark.parametrize('backend', ['Vivado', 'Quartus', 'Altera'])
+@pytest.mark.parametrize('backend', ['Vivado', 'Vitis', 'Quartus', 'Altera'])
 def test_qgru(test_case_id, backend):
     """
     Test proper handling of QGRU.
@@ -751,6 +774,107 @@ def test_qgru(test_case_id, backend):
     y_hls4ml = hls_model.predict(X)
 
     np.testing.assert_allclose(y_qkeras, y_hls4ml.reshape(y_qkeras.shape), atol=0.1)
+    q = quantized_bits(8, 0, alpha=1)
+    assert_rnn_weights_quantized(
+        hls_model.graph[model.layers[0].name],
+        {'weight': q, 'recurrent_weight': q, 'bias': q, 'recurrent_bias': q},
+        backend,
+    )
+
+
+RNN_QUANTIZERS = {
+    'quantized_bits': {
+        'weight': 'quantized_bits(8, 0, alpha=1)',
+        'recurrent_weight': 'quantized_bits(6, 0, alpha=1)',
+        'bias': 'quantized_bits(4, 0, alpha=1)',
+    },
+    'quantized_po2': {
+        'weight': 'quantized_po2(4)',
+        'recurrent_weight': 'quantized_po2(5)',
+        'bias': 'quantized_bits(8, 0, alpha=1)',
+    },
+}
+
+
+def make_qrnn_model(layer, quantizers):
+    """The model of #1579: one QKeras recurrent layer of 8 units on a (6, 4) input, with the given kernel, recurrent
+    and bias quantizers. The weights are drawn from (-0.5, 0.5), so the bias is not zero and the model is the same
+    with every Keras version."""
+    kwargs = dict(
+        kernel_quantizer=quantizers['weight'],
+        recurrent_quantizer=quantizers['recurrent_weight'],
+        bias_quantizer=quantizers['bias'],
+    )
+    layers = {
+        'QSimpleRNN': lambda: QSimpleRNN(8, activation='relu', **kwargs),
+        'QLSTM': lambda: QLSTM(8, activation='tanh', recurrent_activation='sigmoid', **kwargs),
+        'QGRU': lambda: QGRU(8, activation='tanh', recurrent_activation='sigmoid', reset_after=True, **kwargs),
+    }
+    x_in = Input((6, 4))
+    model = Model(x_in, layers[layer]()(x_in))
+    rng = np.random.default_rng(1)
+    model.layers[1].set_weights([rng.uniform(-0.5, 0.5, w.shape).astype('float32') for w in model.layers[1].get_weights()])
+    return model
+
+
+@pytest.mark.parametrize('layer', ['QSimpleRNN', 'QLSTM', 'QGRU'])
+@pytest.mark.parametrize('quantizers', list(RNN_QUANTIZERS))
+@pytest.mark.parametrize('strategy', ['Latency', 'Resource'])
+@pytest.mark.parametrize('backend', ['Vivado', 'Vitis'])
+def test_qrnn_weight_quantizers(test_case_id, layer, quantizers, strategy, backend):
+    """The kernel, recurrent and bias quantizers of a QKeras recurrent layer are applied to its weights, so hls4ml
+    computes with the same weights as QKeras. The three quantizers differ, so each weight has to get its own."""
+    quantizers = RNN_QUANTIZERS[quantizers]
+    model = make_qrnn_model(layer, quantizers)
+    X = np.round(np.random.default_rng(0).uniform(-0.5, 0.5, (50, 6, 4)) * 2**7) * 2**-7
+
+    config = hls4ml.utils.config_from_keras_model(
+        model, granularity='name', backend=backend, default_precision='ap_fixed<24,10>'
+    )
+    config['LayerName'][model.layers[1].name]['Strategy'] = strategy
+    if strategy == 'Resource':
+        # 4 inputs in 2 lanes over 3 cycles: the weights are padded with zeros
+        config['LayerName'][model.layers[1].name]['ReuseFactor'] = 3
+    output_dir = str(test_root_path / test_case_id)
+    hls_model = hls4ml.converters.convert_from_keras_model(model, hls_config=config, backend=backend, output_dir=output_dir)
+    hls_model.compile()
+
+    hls_layer = hls_model.graph[model.layers[1].name]
+    if strategy == 'Latency':  # the Resource strategy lays the weights out in the order its cell reads them
+        expected = {name: get_quantizer(q) for name, q in quantizers.items()}
+        if layer == 'QGRU':
+            expected['recurrent_bias'] = get_quantizer(quantizers['bias'])
+        assert_rnn_weights_quantized(hls_layer, expected, backend)
+    for w in hls_layer.weights.values():
+        if isinstance(w, ExponentWeightVariable):  # every exponent written to the project fits the type, zero padding too
+            assert w._format()[..., 1].min() >= -(2 ** (w.type.precision.width - 1))
+
+    y_qkeras = model.predict(X, verbose=0)
+    y_hls4ml = hls_model.predict(X).reshape(y_qkeras.shape)
+    np.testing.assert_allclose(y_hls4ml, y_qkeras, rtol=0, atol=0.05)
+
+
+def test_qlstm_altera_gate_quantizers(test_case_id):
+    """On Altera the LSTM weights are split per gate, and each piece gets the quantizer and the type of the array it
+    comes from: the kernel, the recurrent kernel or the bias. Conversion only; compiling needs the Altera toolchain."""
+    bits = {'weight': 8, 'recurrent_weight': 4, 'bias': 3}
+    quantizers = {name: f'quantized_bits({b}, 0, alpha=1)' for name, b in bits.items()}
+    model = make_qrnn_model('QLSTM', quantizers)
+    config = hls4ml.utils.config_from_keras_model(
+        model, granularity='name', backend='Altera', default_precision='ap_fixed<24,10>'
+    )
+    output_dir = str(test_root_path / test_case_id)
+    hls_model = hls4ml.converters.convert_from_keras_model(model, hls_config=config, backend='Altera', output_dir=output_dir)
+
+    hls_layer = hls_model.graph[model.layers[1].name]
+    assert_rnn_weights_quantized(hls_layer, {name: get_quantizer(q) for name, q in quantizers.items()}, 'Altera')
+    arrays = dict(zip(('weight', 'recurrent_weight', 'bias'), model.layers[1].get_weights()))
+    for i, gate in enumerate('ifco'):
+        for name, q in quantizers.items():
+            w = hls_layer.weights[f'{name}_{gate}']
+            expected = np.asarray(get_quantizer(q)(arrays[name][..., 8 * i : 8 * (i + 1)]))
+            np.testing.assert_array_equal(w.data, expected.T if expected.ndim == 2 else expected)
+            assert (w.type.precision.width, w.type.precision.integer) == (bits[name], 1)
 
 
 @pytest.mark.parametrize('rnn_layer', ['SimpleRNN', 'LSTM', 'GRU', 'Bidirectional'])