[pre-commit.ci] auto fixes from pre-commit hooks

Author: pre-commit-ci[bot]

nn_utils.py

@@ -1,31 +1,23 @@
-
-
-from io import BytesIO
 import json
 import os
+import pickle as pkl
 import random
-
 from glob import glob
+from io import BytesIO
 from pathlib import Path
-import pickle as pkl
 from typing import Callable
 
+import h5py as h5
 import numpy as np
-
 import tensorflow as tf
-from tensorflow import keras
-
-from tqdm.auto import tqdm
-from matplotlib import pyplot as plt
-
 import zstd
-import h5py as h5
-
-from keras.src.saving.legacy import hdf5_format
-from keras.src.layers.convolutional.base_conv import Conv
-from keras.layers import Dense
-
 from HGQ.bops import trace_minmax
+from keras.layers import Dense
+from keras.src.layers.convolutional.base_conv import Conv
+from keras.src.saving.legacy import hdf5_format
+from matplotlib import pyplot as plt
+from tensorflow import keras
+from tqdm.auto import tqdm
 
 
 class NumpyFloatValuesEncoder(json.JSONEncoder):
@@ -36,14 +28,15 @@ def default(self, obj):
 
 
 class SaveTopN(keras.callbacks.Callback):
-    def __init__(self,
-                 metric_fn: Callable[[dict], float],
-                 n: int,
-                 path: str | Path,
-                 side: str = 'max',
-                 fname_format='epoch={epoch}-metric={metric:.4e}.h5',
-                 cond_fn: Callable[[dict], bool] = lambda x: True,
-                 ):
+    def __init__(
+        self,
+        metric_fn: Callable[[dict], float],
+        n: int,
+        path: str | Path,
+        side: str = 'max',
+        fname_format='epoch={epoch}-metric={metric:.4e}.h5',
+        cond_fn: Callable[[dict], bool] = lambda x: True,
+    ):
         self.n = n
         self.metric_fn = metric_fn
         self.path = Path(path)
@@ -188,9 +181,11 @@ def absorb_batchNorm(model_target, model_original):
         if layer.__class__.__name__ == 'Functional':
             absorb_batchNorm(layer, model_original.get_layer(layer.name))
             continue
-        if (isinstance(layer, Dense) or isinstance(layer, Conv)) and \
-            len(nodes := model_original.get_layer(layer.name)._outbound_nodes) > 0 and \
-                isinstance(nodes[0].outbound_layer, keras.layers.BatchNormalization):
+        if (
+            (isinstance(layer, Dense) or isinstance(layer, Conv))
+            and len(nodes := model_original.get_layer(layer.name)._outbound_nodes) > 0
+            and isinstance(nodes[0].outbound_layer, keras.layers.BatchNormalization)
+        ):
             _gamma, _beta, _mu, _var = model_original.get_layer(layer.name)._outbound_nodes[0].outbound_layer.get_weights()
             _ratio = _gamma / np.sqrt(0.001 + _var)
             _bias = -_gamma * _mu / np.sqrt(0.001 + _var) + _beta
@@ -213,7 +208,7 @@ def absorb_batchNorm(model_target, model_original):
             weights = layer.get_weights()
             new_weights = model_original.get_layer(layer.name).get_weights()
             l = len(new_weights)
-            layer.set_weights([*new_weights, *weights[l:]][:len(weights)])
+            layer.set_weights([*new_weights, *weights[l:]][: len(weights)])
 
 
 def set_seed(seed):
@@ -225,9 +220,10 @@ def set_seed(seed):
     tf.config.experimental.enable_op_determinism()
 
 
-import h5py as h5
 import json
 
+import h5py as h5
+
 
 def get_best_ckpt(save_path: Path, take_min=False):
     ckpts = list(save_path.glob('*.h5'))
@@ -245,13 +241,14 @@ def rank(ckpt: Path):
 
 
 class PeratoFront(keras.callbacks.Callback):
-    def __init__(self,
-                 path: str | Path,
-                 fname_format: str,
-                 metrics_names: list[str],
-                 sides: list[int],
-                 cond_fn: Callable[[dict], bool] = lambda x: True,
-                 ):
+    def __init__(
+        self,
+        path: str | Path,
+        fname_format: str,
+        metrics_names: list[str],
+        sides: list[int],
+        cond_fn: Callable[[dict], bool] = lambda x: True,
+    ):
         self.path = Path(path)
         self.fname_format = fname_format
         os.makedirs(path, exist_ok=True)

part4.1_HG_quantization.ipynb

@@ -75,7 +75,7 @@
     "classes = np.load('classes.npy', allow_pickle=True)\n",
     "\n",
     "# Convert everything to tf.Tensor to avoid casting\n",
-    "with tf.device('/cpu:0'): # type: ignore\n",
+    "with tf.device('/cpu:0'):  # type: ignore\n",
     "    _X_train_val = tf.convert_to_tensor(X_train_val, dtype=tf.float32)\n",
     "    # We don't make explicit y categorical tensor:\n",
     "    # Use SparseCategoricalCrossentropy loss instead.\n",
@@ -128,13 +128,15 @@
     "beta = 3e-6\n",
     "# The bigger the beta, the smaller the models is, at the cost of accuracy.\n",
     "\n",
-    "model = Sequential([\n",
-    "    HQuantize(beta=beta),\n",
-    "    HDense(64, activation='relu', beta=beta),\n",
-    "    HDense(32, activation='relu', beta=beta),\n",
-    "    HDense(32, activation='relu', beta=beta),\n",
-    "    HDense(5, beta=beta),\n",
-    "])"
+    "model = Sequential(\n",
+    "    [\n",
+    "        HQuantize(beta=beta),\n",
+    "        HDense(64, activation='relu', beta=beta),\n",
+    "        HDense(32, activation='relu', beta=beta),\n",
+    "        HDense(32, activation='relu', beta=beta),\n",
+    "        HDense(5, beta=beta),\n",
+    "    ]\n",
+    ")"
    ]
   },
   {
@@ -218,7 +220,7 @@
     "    opt = Adam(learning_rate=0)\n",
     "    loss = SparseCategoricalCrossentropy(from_logits=True)\n",
     "    model.compile(optimizer=opt, loss=loss, metrics=['accuracy'])\n",
-    "    \n",
+    "\n",
     "    model.fit(\n",
     "        _X_train_val,\n",
     "        _y_train_val,\n",
@@ -227,13 +229,14 @@
     "        validation_split=0.25,\n",
     "        shuffle=True,\n",
     "        callbacks=callbacks,\n",
-    "        verbose=0, # type: ignore\n",
+    "        verbose=0,  # type: ignore\n",
     "    )\n",
     "    model.save('model_3.1/model.h5')\n",
     "else:\n",
     "    from keras.models import load_model\n",
+    "\n",
     "    # No need to use custom_objects as the custom layers are already registered\n",
-    "    model:keras.Model = load_model('model_3.1/model.h5') # type: ignore\n"
+    "    model: keras.Model = load_model('model_3.1/model.h5')  # type: ignore"
    ]
   },
   {
@@ -256,6 +259,7 @@
    "outputs": [],
    "source": [
     "from HGQ import trace_minmax, to_proxy_model\n",
+    "\n",
     "trace_minmax(model, X_train_val)"
    ]
   },
@@ -303,9 +307,7 @@
     "import hls4ml\n",
     "import plotting\n",
     "\n",
-    "hls_model = hls4ml.converters.convert_from_keras_model(\n",
-    "    proxy, output_dir='model_3.1/hls4ml_prj', part='xcu250-figd2104-2L-e'\n",
-    ")\n",
+    "hls_model = hls4ml.converters.convert_from_keras_model(proxy, output_dir='model_3.1/hls4ml_prj', part='xcu250-figd2104-2L-e')\n",
     "hls_model.compile()\n",
     "\n",
     "X_test = np.ascontiguousarray(X_test)\n",