Feat: remove quant metadata from quantlayer

docs/tutorials/quant_tensor_quant_conv2d_overview.ipynb

@@ -129,10 +129,10 @@
  }
  ],
  "source": [
- "print(f'Is weight quant enabled: {default_quant_conv.is_weight_quant_enabled}')\n",
- "print(f'Is bias quant enabled: {default_quant_conv.is_bias_quant_enabled}')\n",
+ "print(f'Is weight quant enabled: {default_quant_conv.weight_quant.is_quant_enabled}')\n",
+ "print(f'Is bias quant enabled: {default_quant_conv.bias_quant.is_quant_enabled}')\n",
  "print(f'Is input quant enabled: {default_quant_conv.is_input_quant_enabled}')\n",
- "print(f'Is output quant enabled: {default_quant_conv.is_output_quant_enabled}')"
+ "print(f'Is output quant enabled: {default_quant_conv.output_quant.is_quant_enabled}')"
  ]
  },
  {

notebooks/03_anatomy_of_a_quantizer.ipynb

@@ -775,7 +775,7 @@
  "cell_type": "markdown",
  "metadata": {},
  "source": [
- "Indeed we can verify that `quant_weight_scale()` is equal to `weight.abs().max()`:"
+ "Indeed we can verify that `weight_quant.scale()` is equal to `weight.abs().max()`:"
  ]
  },
  {
@@ -792,7 +792,7 @@
  }
  ],
  "source": [
- "assert_with_message((param_from_max_quant_conv.quant_weight_scale() == param_from_max_quant_conv.weight.abs().max()).item())"
+ "assert_with_message((param_from_max_quant_conv.weight_quant.scale() == param_from_max_quant_conv.weight.abs().max()).item())"
  ]
  },
  {
@@ -1024,7 +1024,7 @@
  }
  ],
  "source": [
- "assert_with_message((quant_conv1.quant_weight_scale() == quant_conv2.quant_weight_scale()).item())"
+ "assert_with_message((quant_conv1.weight_quant.scale() == quant_conv2.weight_quant.scale()).item())"
  ]
  },
  {
@@ -1059,9 +1059,9 @@
  " return module.weight.abs().mean()\n",
  " \n",
  "quant_conv1 = QuantConv2d(3, 2, (3, 3), weight_quant=SharedParamFromMeanWeightQuantizer)\n",
- "old_quant_conv1_scale = quant_conv1.quant_weight_scale()\n",
+ "old_quant_conv1_scale = quant_conv1.weight_quant.scale()\n",
  "quant_conv2 = QuantConv2d(3, 2, (3, 3), weight_quant=quant_conv1.weight_quant)\n",
- "new_quant_conv1_scale = quant_conv1.quant_weight_scale()\n",
+ "new_quant_conv1_scale = quant_conv1.weight_quant.scale()\n",
  "\n",
  "assert_with_message(not (old_quant_conv1_scale == new_quant_conv1_scale).item())"
  ]
@@ -1080,7 +1080,7 @@
  }
  ],
  "source": [
- "assert_with_message((new_quant_conv1_scale == quant_conv2.quant_weight_scale()).item())"
+ "assert_with_message((new_quant_conv1_scale == quant_conv2.weight_quant.scale()).item())"
  ]
  },
  {
@@ -1134,7 +1134,7 @@
  "quant_conv_w_init = QuantConv2d(3, 2, (3, 3), weight_quant=ParamFromMaxWeightQuantizer)\n",
  "torch.nn.init.uniform_(quant_conv_w_init.weight)\n",
  "\n",
- "assert_with_message(not (quant_conv_w_init.weight.abs().max() == quant_conv_w_init.quant_weight_scale()).item())"
+ "assert_with_message(not (quant_conv_w_init.weight.abs().max() == quant_conv_w_init.weight_quant.scale()).item())"
  ]
  },
  {
@@ -1160,7 +1160,7 @@
  "source": [
  "quant_conv_w_init.weight_quant.init_tensor_quant()\n",
  "\n",
- "assert_with_message((quant_conv_w_init.weight.abs().max() == quant_conv_w_init.quant_weight_scale()).item())"
+ "assert_with_message((quant_conv_w_init.weight.abs().max() == quant_conv_w_init.weight_quant.scale()).item())"
  ]
  },
  {

src/brevitas/export/common/handler/base.py

@@ -11,6 +11,7 @@
 
 from brevitas.function.ops import max_int
 from brevitas.function.ops import min_int
+from brevitas.nn.quant_layer import QuantNonLinearActLayer
 
 __all__ = ['BaseHandler', 'BitWidthHandlerMixin', 'ZeroPointHandlerMixin']
 
@@ -130,21 +131,25 @@ def zero_point_with_dtype(cls, signed, bit_width, zero_point):
 
  @classmethod
  def quant_input_zero_point(cls, module):
- signed = module.is_quant_input_signed
- zero_point = module.quant_input_zero_point()
- bit_width = module.quant_input_bit_width()
+ signed = module.input_quant.signed()
+ zero_point = module.input_quant.zero_point()
+ bit_width = module.input_quant.bit_width()
  return cls.zero_point_with_dtype(signed, bit_width, zero_point)
 
  @classmethod
  def quant_weight_zero_point(cls, module):
- signed = module.is_quant_weight_signed
- zero_point = module.quant_weight_zero_point()
- bit_width = module.quant_weight_bit_width()
+ signed = module.weight_quant.is_signed
+ zero_point = module.quant_weight().zero_point
+ bit_width = module.weight_quant.bit_width()
  return cls.zero_point_with_dtype(signed, bit_width, zero_point)
 
  @classmethod
  def quant_output_zero_point(cls, module):
- signed = module.is_quant_output_signed
- zero_point = module.quant_output_zero_point()
- bit_width = module.quant_output_bit_width()
+ if isinstance(module, QuantNonLinearActLayer):
+ quant = module.act_quant
+ else:
+ quant = module.output_quant
+ signed = quant.signed()
+ zero_point = quant.zero_point()
+ bit_width = quant.bit_width()
  return cls.zero_point_with_dtype(signed, bit_width, zero_point)

src/brevitas/export/onnx/standard/qoperator/handler/act.py

@@ -20,14 +20,14 @@ class StdQOpONNXQuantNLALHandler(StdQOpONNXQuantLayerHandler, ABC):
 
  @classmethod
  def validate(cls, module: QuantNLAL):
- if cls.input_quant_supported and module.is_input_quant_enabled:
- assert not module.is_quant_input_narrow_range, "Narrow range quant not supported."
- elif not cls.input_quant_supported and module.is_input_quant_enabled:
+ if cls.input_quant_supported and module.input_quant.is_quant_enabled:
+ assert not module.input_quant.is_narrow_range, "Narrow range quant not supported."
+ elif not cls.input_quant_supported and module.input_quant.is_quant_enabled:
  raise RuntimeError("Input quant not supported.")
- if module.is_act_quant_enabled:
- assert not module.is_quant_act_narrow_range, "Narrow range quant not supported."
- input_bit_width = module.quant_input_bit_width()
- act_bit_width = module.quant_act_bit_width()
+ if module.act_quant.is_quant_enabled:
+ assert not module.act_quant.is_narrow_range, "Narrow range quant not supported."
+ input_bit_width = module.input_quant.bit_width()
+ act_bit_width = module.act_quant.bit_width()
  if input_bit_width is not None:
  cls.validate_8b_bit_width(input_bit_width, le_then=True)
  if act_bit_width is not None:

src/brevitas/export/onnx/standard/qoperator/handler/base.py

@@ -13,6 +13,7 @@
 from brevitas.export.onnx.standard.function import DequantizeLinearFn
 from brevitas.export.onnx.standard.function import IntClipFn
 from brevitas.export.onnx.standard.function import QuantizeLinearFn
+from brevitas.nn.quant_layer import QuantNonLinearActLayer
 
 
 class StdQOpONNXQuantLayerHandler(ONNXBaseHandler,
@@ -48,64 +49,68 @@ def torch_8b_dtype(cls, is_signed):
 
  @classmethod
  def quant_output_shape(cls, module):
- cached_out = module._cached_out
+ cached_out = module.output_quant._cached_act
  if cached_out is None:
  raise RuntimeError("Caching of outputs is required to export")
  return cached_out.shape
 
  @classmethod
  def output_quant_symbolic_kwargs(cls, module):
- if module.is_output_quant_enabled:
+ quant_proxy = module.act_quant if isinstance(
+ module, QuantNonLinearActLayer) else module.output_quant
+ if quant_proxy.is_quant_enabled:
  return {
- 'output_scale': module.quant_output_scale(),
+ 'output_scale': quant_proxy.scale(),
  'output_zero_point': cls.quant_output_zero_point(module),
- 'output_dtype': cls.torch_8b_dtype(module.is_quant_output_signed),
- 'output_axis': cls.quant_axis(module.quant_output_scale())}
+ 'output_dtype': cls.torch_8b_dtype(quant_proxy.signed()),
+ 'output_axis': cls.quant_axis(quant_proxy.scale())}
  else:
  return None
 
  @classmethod
  def output_clip_symbolic_kwargs(cls, module):
- if module.is_output_quant_enabled:
- narrow = module.is_quant_output_narrow_range
- signed = module.is_quant_output_signed
- bit_width = module.quant_output_bit_width()
+ quant_proxy = module.act_quant if isinstance(
+ module, QuantNonLinearActLayer) else module.output_quant
+ if quant_proxy.is_quant_enabled:
+ narrow = quant_proxy.is_narrow_range
+ signed = quant_proxy.signed()
+ bit_width = quant_proxy.bit_width()
  return cls.int_clip_symbolic_kwargs(narrow, signed, bit_width)
  else:
  return None
 
  @classmethod
  def input_clip_symbolic_kwargs(cls, module):
- if module.is_input_quant_enabled:
- narrow = module.is_quant_input_narrow_range
- signed = module.is_quant_input_signed
- bit_width = module.quant_input_bit_width()
+ if module.input_quant.is_quant_enabled:
+ narrow = module.input_quant.is_narrow_range
+ signed = module.input_quant.signed()
+ bit_width = module.input_quant.bit_width()
  return cls.int_clip_symbolic_kwargs(narrow, signed, bit_width)
  else:
  return None
 
  @classmethod
  def output_dequant_symbolic_kwargs(cls, module):
  return {
- 'input_scale': module.quant_output_scale(),
+ 'input_scale': module.output_quant.scale(),
  'input_zero_point': cls.quant_output_zero_point(module),
- 'input_axis': cls.quant_axis(module.quant_output_scale())}
+ 'input_axis': cls.quant_axis(module.output_quant.scale())}
 
  @classmethod
  def input_quant_symbolic_kwargs(cls, module):
- if module.is_input_quant_enabled:
+ if module.input_quant.is_quant_enabled:
  return {
- 'output_scale': module.quant_input_scale(),
+ 'output_scale': module.input_quant.scale(),
  'output_zero_point': cls.quant_input_zero_point(module),
- 'output_dtype': cls.torch_8b_dtype(module.is_quant_input_signed),
- 'output_axis': cls.quant_axis(module.quant_input_scale())}
+ 'output_dtype': cls.torch_8b_dtype(module.input_quant.signed()),
+ 'output_axis': cls.quant_axis(module.input_quant.scale())}
  else:
  return None
 
  @classmethod
  def input_dequant_symbolic_kwargs(cls, module):
- if module._cached_inp is not None:
- return cls.dequant_symbolic_kwargs_from_cached_io(module._cached_inp)
+ if module.input_quant._cached_act is not None:
+ return cls.dequant_symbolic_kwargs_from_cached_io(module.input_quant._cached_act)
  else:
  return None
 

src/brevitas/export/onnx/standard/qoperator/handler/parameter.py

@@ -26,7 +26,7 @@ class StdQOpONNXQuantWBIOLHandler(StdQOpONNXQuantLayerHandler, ABC):
  @staticmethod
  def int_weight(module: QuantWBIOL):
  int_weight = module.int_weight(float_datatype=False).detach()
- if module.is_quant_weight_signed:
+ if module.weight_quant.is_signed:
  return int_weight.type(torch.int8)
  else:
  return int_weight.type(torch.uint8)
@@ -41,20 +41,20 @@ def int_bias(module: QuantWBIOL):
 
  @classmethod
  def validate(cls, module: QuantWBIOL, requires_quant_bias=True):
- assert module.is_weight_quant_enabled, 'Weight quant required'
- assert module.is_output_quant_enabled, 'Output quant required'
+ assert module.weight_quant.is_quant_enabled, 'Weight quant required'
+ assert module.output_quant.is_quant_enabled, 'Output quant required'
  # Handling narrow_range is across the network is difficult do to the fact that
  # it's not part of QuantTensor, and so it can't be cached
- assert not module.is_quant_output_narrow_range, 'Narrow output quant not supported'
- if module.is_input_quant_enabled:
- assert not module.is_quant_input_narrow_range, 'Narrow output quant not supported'
- cls.validate_8b_bit_width(module.quant_weight_bit_width(), le_then=True)
- cls.validate_8b_bit_width(module.quant_input_bit_width(), le_then=True)
- cls.validate_8b_bit_width(module.quant_output_bit_width(), le_then=True)
+ assert not module.output_quant.is_narrow_range, 'Narrow output quant not supported'
+ if module.input_quant.is_quant_enabled:
+ assert not module.input_quant.is_narrow_range, 'Narrow output quant not supported'
+ cls.validate_8b_bit_width(module.weight_quant.bit_width(), le_then=True)
+ cls.validate_8b_bit_width(module.input_quant.bit_width(), le_then=True)
+ cls.validate_8b_bit_width(module.output_quant.bit_width(), le_then=True)
  if module.bias is not None and requires_quant_bias:
- assert module.is_bias_quant_enabled
- assert module.is_quant_bias_signed
- cls.validate_32b_bit_width(module.quant_bias_bit_width(), le_then=True)
+ assert module.bias_quant.is_quant_enabled
+ assert module.bias_quant.is_signed
+ cls.validate_32b_bit_width(module.bias_quant.bit_width(), le_then=True)
 
  def prepare_for_export(self, module: Union[QuantConv1d, QuantConv2d]):
  self.validate(module)
@@ -108,14 +108,14 @@ class StdQOpONNXQuantConvNdHandler(StdQOpONNXQuantWBIOLHandler, ABC):
 
  def op_symbolic_kwargs(self, module: Union[QuantConv1d, QuantConv2d]):
  conv_symbolic_kwargs = {
- 'input_scale': module.quant_input_scale(),
+ 'input_scale': module.input_quant.scale(),
  'input_zero_point': self.quant_input_zero_point(module),
  'int_weight': self.int_weight(module),
- 'weight_scale': to_0dim_if_scalar(module.quant_weight_scale().flatten()),
+ 'weight_scale': to_0dim_if_scalar(module.quant_weight().scale.flatten()),
  'weight_zero_point': to_0dim_if_scalar(self.quant_weight_zero_point(module).flatten()),
- 'output_scale': module.quant_output_scale(),
+ 'output_scale': module.output_quant.scale(),
  'output_zero_point': self.quant_output_zero_point(module),
- 'output_dtype': self.torch_8b_dtype(module.is_quant_output_signed),
+ 'output_dtype': self.torch_8b_dtype(module.output_quant.signed()),
  'int_bias': self.int_bias(module),
  'out_shape': self.quant_output_shape(module),
  'kernel_size': list(module.kernel_size),
@@ -145,14 +145,14 @@ class StdQOpONNXQuantLinearHandler(StdQOpONNXQuantWBIOLHandler):
  # Convert linear to conv1d to handle bias
  def op_symbolic_kwargs(self, module: QuantLinear):
  conv_symbolic_kwargs = {
- 'input_scale': module.quant_input_scale(),
+ 'input_scale': module.input_quant.scale(),
  'input_zero_point': self.quant_input_zero_point(module),
  'int_weight': self.int_weight(module).view(module.out_features, module.in_features, 1),
- 'weight_scale': to_0dim_if_scalar(module.quant_weight_scale().flatten()),
+ 'weight_scale': to_0dim_if_scalar(module.quant_weight().scale.flatten()),
  'weight_zero_point': to_0dim_if_scalar(self.quant_weight_zero_point(module).flatten()),
- 'output_scale': module.quant_output_scale(),
+ 'output_scale': module.output_quant.scale(),
  'output_zero_point': self.quant_output_zero_point(module),
- 'output_dtype': self.torch_8b_dtype(module.is_quant_output_signed),
+ 'output_dtype': self.torch_8b_dtype(module.output_quant.signed()),
  'int_bias': self.int_bias(module),
  'out_shape': self.quant_output_shape(module) + (1,),
  'kernel_size': [1],

src/brevitas/export/torch/qoperator/handler/act.py

@@ -23,13 +23,13 @@ def explicit_output_dtype(cls) -> bool:
 
  @classmethod
  def validate(cls, module: QuantNLAL):
- assert not module.is_input_quant_enabled, 'Input quantization not supported'
- cls.validate_8b_bit_width(module.quant_act_bit_width(), le_then=False)
+ assert not module.input_quant.is_quant_enabled, 'Input quantization not supported'
+ cls.validate_8b_bit_width(module.act_quant.bit_width(), le_then=False)
 
  def prepare_for_export(self, module: QuantNLAL):
  self.validate(module)
  self.qf_impl, self.qf_kwargs = self.prepare_qf(module)
- if module.is_act_quant_enabled:
+ if module.act_quant.is_quant_enabled:
  self.output_quant_impl, self.output_quant_kwargs = self.prepare_output_quant(module)
  self.return_quant_tensor = module.return_quant_tensor
 

src/brevitas/export/torch/qoperator/handler/base.py

@@ -10,6 +10,7 @@
 from brevitas.export.common.handler.base import BaseHandler
 from brevitas.export.common.handler.base import BitWidthHandlerMixin
 from brevitas.export.common.handler.base import ZeroPointHandlerMixin
+from brevitas.nn.quant_layer import QuantNonLinearActLayer
 
 SCALAR_SHAPE = ()
 
@@ -55,17 +56,21 @@ def gen_quant_impl_kwargs(
 
  @classmethod
  def prepare_input_quant(cls, module):
- scale = module.quant_input_scale()
+ scale = module.input_quant.scale()
  zero_point = cls.quant_input_zero_point(module)
- signed = module.is_quant_input_signed
+ signed = module.input_quant.signed()
  quant_impl, quant_kwargs = cls.gen_quant_impl_kwargs(scale, zero_point, signed)
  return quant_impl, quant_kwargs
 
  @classmethod
  def prepare_output_quant(cls, module):
- scale = module.quant_output_scale()
+ if isinstance(module, QuantNonLinearActLayer):
+ quant = module.act_quant
+ else:
+ quant = module.output_quant
+ scale = quant.scale()
  zero_point = cls.quant_output_zero_point(module)
- signed = module.is_quant_output_signed
+ signed = quant.signed()
  incl_dtype = cls.explicit_output_dtype()
  quant_impl, quant_kwargs = cls.gen_quant_impl_kwargs(scale, zero_point, signed, incl_dtype)
  return quant_impl, quant_kwargs