Merge pull request #3 from vloncar/opt1

Cleanup docstrings

Author: bo3z

docs/advanced/model_optimization.rst

@@ -1,13 +1,14 @@
-========================
-hls4ml Optimization API
-========================
+=================================
+Hardware-aware Optimization API
+=================================
 
 Pruning and weight sharing are effective techniques to reduce model footprint and computational requirements. The hls4ml Optimization API introduces hardware-aware pruning and weight sharing.
 By defining custom objectives, the algorithm solves a Knapsack optimization problem aimed at maximizing model performance, while keeping the target resource(s) at a minimum. Out-of-the box objectives include network sparsity, GPU FLOPs, Vivado DSPs, memory utilization etc.
 
 The code block below showcases three use cases of the hls4ml Optimization API - network sparsity (unstructured pruning), GPU FLOPs (structured pruning) and Vivado DSP utilization (pattern pruning). First, we start with unstructured pruning:
 
 .. code-block:: Python
+
     from sklearn.metrics import accuracy_score
     from tensorflow.keras.optimizers import Adam
     from tensorflow.keras.metrics import CategoricalAccuracy
@@ -71,7 +72,9 @@ In a similar manner, it is possible to target GPU FLOPs or Vivado DSPs. However,
 Instead, it is the sparsity of the target resource. As an example: Starting with a network utilizing 512 DSPs and a final sparsity of 50%; the optimized network will use 256 DSPs.
 
 To optimize GPU FLOPs, the code is similar to above:
+
 .. code-block:: Python
+
     from hls4ml.optimization.objectives.gpu_objectives import GPUFLOPEstimator
 
     # Optimize model
@@ -91,7 +94,9 @@ To optimize GPU FLOPs, the code is similar to above:
     print(optimized_model.summary())
 
 Finally, optimizing Vivado DSPs is possible, given a hls4ml config:
+
 .. code-block:: Python
+
     from hls4ml.utils.config import config_from_keras_model
     from hls4ml.optimization.objectives.vivado_objectives import VivadoDSPEstimator
 
@@ -121,7 +126,9 @@ Finally, optimizing Vivado DSPs is possible, given a hls4ml config:
 
 There are two more Vivado "optimizers" - VivadoFFEstimator, aimed at reducing register utilisation and VivadoMultiObjectiveEstimator, aimed at optimising BRAM and DSP utilisation.
 Note, to ensure DSPs are optimized, "unrolled" Dense multiplication must be used before synthesing HLS, by modifying the config:
+
 .. code-block:: Python
+
     hls_config = config_from_keras_model(optimized_model)
     hls_config['Model']['DenseResourceImplementation'] = 'Unrolled'
     # Any addition hls4ml config, such as strategy, reuse factor etc...

docs/index.rst

@@ -25,6 +25,7 @@
     advanced/fifo_depth
     advanced/extension
     advanced/accelerator
+    advanced/model_optimization
 
 .. toctree::
     :hidden:
@@ -34,6 +35,7 @@
     autodoc/hls4ml.backends
     autodoc/hls4ml.converters
     autodoc/hls4ml.model
+    autodoc/hls4ml.optimization
     autodoc/hls4ml.report
     autodoc/hls4ml.utils
     autodoc/hls4ml.writer

hls4ml/optimization/__init__.py

@@ -37,41 +37,42 @@ def optimize_keras_for_hls4ml(
     Top-level function for optimizing a Keras model, given hls4ml config and a hardware objective(s)
 
     Args:
-    - keras_model (keras.Model): Model to be optimized
-    - hls_config (dict): hls4ml configuration, obtained from hls4ml.utils.config.config_from_keras_model(...)
-    - objective (hls4ml.optimization.objectives.ObjectiveEstimator):
+        keras_model (keras.Model): Model to be optimized
+        hls_config (dict): hls4ml configuration, obtained from hls4ml.utils.config.config_from_keras_model(...)
+        objective (hls4ml.optimization.objectives.ObjectiveEstimator):
         Parameter, hardware or user-defined objective of optimization
-    - scheduler (hls4ml.optimization.schduler.OptimizationScheduler):
+        scheduler (hls4ml.optimization.scheduler.OptimizationScheduler):
         Sparsity scheduler, choose between constant, polynomial and binary
-    - X_train (np.array): Training inputs
-    - y_train (np.array): Training labels
-    - X_val (np.array): Validation inputs
-    - y_val (np.array): Validation labels
-    - batch_size (int): Batch size during training
-    - epochs (int): Maximum number of epochs to fine-tune model, in one iteration of pruning
-    - optimizer (keras.optimizers.Optimizer or equivalent-string description): Optimizer used during training
-    - loss_fn (keras.losses.Loss or equivalent loss description): Loss function used during training
-    - 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
+        X_train (np.array): Training inputs
+        y_train (np.array): Training labels
+        X_val (np.array): Validation inputs
+        y_val (np.array): Validation labels
+        batch_size (int): Batch size during training
+        epochs (int): Maximum number of epochs to fine-tune model, in one iteration of pruning
+        optimizer (keras.optimizers.Optimizer or equivalent-string description): Optimizer used during training
+        loss_fn (keras.losses.Loss or equivalent loss description): Loss function used during training
+        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
+        callbacks (list of keras.callbacks.Callback) Currently not supported, developed in future versions
+        ranking_metric (string): Metric used for ranking weights and structures;
+            currently supported l1, l2, saliency and Oracle
+        local (boolean): Layer-wise or global pruning
+        verbose (boolean): Display debug logs during model optimization
+        rewinding_epochs (int): Number of epochs to retrain model without weight freezing,
+            allows regrowth of previously pruned weights
+        cutoff_bad_trials (int): After how many bad trials (performance below threshold),
+            should model pruning / weight sharing stop
+        directory (string): Directory to store temporary results
+        tuner (str): Tuning algorithm, choose between Bayesian, Hyperband and None
+        knapsack_solver (str): Algorithm to solve Knapsack problem when optimizing;
+            default usually works well; for very large networks, greedy algorithm might be more suitable
+        regularization_range (list): List of suitable hyperparameters for weight decay
 
-    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
-    - local (boolean): Layer-wise or global pruning
-    - verbose (boolean): Display debug logs during model optimization
-    - rewinding_epochs (int): Number of epochs to retrain model without weight freezing,
-        allows regrowth of previously pruned weights
-    - cutoff_bad_trials (int): After how many bad trials (performance below threshold),
-        should model pruning / weight sharing stop
-    - directory (string): Directory to store temporary results
-    - tuner (str): Tuning alogorithm, choose between Bayesian, Hyperband and None
-    - knapsack_solver (str): Algorithm to solve Knapsack problem when optimizing;
-        default usually works well; for very large networks, greedy algorithm might be more suitable
-    - regularization_range (list): List of suitable hyperparameters for weight decay
+    Returns:
+        keras.Model: Optimized model
     '''
 
     # Extract model attributes

hls4ml/optimization/attributes.py

@@ -11,14 +11,14 @@ class hls4mlAttributes:
     A class for storing hls4ml information of a single layer
 
     Args:
-    - n_in (int): Number of inputs (rows) for Dense matrix multiplication
-    - n_out (int): Number of outputs (cols) for Dense matrix multiplication
-    - io_type (string): io_parallel or io_stream
-    - strategy (string): Resource or Latency
-    - weight_precision (FixedPrecisionType): Layer weight precision
-    - output_precision (FixedPrecisionType): Layer output precision
-    - reuse_factor (int): Layer reuse factor
-    - parallelization_factor (int): Layer parallelization factor - [applicable to io_parallel Conv2D]
+        n_in (int): Number of inputs (rows) for Dense matrix multiplication
+        n_out (int): Number of outputs (cols) for Dense matrix multiplication
+        io_type (string): io_parallel or io_stream
+        strategy (string): Resource or Latency
+        weight_precision (FixedPrecisionType): Layer weight precision
+        output_precision (FixedPrecisionType): Layer output precision
+        reuse_factor (int): Layer reuse factor
+        parallelization_factor (int): Layer parallelization factor - [applicable to io_parallel Conv2D]
     '''
 
     def __init__(
@@ -51,12 +51,12 @@ 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
+        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
 
     Notes:
         - In the case of hls4ml, pattern_offset is equivalent to the number of weights processed in parallel
@@ -88,16 +88,16 @@ class LayerAttributes:
     A class for storing layer information
 
     Args:
-        - name (string): Layer name
-        - layer_type (keras.Layer): Layer type (e.g. Dense, Conv2D etc.)
-        - inbound_layers (list): List of parent nodes, identified by name
-        - weight_shape (tuple): Layer weight shape
-        - input_shape (tuple): Layer input shape
-        - output_shape (tuple): Layer output shape
-        - optimizable (bool): Should optimizations (pruning, weight sharing) be applied to this layer
-        - optimization_attributes (OptimizationAttributes): Type of optimization,
+        name (string): Layer name
+        layer_type (keras.Layer): Layer type (e.g. Dense, Conv2D etc.)
+        inbound_layers (list): List of parent nodes, identified by name
+        weight_shape (tuple): Layer weight shape
+        input_shape (tuple): Layer input shape
+        output_shape (tuple): Layer output shape
+        optimizable (bool): Should optimizations (pruning, weight sharing) be applied to this layer
+        optimization_attributes (OptimizationAttributes): Type of optimization,
             pruning or weight sharing, block shape and pattern offset
-        - args (dict): Additional information,
+        args (dict): Additional information,
             e.g. hls4mlAttributes; dictionary so it can be generic enough for different platforms
     '''
 
@@ -147,10 +147,10 @@ def get_attributes_from_keras_model(model):
     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
+        model (keras.model): Model to extract attributes from
 
-    Return:
-        - model_attributes (dict): Each key corresponds to a layer name, values are instances of LayerAttribute
+    Returns:
+        model_attributes (dict): Each key corresponds to a layer name, values are instances of LayerAttribute
     '''
     is_sequential = model.__class__.__name__ == 'Sequential'
     model_attributes = {}
@@ -188,11 +188,11 @@ def get_attributes_from_keras_model_and_hls4ml_config(model, config):
     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
+        model (keras.model): Model to extract attributes from
+        config (dict): hls4ml dictionary
 
-    Return:
-        - model_attributes (dict): Each key corresponds to a layer name, values are LayerAttribute instances
+    Returns:
+        model_attributes (dict): Each key corresponds to a layer name, values are LayerAttribute instances
     '''
 
     # Extract Keras attributes

hls4ml/optimization/keras/__init__.py

@@ -50,45 +50,43 @@ def optimize_model(
     Top-level function for optimizing a Keras model, given objectives
 
     Args:
-    - model (keras.Model): Model to be optimized
-    - model_attributes (dict): Layer-wise model attributes,
-        obtained from hls4ml.optimization.get_attributes_from_keras_model(...)
-    - objective (hls4ml.optimization.objectives.ObjectiveEstimator):
-        Parameter, hardware or user-defined objective of optimization
-    - scheduler (hls4ml.optimization.schduler.OptimizationScheduler):
-        Sparsity scheduler, choose between constant, polynomial and binary
-    - X_train (np.array): Training inputs
-    - y_train (np.array): Training labels
-    - X_val (np.array): Validation inputs
-    - y_val (np.array): Validation labels
-    - batch_size (int): Batch size during training
-    - epochs (int): Maximum number of epochs to fine-tune model, in one iteration of pruning
-    - optimizer (keras.optimizers.Optimizer or equivalent-string description):
-        Optimizer used during training
-    - loss_fn (keras.losses.Loss or equivalent loss description):
-        Loss function used during training
-    - 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
-    - local (boolean): Layer-wise or global pruning
-    - verbose (boolean): Display debug logs during model optimization
-    - rewinding_epochs (int): Number of epochs to retrain model without weight freezing,
-        allows regrowth of previously pruned weights
-    - cutoff_bad_trials (int): After how many bad trials (performance below threshold),
-        should model pruning / weight sharing stop
-    - directory (string): Directory to store temporary results
-    - tuner (str): Tuning alogorithm, choose between Bayesian, Hyperband and None
-    - knapsack_solver (str): Algorithm to solve Knapsack problem when optimizing;
-        default usually works well; for very large networks, greedy algorithm might be more suitable
-    - regularization_range (list): List of suitable hyperparameters for weight decay
+        model (keras.Model): Model to be optimized
+        model_attributes (dict): Layer-wise model attributes,
+            obtained from hls4ml.optimization.get_attributes_from_keras_model(...)
+        objective (hls4ml.optimization.objectives.ObjectiveEstimator):
+            Parameter, hardware or user-defined objective of optimization
+        scheduler (hls4ml.optimization.scheduler.OptimizationScheduler):
+            Sparsity scheduler, choose between constant, polynomial and binary
+        X_train (np.array): Training inputs
+        y_train (np.array): Training labels
+        X_val (np.array): Validation inputs
+        y_val (np.array): Validation labels
+        batch_size (int): Batch size during training
+        epochs (int): Maximum number of epochs to fine-tune model, in one iteration of pruning
+        optimizer (keras.optimizers.Optimizer or equivalent-string description): Optimizer used during training
+        loss_fn (keras.losses.Loss or equivalent loss description): Loss function used during training
+        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
+        callbacks (list of keras.callbacks.Callback) Currently not supported, developed in future versions
+        ranking_metric (string): Metric used for ranking weights and structures;
+            currently supported l1, l2, saliency and Oracle
+        local (boolean): Layer-wise or global pruning
+        verbose (boolean): Display debug logs during model optimization
+        rewinding_epochs (int): Number of epochs to retrain model without weight freezing,
+            allows regrowth of previously pruned weights
+        cutoff_bad_trials (int): After how many bad trials (performance below threshold),
+            should model pruning / weight sharing stop
+        directory (string): Directory to store temporary results
+        tuner (str): Tuning algorithm, choose between Bayesian, Hyperband and None
+        knapsack_solver (str): Algorithm to solve Knapsack problem when optimizing;
+            default usually works well; for very large networks, greedy algorithm might be more suitable
+        regularization_range (list): List of suitable hyperparameters for weight decay
+
+    Returns:
+        keras.Model: Optimized model
     '''
 
     if not isinstance(scheduler, OptimizationScheduler):
@@ -213,7 +211,7 @@ def optimize_model(
 
         # Mask gradients
         # Before training the model at the next sparsity level, reset internal states
-        # Furthemore, modern optimizers (e.g. Adam) accumulate gradients during backprop
+        # Furthermore, modern optimizers (e.g. Adam) accumulate gradients during backprop
         # Therefore, even if the gradient for a weight is zero, it might be updated, due to previous gradients
         # Avoid this by resetting the internal variables of an optimizer
         optimizable_model.reset_metrics()
@@ -329,7 +327,7 @@ def __call__(self, X, y, s):
             - y (tf.Tensor): Output data
             - s (float): Sparsity
 
-        Return:
+        Returns:
             - loss (tf.Varilable): Model loss with input X and output y
         '''
         grads = []