Feat (ONNX Export): move dynamic quantizers and add example for expor…

…t to ONNX
Xilinx · Feb 26, 2024 · 1c8d328 · 1c8d328
1 parent 6079b12
commit 1c8d328
Show file tree

Hide file tree

Showing 5 changed files with 224 additions and 49 deletions.
diff --git a/notebooks/ONNX_export_tutorial.ipynb b/notebooks/ONNX_export_tutorial.ipynb
diff --git a/...xamples/common/generative/quant_blocks.py → ...brevitas/quant/generative/quant_blocks.py b/...xamples/common/generative/quant_blocks.py → ...brevitas/quant/generative/quant_blocks.py
diff --git a/..._examples/common/generative/quantizers.py → src/brevitas/quant/generative/quantizers.py b/..._examples/common/generative/quantizers.py → src/brevitas/quant/generative/quantizers.py
diff --git a/src/brevitas_examples/common/generative/quantize.py b/src/brevitas_examples/common/generative/quantize.py
@@ -19,6 +19,14 @@
 from brevitas.quant.fixed_point import Int8WeightPerChannelFixedPointMSE
 from brevitas.quant.fixed_point import Int8WeightPerTensorFixedPoint
 from brevitas.quant.fixed_point import Int8WeightPerTensorFixedPointMSE
+from brevitas.quant.generative.quantizers import Fp8e4m3WeightSymmetricGroupQuant
+from brevitas.quant.generative.quantizers import Int8DynamicActPerGroupFloat
+from brevitas.quant.generative.quantizers import Int8DynamicActPerRowFloat
+from brevitas.quant.generative.quantizers import Int8DynamicActPerTensorFloat
+from brevitas.quant.generative.quantizers import IntWeightSymmetricGroupQuant
+from brevitas.quant.generative.quantizers import ShiftedUint8DynamicActPerRowFloat
+from brevitas.quant.generative.quantizers import ShiftedUint8DynamicActPerTensorFloat
+from brevitas.quant.generative.quantizers import ShiftedUintWeightAsymmetricGroupQuant
 from brevitas.quant.scaled_int import Int8ActPerTensorFloat
 from brevitas.quant.scaled_int import Int8ActPerTensorFloatMSE
 from brevitas.quant.scaled_int import Int8WeightPerChannelFloat
@@ -31,16 +39,9 @@
 from brevitas.quant.shifted_scaled_int import ShiftedUint8WeightPerChannelFloatMSE
 from brevitas.quant.shifted_scaled_int import ShiftedUint8WeightPerTensorFloat
 from brevitas.quant.shifted_scaled_int import ShiftedUint8WeightPerTensorFloatMSE
-from brevitas_examples.common.generative.nn import LoRACompatibleQuantConv2d
-from brevitas_examples.common.generative.nn import LoRACompatibleQuantLinear
-from brevitas_examples.common.generative.quantizers import Fp8e4m3WeightSymmetricGroupQuant
-from brevitas_examples.common.generative.quantizers import Int8DynamicActPerGroupFloat
-from brevitas_examples.common.generative.quantizers import Int8DynamicActPerRowFloat
-from brevitas_examples.common.generative.quantizers import Int8DynamicActPerTensorFloat
-from brevitas_examples.common.generative.quantizers import IntWeightSymmetricGroupQuant
-from brevitas_examples.common.generative.quantizers import ShiftedUint8DynamicActPerRowFloat
-from brevitas_examples.common.generative.quantizers import ShiftedUint8DynamicActPerTensorFloat
-from brevitas_examples.common.generative.quantizers import ShiftedUintWeightAsymmetricGroupQuant
+
+from .nn import LoRACompatibleQuantConv2d
+from .nn import LoRACompatibleQuantLinear
 
 WEIGHT_QUANT_MAP = {
  'int': {

diff --git a/tests/brevitas_ort/common.py b/tests/brevitas_ort/common.py
@@ -20,6 +20,7 @@
 from brevitas.quant.fixed_point import Int8ActPerTensorFixedPoint
 from brevitas.quant.fixed_point import Int8WeightPerChannelFixedPoint
 from brevitas.quant.fixed_point import Int8WeightPerTensorFixedPoint
+from brevitas.quant.generative.quantizers import ShiftedUint8DynamicActPerTensorFloat
 from brevitas.quant.scaled_int import Int8AccumulatorAwareWeightQuant
 from brevitas.quant.scaled_int import Int8ActPerTensorFloat
 from brevitas.quant.scaled_int import Int8WeightPerChannelFloat
@@ -28,7 +29,6 @@
 from brevitas.quant.shifted_scaled_int import ShiftedUint8WeightPerChannelFloat
 from brevitas.quant.shifted_scaled_int import ShiftedUint8WeightPerTensorFloat
 from brevitas.quant_tensor import QuantTensor
-from brevitas_examples.common.generative.quantizers import ShiftedUint8DynamicActPerTensorFloat
 
 SEED = 123456
 OUT_CH = 16