Remove unused channels w/ Keras Surgeon

bo3z · bo3z · commit e655ab604968 · 2023-06-13T11:53:11.000+01:00
diff --git a/hls4ml/optimization/keras/reduction.py b/hls4ml/optimization/keras/reduction.py
@@ -0,0 +1,59 @@
+import numpy as np
+from tensorflow.keras.models import Sequential
+from tensorflow.keras.layers import Dense, Conv2D
+from hls4ml.optimization.keras.utils import get_last_layer_with_weights
+
+'''
+Function for removing zero neurons & filters from a model and rewiring the model graph
+This function is built on top of Keras Surgeon available at: https://github.com/BenWhetton/keras-surgeon 
+Keras Surgeon is no longer under active development and does not work for TensorFlow 2.3+ and QKeras
+The baseline version was forked and updated, available at: https://github.com/bo3z/keras-surgeon
+
+Args:
+    - model (keras.model): Input model
+    
+Return:
+    - reduced (keras.model): Modified model, with redundant structures removed
+
+    '''
+def reduce_model(model):
+    # TODO - Should we make Keras Surgeon a hard requirement in setup.cfg? If so, needs to be installed from git, @bo3z fork
+    try:
+        from kerassurgeon import Surgeon
+    except ModuleNotFoundError:
+        raise Exception('Keras Surgeon not installed. Unable to reduce model footprint '\
+                         'Please install up-to-date Keras Surgeon compatible wit TensorFlow 2.3+ and QKeras '\
+                         'Installation from git: https://github.com/bo3z/keras-surgeon')
+    
+    # Initiate surgeon
+    surgeon = Surgeon(model)
+
+    # Iterate through layers and identify neurons (columns) and filters (tensors, W x H x C) to be removed
+    last_idx = get_last_layer_with_weights(model)
+    for idx, layer in enumerate(model.layers):
+        # Last layer with weights cannot be removed, as it maps to data set labels
+        if (idx == last_idx):
+            break
+
+        # Currently supported Dense and Conv2D; these two can be combined in a single if-statement
+        # Keras Surgeon has a full range of support for Conv1D / Conv3D, reucurrent etc. - might extend in the future
+        if isinstance(layer, Dense):
+            weights = layer.get_weights()[0]
+            zeros = np.where(~weights.any(axis=0))[0].tolist()
+            surgeon.add_job('delete_channels', layer, channels=zeros)
+
+        elif isinstance(layer, Conv2D):
+            weights = layer.get_weights()[0]
+            zeros = np.where(~weights.reshape(-1, weights.shape[-1]).any(axis=0))[0].tolist()
+            surgeon.add_job('delete_channels', layer, channels=zeros)
+    
+    # Reduce model
+    reduced = surgeon.operate()
+
+    # By default, Keras surgeon returns a Functional model
+    # If the original was a Sequential, convert back
+    is_sequential =  model.__class__.__name__ == 'Sequential'
+    if is_sequential:
+        return Sequential(layers=reduced.layers)
+    else:
+        return reduced
diff --git a/test/pytest/optimization/keras/test_reduction.py b/test/pytest/optimization/keras/test_reduction.py
@@ -0,0 +1,134 @@
+from qkeras import quantized_bits
+from qkeras import QDense, QActivation, QConv2D
+from tensorflow.keras.models import Sequential
+from tensorflow.keras.layers import Dense, Conv2D, MaxPooling2D, Softmax, BatchNormalization, ReLU, Flatten, AveragePooling2D
+from hls4ml.optimization.keras.reduction import reduce_model
+from hls4ml.optimization.keras.utils import get_model_sparsity
+
+'''
+Set some neurons / filters to zero and verify that these are removed
+Even is some neurons (columns) in the output layer are zero, these should not be removed (to match data set labels)
+Test verify the above property, by setting some zeros in the last layer and verifying these remain in place
+'''
+
+def test_keras_model_reduction():
+    model = Sequential()
+    model.add(Conv2D(8, (3, 3), input_shape=(64, 64, 1), name='conv2d_1', padding='same'))
+    model.add(MaxPooling2D())
+    model.add(BatchNormalization())
+    model.add(ReLU())
+    model.add(Conv2D(32, (5, 5), padding='same', name='conv2d_2'))
+    model.add(AveragePooling2D())
+    model.add(BatchNormalization())
+    model.add(ReLU())
+    model.add(Flatten())
+    model.add(Dense(32, input_shape=(16, ), name = 'dense_1', activation='relu'))
+    model.add(BatchNormalization())
+    model.add(Dense(14, name = 'dense_2', activation='relu'))
+    model.add(BatchNormalization())
+    model.add(Dense(5, name = 'dense_3'))
+    model.add(Softmax())
+
+    indices = {
+        'conv2d_1': [2, 4, 7],
+        'conv2d_2': [0, 1, 2, 3, 4, 5],
+        'dense_1': [0, 5, 17, 28],
+        'dense_2': [1, 9, 4],
+        'dense_3': [3],
+    }
+    for layer in model.layers:
+        if isinstance(layer, Dense):
+            weights = layer.get_weights()
+            weights[0][:, indices[layer.name]] = 0
+            layer.set_weights(weights)
+        if isinstance(layer, Conv2D):
+            weights = layer.get_weights()
+            weights[0][:, :, :, indices[layer.name]] = 0
+            layer.set_weights(weights)
+   
+    sparsity, _ = get_model_sparsity(model)
+    assert(sparsity > 0)
+
+    reduced = reduce_model(model)
+    assert(reduced.get_layer('conv2d_1').get_weights()[0].shape == (3, 3, 1, 5))
+    assert(reduced.get_layer('conv2d_2').get_weights()[0].shape == (5, 5, 5, 26))
+    assert(reduced.get_layer('dense_1').get_weights()[0].shape == (6656, 28))
+    assert(reduced.get_layer('dense_2').get_weights()[0].shape == (28, 11))
+    assert(reduced.get_layer('dense_3').get_weights()[0].shape == (11, 5))
+
+    _, layer_sparsity = get_model_sparsity(reduced)
+    assert(layer_sparsity['conv2d_1'] == 0)
+    assert(layer_sparsity['conv2d_2'] == 0)
+    assert(layer_sparsity['dense_1'] == 0)
+    assert(layer_sparsity['dense_2'] == 0)
+    assert(layer_sparsity['dense_3'] > 0)
+
+def test_qkeras_model_reduction():
+    bits = 8
+    activation = 'quantized_relu(4)'
+    quantizer = quantized_bits(bits, 0)
+    
+    model = Sequential()
+    model.add(QConv2D(8, (3, 3), input_shape=(64, 64, 1), name='qconv2d_1', padding='same', kernel_quantizer=quantizer))
+    model.add(MaxPooling2D())
+    model.add(BatchNormalization())
+    model.add(QActivation(activation, name='qrelu_1'))
+    model.add(QConv2D(32, (5, 5), padding='same', name='qconv2d_2', kernel_quantizer=quantizer))
+    model.add(AveragePooling2D())
+    model.add(BatchNormalization())
+    model.add(QActivation(activation, name='qrelu_2'))
+    model.add(Flatten())
+    model.add(QDense(32, input_shape=(16, ), name = 'qdense_1', kernel_quantizer=quantizer))
+    model.add(QActivation(activation, name='qrelu_3'))
+    model.add(BatchNormalization())
+    model.add(QDense(14, name = 'qdense_2', kernel_quantizer=quantizer))
+    model.add(QActivation(activation, name='qrelu_4'))
+    model.add(BatchNormalization())
+    model.add(QDense(5, name = 'qdense_3', kernel_quantizer=quantizer))
+    model.add(Softmax())
+
+    indices = {
+        'qconv2d_1': [2, 4, 7],
+        'qconv2d_2': [0, 1, 2, 3, 4, 5],
+        'qdense_1': [0, 5, 17, 28],
+        'qdense_2': [1, 9, 4],
+        'qdense_3': [3],
+    }
+    for layer in model.layers:
+        if isinstance(layer, QDense):
+            weights = layer.get_weights()
+            weights[0][:, indices[layer.name]] = 0
+            layer.set_weights(weights)
+        if isinstance(layer, QConv2D):
+            weights = layer.get_weights()
+            weights[0][:, :, :, indices[layer.name]] = 0
+            layer.set_weights(weights)
+   
+    sparsity, _ = get_model_sparsity(model)
+    assert(sparsity > 0)
+
+    reduced = reduce_model(model)
+    assert(reduced.get_layer('qconv2d_1').get_weights()[0].shape == (3, 3, 1, 5))
+    assert(reduced.get_layer('qconv2d_2').get_weights()[0].shape == (5, 5, 5, 26))
+    assert(reduced.get_layer('qdense_1').get_weights()[0].shape == (6656, 28))
+    assert(reduced.get_layer('qdense_2').get_weights()[0].shape == (28, 11))
+    assert(reduced.get_layer('qdense_3').get_weights()[0].shape == (11, 5))
+
+    _, layer_sparsity = get_model_sparsity(reduced)
+    assert(layer_sparsity['qconv2d_1'] == 0)
+    assert(layer_sparsity['qconv2d_2'] == 0)
+    assert(layer_sparsity['qdense_1'] == 0)
+    assert(layer_sparsity['qdense_2'] == 0)
+    assert(layer_sparsity['qdense_3'] > 0)
+
+    # Verify network surgery has no impact on quantization
+    assert(isinstance(reduced.get_layer('qrelu_1'), QActivation))
+    assert(isinstance(reduced.get_layer('qrelu_2'), QActivation))
+    assert(isinstance(reduced.get_layer('qrelu_3'), QActivation))
+    assert(isinstance(reduced.get_layer('qrelu_4'), QActivation))
+    assert(reduced.get_layer('qconv2d_1').kernel_quantizer['config']['bits'] == bits)
+    assert(reduced.get_layer('qconv2d_2').kernel_quantizer['config']['bits'] == bits)
+    assert(reduced.get_layer('qdense_1').kernel_quantizer['config']['bits'] == bits)
+    assert(reduced.get_layer('qdense_2').kernel_quantizer['config']['bits'] == bits)
+    assert(reduced.get_layer('qdense_3').kernel_quantizer['config']['bits'] == bits)
+    
