make Brevitas use FixedPoint quantization, add back bias

Author: jmitrevs

2_quantization/2b_brevitas.ipynb

@@ -82,8 +82,9 @@
     "- `QuantLinear` — linear layer with quantized weights (and optionally biases)\n",
     "- `QuantReLU` — ReLU that clips and quantizes activations to a fixed bit-width\n",
     "\n",
-    "We use `weight_bit_width=6` and `bit_width=6` throughout, matching the 6-bit precision from Part 2a.\n",
-    "Unlike regular PyTorch, every forward pass applies the quantization, so the model trains with the actual fixed-point precisions used on the FPGA."
+    "We use `weight_bit_width=6`, `bias_bit_width=6`, and `bit_width=6` throughout, matching the 6-bit precision from Part 2a.\n",
+    "Unlike regular PyTorch, every forward pass applies the quantization, so the model trains with the actual fixed-point precisions used on the FPGA. Additionally, we specify FixedPoint quantizers to enforce power of two scales, which are generally better handled by hls4ml.\n",
+    "Note, these particular quantizers have a maximum bit width of 8-bits, though narrower is allowed."
    ]
   },
   {
@@ -93,16 +94,22 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "from brevitas.quant import Int8BiasPerTensorFixedPointInternalScaling, Int8WeightPerTensorFixedPoint, Uint8ActPerTensorFixedPoint\n",
+    "\n",
     "class JetTaggerBrevitas(nn.Module):\n",
     "    def __init__(self):\n",
     "        super().__init__()\n",
-    "        self.fc1 = qnn.QuantLinear(16, 64, bias=False, weight_bit_width=6)\n",
-    "        self.relu1 = qnn.QuantReLU(bit_width=6)\n",
-    "        self.fc2 = qnn.QuantLinear(64, 32, bias=False, weight_bit_width=6)\n",
-    "        self.relu2 = qnn.QuantReLU(bit_width=6)\n",
-    "        self.fc3 = qnn.QuantLinear(32, 32, bias=False, weight_bit_width=6)\n",
-    "        self.relu3 = qnn.QuantReLU(bit_width=6)\n",
-    "        self.output = qnn.QuantLinear(32, 5, bias=False, weight_bit_width=6)\n",
+    "        self.fc1 = qnn.QuantLinear(16, 64, weight_quant=Int8WeightPerTensorFixedPoint, weight_bit_width=6,\n",
+    "                                   bias_quant=Int8BiasPerTensorFixedPointInternalScaling, bias_bit_width=6)\n",
+    "        self.relu1 = qnn.QuantReLU(act_quant=Uint8ActPerTensorFixedPoint, bit_width=6)\n",
+    "        self.fc2 = qnn.QuantLinear(64, 32, weight_quant=Int8WeightPerTensorFixedPoint, weight_bit_width=6,\n",
+    "                                   bias_quant=Int8BiasPerTensorFixedPointInternalScaling, bias_bit_width=6)\n",
+    "        self.relu2 = qnn.QuantReLU(act_quant=Uint8ActPerTensorFixedPoint, bit_width=6)\n",
+    "        self.fc3 = qnn.QuantLinear(32, 32, weight_quant=Int8WeightPerTensorFixedPoint, weight_bit_width=6,\n",
+    "                                   bias_quant=Int8BiasPerTensorFixedPointInternalScaling, bias_bit_width=6)\n",
+    "        self.relu3 = qnn.QuantReLU(act_quant=Uint8ActPerTensorFixedPoint, bit_width=6)\n",
+    "        self.output = qnn.QuantLinear(32, 5, weight_quant=Int8WeightPerTensorFixedPoint, weight_bit_width=6,\n",
+    "                                      bias_quant=Int8BiasPerTensorFixedPointInternalScaling, bias_bit_width=6)\n",
     "\n",
     "    def forward(self, x):\n",
     "        x = self.relu1(self.fc1(x))\n",
@@ -233,7 +240,7 @@
    "source": [
     "## Convert to hls4ml via QONNX\n",
     "\n",
-    "`config_from_onnx_model` reads the Quant nodes embedded in the QONNX graph and derives per-layer precision automatically. "
+    "`config_from_onnx_model` reads the Quant nodes embedded in the QONNX graph and derives per-layer precision automatically."
    ]
   },
   {
@@ -251,10 +258,6 @@
     "    backend='Vitis',\n",
     ")\n",
     "\n",
-    "for layer in config.get('LayerName', {}):\n",
-    "    if layer.startswith('Softmax'):\n",
-    "        config['LayerName'][layer]['Implementation'] = 'legacy'\n",
-    "\n",
     "print('-----------------------------------')\n",
     "plotting.print_dict(config)\n",
     "print('-----------------------------------')\n",
@@ -408,7 +411,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "hls4ml-tutorial",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
@@ -422,7 +425,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.16"
+   "version": "3.12.13"
   }
  },
  "nbformat": 4,