part of pre-commit of Optimization API pt.1

Author: bo3z

hls4ml/optimization/__init__.py

@@ -1,13 +1,35 @@
 import numpy as np
-from hls4ml.optimization.keras import optimize_model
+
 from hls4ml.optimization.attributes import get_attributes_from_keras_model_and_hls4ml_config
+from hls4ml.optimization.keras import optimize_model
+
 
 def optimize_keras_for_hls4ml(
-    keras_model, hls_config, objective, scheduler, X_train, y_train, X_val, y_val, 
-    batch_size, epochs, optimizer, loss_fn, validation_metric, increasing, rtol, 
-    callbacks=[], ranking_metric='l1', local=False, verbose=False, rewinding_epochs=1, cutoff_bad_trials=3,
-    directory='hls4ml-optimization', tuner='Bayesian', knapsack_solver='CBC_MIP',
-    regularization_range=np.logspace(-6, -2, num=16).tolist()
+    keras_model,
+    hls_config,
+    objective,
+    scheduler,
+    X_train,
+    y_train,
+    X_val,
+    y_val,
+    batch_size,
+    epochs,
+    optimizer,
+    loss_fn,
+    validation_metric,
+    increasing,
+    rtol,
+    callbacks=[],
+    ranking_metric='l1',
+    local=False,
+    verbose=False,
+    rewinding_epochs=1,
+    cutoff_bad_trials=3,
+    directory='hls4ml-optimization',
+    tuner='Bayesian',
+    knapsack_solver='CBC_MIP',
+    regularization_range=np.logspace(-6, -2, num=16).tolist(),
 ):
     '''
     Top-level function for optimizing a Keras model, given hls4ml config and a hardware objective(s)
@@ -28,7 +50,7 @@ def optimize_keras_for_hls4ml(
     - validation_metric (keras.metrics.Metric or equivalent loss description): Validation metric, used as a baseline
     - increasing (boolean): If the metric improves with increased values; e.g. accuracy -> increasing = True, MSE -> increasing = False
     - rtol (float): Relative tolerance; pruning stops when pruned_validation_metric < (or >) rtol * baseline_validation_metric
-    
+
     Kwargs:
     - callbacks (list of keras.callbacks.Callback) Currently not supported, developed in future versions
     - ranking_metric (string): Metric used for rannking weights and structures; currently supported l1, l2, saliency and Oracle
@@ -47,11 +69,29 @@ def optimize_keras_for_hls4ml(
 
     # Optimize model
     return optimize_model(
-        keras_model, model_attributes, objective, scheduler,
-        X_train, y_train, X_val, y_val, batch_size, epochs, 
-        optimizer, loss_fn, validation_metric, increasing, rtol, 
-        callbacks=callbacks, ranking_metric=ranking_metric, local=local, verbose=verbose, 
-        rewinding_epochs=rewinding_epochs, cutoff_bad_trials=cutoff_bad_trials,
-        directory=directory, tuner=tuner, knapsack_solver=knapsack_solver,
-        regularization_range=regularization_range
+        keras_model,
+        model_attributes,
+        objective,
+        scheduler,
+        X_train,
+        y_train,
+        X_val,
+        y_val,
+        batch_size,
+        epochs,
+        optimizer,
+        loss_fn,
+        validation_metric,
+        increasing,
+        rtol,
+        callbacks=callbacks,
+        ranking_metric=ranking_metric,
+        local=local,
+        verbose=verbose,
+        rewinding_epochs=rewinding_epochs,
+        cutoff_bad_trials=cutoff_bad_trials,
+        directory=directory,
+        tuner=tuner,
+        knapsack_solver=knapsack_solver,
+        regularization_range=regularization_range,
     )

hls4ml/optimization/attributes.py

@@ -1,8 +1,10 @@
-import hls4ml
 import numpy as np
+
+import hls4ml
+from hls4ml.model.types import FixedPrecisionType, IntegerPrecisionType
 from hls4ml.optimization.config import SUPPORTED_STRUCTURES
 from hls4ml.optimization.keras.config import SUPPORTED_LAYERS
-from hls4ml.model.types import FixedPrecisionType, IntegerPrecisionType
+
 
 class hls4mlAttributes:
     '''
@@ -18,19 +20,22 @@ class hls4mlAttributes:
     - reuse_factor (int): Layer reuse factor
     - parallelization_factor (int): Layer parallelization factor - [applicable to io_parallel Conv2D]
     '''
-    def __init__(self, n_in, n_out, io_type, strategy, weight_precision, output_precision, reuse_factor, parallelization_factor=1):    
+
+    def __init__(
+        self, n_in, n_out, io_type, strategy, weight_precision, output_precision, reuse_factor, parallelization_factor=1
+    ):
         if not isinstance(weight_precision, (FixedPrecisionType, IntegerPrecisionType)):
             raise Exception('Layer weight precision is not in valid format')
-    
+
         if not isinstance(output_precision, (FixedPrecisionType, IntegerPrecisionType)):
             raise Exception('Layer weight precision is not in valid format')
 
         if not strategy in ('Latency', 'latency', 'Resource', 'resource'):
             raise Exception('Unknown layer strategy')
-        
+
         if not io_type in ('io_parallel', 'io_stream'):
             raise Exception('Unknown IO type')
-        
+
         self.n_in = n_in
         self.n_out = n_out
         self.io_type = io_type
@@ -40,37 +45,48 @@ def __init__(self, n_in, n_out, io_type, strategy, weight_precision, output_prec
         self.reuse_factor = reuse_factor
         self.parallelization_factor = parallelization_factor
 
+
 class OptimizationAttributes:
     '''
     A class for storing layer optimization attributes
-    
+
     Args:
         - structure_type (enum): Targeted structure - unstructured, structured, pattern, block
         - pruning (boolean): Should pruning be applied to the layer
         - weight_sharing (boolean): Should weight sharing be applied to the layer
         - block_shape (tuple): Block shape if structure_type == block
-        - pattern_offset (int): Length of each pattern if structure_type == pattern 
-        - consecutive_patterns (int): How many consecutive patterns are grouped together if structure_type == pattern 
-    
+        - pattern_offset (int): Length of each pattern if structure_type == pattern
+        - consecutive_patterns (int): How many consecutive patterns are grouped together if structure_type == pattern
+
     Notes:
         - In the case of hls4ml, pattern_offset is equivalent to the number of weights processed in parallel
         - The pattern_offset is n_in * n_out / reuse_factor; default case (=1) is equivalent to no unrolling
     '''
-    def __init__(self, structure_type=SUPPORTED_STRUCTURES.UNSTRUCTURED, pruning=False, weight_sharing=False, block_shape=(1, 1), pattern_offset=1, consecutive_patterns=1):
+
+    def __init__(
+        self,
+        structure_type=SUPPORTED_STRUCTURES.UNSTRUCTURED,
+        pruning=False,
+        weight_sharing=False,
+        block_shape=(1, 1),
+        pattern_offset=1,
+        consecutive_patterns=1,
+    ):
         if not isinstance(structure_type, SUPPORTED_STRUCTURES):
             raise Exception(f'{self.__class__.__name__} unknown structure type')
-    
+
         self.structure_type = structure_type
         self.pruning = pruning
         self.weight_sharing = weight_sharing
         self.block_shape = block_shape
         self.pattern_offset = pattern_offset
         self.consecutive_patterns = consecutive_patterns
 
+
 class LayerAttributes:
     '''
     A class for storing layer information
-    
+
     Args:
         - name (string): Layer name
         - layer_type (keras.Layer): Layer type (e.g. Dense, Conv2D etc.)
@@ -82,77 +98,99 @@ class LayerAttributes:
         - optimization_attributes (OptimizationAttributes): Type of optimization, optimization vs weight sharing, block shape and pattern offset
         - args (dict): Additional information, e.g. hls4mlAttributes; dictionary so it can be generic enough for different platforms
     '''
-    def __init__(self, name, layer_type, inbound_layers, weight_shape, input_shape,output_shape,optimizable, optimization_attributes, args):
+
+    def __init__(
+        self,
+        name,
+        layer_type,
+        inbound_layers,
+        weight_shape,
+        input_shape,
+        output_shape,
+        optimizable,
+        optimization_attributes,
+        args,
+    ):
         self.name = name
         self.layer_type = layer_type
         self.inbound_layers = inbound_layers
         self.weight_shape = weight_shape
         self.input_shape = input_shape
         self.output_shape = output_shape
         self.optimizable = optimizable
-        self.optimization_attributes = optimization_attributes 
+        self.optimization_attributes = optimization_attributes
         self.args = args
 
     def update_args(self, updates):
         self.args.update(updates)
 
     def __str__(self):
-        return  f'name: {self.name}, ' \
-                f'layer_type: {self.layer_type}, ' \
-                f'inbound_layers: {self.inbound_layers}, ' \
-                f'weight_shape: {self.weight_shape}, ' \
-                f'input_shape: {self.input_shape}, ' \
-                f'output_shape: {self.output_shape}, ' \
-                f'optimizable: {self.optimizable}, ' \
-                f'optimization_attributes: {self.optimization_attributes}, ' \
-                f'args: {self.args}, ' \
+        return (
+            f'name: {self.name}, '
+            f'layer_type: {self.layer_type}, '
+            f'inbound_layers: {self.inbound_layers}, '
+            f'weight_shape: {self.weight_shape}, '
+            f'input_shape: {self.input_shape}, '
+            f'output_shape: {self.output_shape}, '
+            f'optimizable: {self.optimizable}, '
+            f'optimization_attributes: {self.optimization_attributes}, '
+            f'args: {self.args}, '
+        )
+
 
 def get_attributes_from_keras_model(model):
     '''
     Given a Keras model, builds a dictionary of class attributes
     Additional arguments (e.g. reuse factor), depend on the target hardware platform and are inserted later
     Per-layer pruning sype (structured, pattern etc.), depend on the pruning objective and are inserted later
-    
+
     Args:
         - model (keras.model): Model to extract attributes from
 
     Return:
-        - model_attributes (dict): Each key corresponds to a layer name, values are instances of LayerAttribute 
+        - model_attributes (dict): Each key corresponds to a layer name, values are instances of LayerAttribute
     '''
-    is_sequential =  model.__class__.__name__ == 'Sequential'
+    is_sequential = model.__class__.__name__ == 'Sequential'
     model_attributes = {}
 
     for i, layer in enumerate(model.layers):
         inbound_layers = []
         if is_sequential and i > 0:
-            inbound_layers.append(model.layers[i-1])
+            inbound_layers.append(model.layers[i - 1])
         elif not is_sequential:
             nodes = model.get_config()['layers'][i]['inbound_nodes']
             if len(nodes) > 0:
                 inbound_layers.append(node[0] for node in nodes[0])
-        
+
         layer_weights = layer.get_weights()
         weight_shape = layer_weights[0].shape if len(layer_weights) > 0 else ()
 
         model_attributes[layer.name] = LayerAttributes(
-            layer.name, layer.__class__, inbound_layers,
-            weight_shape, layer.input_shape[1:], layer.output_shape[1:],
-            False, OptimizationAttributes(), {}
+            layer.name,
+            layer.__class__,
+            inbound_layers,
+            weight_shape,
+            layer.input_shape[1:],
+            layer.output_shape[1:],
+            False,
+            OptimizationAttributes(),
+            {},
         )
-    
+
     return model_attributes
 
+
 def get_attributes_from_keras_model_and_hls4ml_config(model, config):
     '''
     Given a Keras model and hls4ml configuration, builds a dictionary of class attributes
     Per-layer pruning sype (structured, pruning etc.), depend on the pruning objective and are inserted later
-    
+
     Args:
         - model (keras.model): Model to extract attributes from
-        - config (dict): hls4ml dictionary 
-    
+        - config (dict): hls4ml dictionary
+
     Return:
-        - model_attributes (dict): Each key corresponds to a layer name, values are LayerAttribute instances 
+        - model_attributes (dict): Each key corresponds to a layer name, values are LayerAttribute instances
     '''
 
     # Extract Keras attributes
@@ -163,30 +201,32 @@ def get_attributes_from_keras_model_and_hls4ml_config(model, config):
     default_reuse_factor = config['Model']['ReuseFactor']
     default_strategy = config['Model']['Strategy']
     default_precision = config['Model']['Precision']
-    
+
     # Build dictionary
     for layer in model_attributes:
         if model_attributes[layer].layer_type in SUPPORTED_LAYERS:
             n_in, n_out = __get_layer_mult_size(model_attributes[layer])
             layer_config = config['LayerName'][layer] if layer in config['LayerName'] else {}
-            reuse_factor = layer_config['ReuseFactor'] if 'ReuseFactor' in layer_config else default_reuse_factor 
-            parallelization_factor =  layer_config['ParallelizationFactor'] if 'ParallelizationFactor' in layer_config else 1
+            reuse_factor = layer_config['ReuseFactor'] if 'ReuseFactor' in layer_config else default_reuse_factor
+            parallelization_factor = layer_config['ParallelizationFactor'] if 'ParallelizationFactor' in layer_config else 1
             strategy = layer_config['Strategy'] if 'Strategy' in layer_config else default_strategy
-            weight_precision = layer_config['Precision']['weight'] if 'weight' in layer_config['Precision'] else default_precision
+            weight_precision = (
+                layer_config['Precision']['weight'] if 'weight' in layer_config['Precision'] else default_precision
+            )
             weight_precision = hls4ml.backends.fpga.fpga_backend.FPGABackend.convert_precision_string(weight_precision)
-            output_precision = layer_config['Precision']['result'] if 'result' in layer_config['Precision'] else default_precision
+            output_precision = (
+                layer_config['Precision']['result'] if 'result' in layer_config['Precision'] else default_precision
+            )
             output_precision = hls4ml.backends.fpga.fpga_backend.FPGABackend.convert_precision_string(output_precision)
-        
+
             hls4ml_attributes = hls4mlAttributes(
-                n_in, n_out,
-                io_type, strategy,
-                weight_precision, output_precision,
-                reuse_factor, parallelization_factor
+                n_in, n_out, io_type, strategy, weight_precision, output_precision, reuse_factor, parallelization_factor
             )
             model_attributes[layer].update_args({'hls4ml_attributes': hls4ml_attributes})
-    
+
     return model_attributes
 
+
 def __get_layer_mult_size(attributes):
     '''
     Helper function to calculate layer multiplication size

hls4ml/optimization/config.py

@@ -16,7 +16,7 @@
     - Description: Zeroes out or quantizes all the weights in a structure:
         - Dense: Neurons, determined by their outgoing connections (columns in Keras weight tensors)
         - Conv2D: Filters (structures of size filt_width x filt_height x n_chan)
-        - Notes: 
+        - Notes:
             - For Dense, it was also possible optimize by incoming connections (rows); however, removing zero neurons becomes harder
             - For Conv2D, significant literature explored pruning channels; currently not supported
     - Supports: All layers in SUPPORTED_LAYERS (hls4ml.optimization.keras)
@@ -36,8 +36,10 @@
     - Supports: All rank-2 (e.g. Dense, but not Conv2D) layers in SUPPORTED_LAYERS (hls4ml.optimization.keras)
 
 '''
+
+
 class SUPPORTED_STRUCTURES(Enum):
     UNSTRUCTURED = 'unstructured'
-    STRUCTURED   = 'structured'
-    PATTERN      = 'pattern'
-    BLOCK        = 'block' 
+    STRUCTURED = 'structured'
+    PATTERN = 'pattern'
+    BLOCK = 'block'