diff --git a/src/pquant/__init__.py b/src/pquant/__init__.py index 6c27f9d..7eaa9ed 100644 --- a/src/pquant/__init__.py +++ b/src/pquant/__init__.py @@ -36,6 +36,7 @@ load_torch_hgq_model, post_training_prune, ) + from .core.torch.onnx import convert_to_onnx from .core.torch.tracing import check_quantization, print_quantization_check from .core.torch.train import train_model @@ -68,6 +69,7 @@ _forwards.append("check_quantization") _forwards.append("print_quantization_check") _forwards.append("PQConfig") + _forwards.append("convert_to_onnx") __all__ = _forwards else: @@ -93,6 +95,7 @@ get_model_losses, post_training_prune, ) + from .core.keras.onnx import convert_to_onnx from .core.keras.train import train_model _forwards = ["activations", "layers", "quantizer"] @@ -119,4 +122,5 @@ _forwards.append("load_from_file") _forwards.append("load_from_dictionary") _forwards.append("PQConfig") + _forwards.append("convert_to_onnx") __all__ = _forwards diff --git a/src/pquant/core/keras/convert_to_onnx.py b/src/pquant/core/keras/convert_to_onnx.py deleted file mode 100644 index 35263c2..0000000 --- a/src/pquant/core/keras/convert_to_onnx.py +++ /dev/null @@ -1,1478 +0,0 @@ -""" -Convert a PQuant Keras functional model to ONNX or QONNX format. - -Pass ``use_qonnx=True`` to emit QONNX ``Quant`` custom nodes (requires the -qonnx runtime). Pass ``use_qonnx=False`` (default) to emit standard -``Clip + QuantizeLinear + DequantizeLinear`` nodes runnable with plain -onnxruntime. - -Fixed-point (k, i, f) mapping ------------------------------- -QONNX: - scale = 2^(-f) - zero_point = 0 - bit_width = k + i + f - signed = int(k) - -Standard ONNX (QDQ): - scale = 2^(-f) - zero_point = 0 (int8 signed, uint8 unsigned) - clip range = [-2^i, 2^i - 2^(-f)] signed - = [0, 2^i - 2^(-f)] unsigned - -Keras weight layout (kernel always stored as HWIO regardless of data_format): - Conv2D kernel: [kH, kW, in/g, out] → [out, in/g, kH, kW] for ONNX - Conv1D kernel: [kL, in/g, out] → [out, in/g, kL] for ONNX - DepthwiseConv2D kernel: [kH, kW, in, dm] → [in*dm, 1, kH, kW] for ONNX - Dense kernel: [in, out] → stored as [out, in] for Gemm (transB=1) - -Data format: - channels_first Data flows as NCHW; Conv/Pool ONNX ops work naturally. - channels_last Transpose(NHWC→NCHW) is inserted before each Conv/Pool/BN - node and Transpose(NCHW→NHWC) is inserted after. The - logical data format in the ONNX graph therefore stays NHWC - at every inter-layer edge; only the PQ ops run internally in - NCHW. ONNX-aware optimisers (e.g. onnxsim) can fold the - redundant back-to-back transposes away. -""" - -import functools -import logging - -import keras -import numpy as np -import onnx -import onnx.helper as oh -import onnx.numpy_helper as onh -from keras import ops -from onnx import TensorProto - -from pquant.core.keras.activations import PQActivation -from pquant.core.keras.layers import ( - PQBatchNormalization, - PQConv1d, - PQConv2d, - PQDense, - PQDepthwiseConv2d, - PQMultiheadAttention, -) - -# --------------------------------------------------------------------------- -# QONNX Quant node -# --------------------------------------------------------------------------- - -ROUND_MODE_MAP = { - "TRN": "FLOOR", - "RND": "ROUND", - "RND_CONV": "ROUND", - "TRN_ZERO": "TRUNCATE", - "RND_ZERO": "ROUND", - "RND_MIN_INF": "FLOOR", - "RND_INF": "ROUND", -} - - -def _quant_node(name_prefix, input_name, rounding_mode, k, i, f, initializers, overflow_mode="SAT"): - """Build a QONNX Quant node. k/i/f are numpy arrays. Returns ([node], output_name).""" - k_val = int(float(np.array(k).ravel()[0])) - f_arr = np.array(f, dtype=np.float32) - i_arr = np.array(i, dtype=np.float32) - if f_arr.size > 1: - i_arr = i_arr.ravel().max() - f_arr = f_arr.ravel().min() - i_val = float(i_arr) - f_val = float(f_arr) - scale = float(2.0 ** (-f_val)) - bit_width = float(k_val + i_val + f_val) - qonnx_rnd = ROUND_MODE_MAP.get(rounding_mode, "ROUND") - # SAT_SYM excludes the most-negative value → QONNX narrow=1 - narrow = 1 if (k_val == 1 and overflow_mode == "SAT_SYM") else 0 - - scale_name = f"{name_prefix}_scale" - zp_name = f"{name_prefix}_zero_point" - bw_name = f"{name_prefix}_bit_width" - out_name = f"{name_prefix}_quantized" - - initializers.append(onh.from_array(np.array(scale, dtype=np.float32), name=scale_name)) - initializers.append(onh.from_array(np.array(0.0, dtype=np.float32), name=zp_name)) - initializers.append(onh.from_array(np.array(bit_width, dtype=np.float32), name=bw_name)) - - node = oh.make_node( - op_type="Quant", - inputs=[input_name, scale_name, zp_name, bw_name], - outputs=[out_name], - domain="qonnx.custom_op.general", - signed=k_val, - narrow=narrow, - rounding_mode=qonnx_rnd, - ) - return [node], out_name - - -# --------------------------------------------------------------------------- -# Standard ONNX QDQ triple -# --------------------------------------------------------------------------- - - -def _qdq_node( - name_prefix, input_name, rounding_mode, k, i, f, initializers, overflow_mode="SAT", include_clip=True -): # noqa: ARG001 - """Build QuantizeLinear+DequantizeLinear nodes, optionally preceded by a Clip. - - Returns ([nodes], output_name). Set include_clip=False to skip the Clip node - (safe when values are guaranteed to be in-range at inference time, since - QuantizeLinear saturates naturally). - """ - k_val = int(float(np.array(k).ravel()[0])) - i_val = float(np.array(i, dtype=np.float32).ravel()[0]) - f_val = float(np.array(f, dtype=np.float32).ravel()[0]) - scale = float(2.0 ** (-f_val)) - signed = k_val == 1 - - clip_max = float(2.0**i_val - 2.0 ** (-f_val)) - if not signed: - clip_min = 0.0 - elif overflow_mode == "SAT_SYM": - clip_min = -clip_max # symmetric: -(2^i - 2^(-f)) - else: - clip_min = float(-(2.0**i_val)) # SAT: -2^i - zp_val = np.int8(0) if signed else np.uint8(0) - - scale_name = f"{name_prefix}_scale" - zp_name = f"{name_prefix}_zero_point" - quantized_name = f"{name_prefix}_quantized" - out_name = f"{name_prefix}_dequantized" - - initializers += [ - onh.from_array(np.array(scale, dtype=np.float32), name=scale_name), - onh.from_array(np.array(zp_val), name=zp_name), - ] - - if include_clip: - clip_min_name = f"{name_prefix}_clip_min" - clip_max_name = f"{name_prefix}_clip_max" - clipped_name = f"{name_prefix}_clipped" - initializers += [ - onh.from_array(np.array(clip_min, dtype=np.float32), name=clip_min_name), - onh.from_array(np.array(clip_max, dtype=np.float32), name=clip_max_name), - ] - nodes = [ - oh.make_node("Clip", inputs=[input_name, clip_min_name, clip_max_name], outputs=[clipped_name]), - oh.make_node("QuantizeLinear", inputs=[clipped_name, scale_name, zp_name], outputs=[quantized_name]), - ] - else: - nodes = [ - oh.make_node("QuantizeLinear", inputs=[input_name, scale_name, zp_name], outputs=[quantized_name]), - ] - - nodes.append(oh.make_node("DequantizeLinear", inputs=[quantized_name, scale_name, zp_name], outputs=[out_name])) - return nodes, out_name - - -# --------------------------------------------------------------------------- -# integer weight storage helper -# --------------------------------------------------------------------------- - - -def _int_weight_node(name_prefix, weight_np, k, i, f, initializers): # noqa: ARG001 (i unused) - """ - Store a weight tensor as int8/uint8 + DequantizeLinear. - - weight_np must already be in ONNX layout (transposed from Keras) and on the - fixed-point grid after apply_final_compression(). - - k/i/f are numpy arrays (may be per-tensor scalar or per-channel 1-D after - caller has already squeezed/reshaped appropriately). - - Granularity: - - per-tensor (f is scalar): single scale. - - per-channel (f is 1-D of length out_channels): axis=0 on weight tensor. - - per-weight (fully per-element): falls back to float32 storage. - - Returns ([node], output_name). - """ - k_np = np.array(k, dtype=np.float32) - f_np = np.array(f, dtype=np.float32) - k_val = int(float(k_np.ravel()[0])) - dtype = np.int8 if k_val == 1 else np.uint8 - out_channels = weight_np.shape[0] - out_name = f"{name_prefix}_dequantized" - - if f_np.size == 1: - # per-tensor - scale_np = np.array(float(2.0 ** (-float(f_np.ravel()[0]))), dtype=np.float32) - int_w = np.round(weight_np / float(scale_np)).astype(dtype) - per_ch = False - else: - f_1d = f_np.ravel() - if f_1d.size == out_channels: - # per-channel: one f value per output channel - scale_1d = (2.0 ** (-f_1d)).astype(np.float32) - bcast = scale_1d.reshape((out_channels,) + (1,) * (weight_np.ndim - 1)) - int_w = np.round(weight_np / bcast).astype(dtype) - scale_np = scale_1d - per_ch = True - else: - # per-weight: ONNX cannot represent; fall back to float32 - float_name = f"{name_prefix}_float" - initializers.append(onh.from_array(weight_np, name=float_name)) - return [], float_name - - int_name = f"{name_prefix}_int" - scale_name = f"{name_prefix}_dq_scale" - zp_name = f"{name_prefix}_dq_zp" - - zp_np = np.zeros(out_channels if per_ch else 1, dtype=dtype) - initializers += [ - onh.from_array(int_w, name=int_name), - onh.from_array(scale_np, name=scale_name), - onh.from_array(zp_np if per_ch else np.array(dtype(0)), name=zp_name), - ] - node_kwargs = {"axis": 0} if per_ch else {} - node = oh.make_node("DequantizeLinear", inputs=[int_name, scale_name, zp_name], outputs=[out_name], **node_kwargs) - return [node], out_name - - -# --------------------------------------------------------------------------- -# helpers -# --------------------------------------------------------------------------- - - -def _np(tensor): - """Convert a Keras tensor / variable / scalar to a float32 numpy array.""" - return np.array(tensor, dtype=np.float32) - - -def _bn_transpose_info(layer): - """ - Return (need_transpose, perm_fwd, perm_bwd) for a BatchNormalization layer. - - ONNX BN (opset < 14) always normalises on axis 1 (NCHW). If the Keras - layer uses axis=-1 (channels_last), we must insert Transpose nodes around - the BN op. We infer ndim from the layer's stored input_shape. - """ - axis = getattr(layer, "axis", 1) - stored = getattr(layer, "input_shape", None) - ndim = len(stored) if stored is not None else 4 - eff_axis = axis if axis >= 0 else (ndim + axis) - - if eff_axis == 1 or ndim <= 2: - # channels already at position 1, or 2-D input — no transpose needed - return False, None, None - - if ndim == 4 and eff_axis == 3: - return True, [0, 3, 1, 2], [0, 2, 3, 1] - - if ndim == 3 and eff_axis == 2: - return True, [0, 2, 1], [0, 2, 1] - - # Fallback: general permutation that moves eff_axis to position 1 - perm_fwd = [0, eff_axis] + [i for i in range(1, ndim) if i != eff_axis] - # Inverse permutation - perm_bwd = [0] * ndim - for i, p in enumerate(perm_fwd): - perm_bwd[p] = i - return True, perm_fwd, perm_bwd - - -def _maybe_quant_input(layer, prefix, current, nodes, initializers, quant_fn): - if ( - getattr(layer, "input_quantizer", None) is not None - and getattr(layer, "quantize_input", True) - and getattr(layer, "enable_quantization", True) - ): - q = layer.input_quantizer - k, i, f = q.get_quantization_bits() - new_nodes, current = quant_fn( - f"{prefix}_in", - current, - q.round_mode, - _np(k), - _np(i), - _np(f), - initializers, - overflow_mode=getattr(q, "overflow", "SAT"), - ) - nodes.extend(new_nodes) - return current - - -def _maybe_quant_output(layer, prefix, current, nodes, initializers, quant_fn): - if ( - getattr(layer, "output_quantizer", None) is not None - and getattr(layer, "quantize_output", False) - and getattr(layer, "enable_quantization", True) - ): - q = layer.output_quantizer - k, i, f = q.get_quantization_bits() - new_nodes, current = quant_fn( - f"{prefix}_out", - current, - q.round_mode, - _np(k), - _np(i), - _np(f), - initializers, - overflow_mode=getattr(q, "overflow", "SAT"), - ) - nodes.extend(new_nodes) - return current - - -def _add_transpose(name, input_name, perm, nodes): - """Emit a Transpose node and return the output name.""" - out = f"{name}_transpose_{''.join(str(p) for p in perm)}" - nodes.append(oh.make_node("Transpose", inputs=[input_name], outputs=[out], perm=list(perm))) - return out - - -def _channels_last(layer): - return getattr(layer, "data_format", "channels_first") == "channels_last" - - -def _weight_f_for_onnx(f_np, out_channels): - """Squeeze/ravel a Keras per-channel f array to shape (out_channels,) for ONNX.""" - f_flat = f_np.ravel() - if f_flat.size == 1: - return f_flat # scalar, return as-is - if f_flat.size == out_channels: - return f_flat - # Per-element or mismatched: take the minimum to avoid overflow - return np.array([f_flat.min()], dtype=np.float32) - - -# --------------------------------------------------------------------------- -# per-layer graph builders -# --------------------------------------------------------------------------- - - -def _add_dense(layer, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights): - """Dense / PQDense. Keras kernel [in, out] stored as [out, in]; Gemm uses transB=1. - - Storing the weight pre-transposed means axis=0 is always the output dimension, - which is required for per-channel DequantizeLinear and avoids a runtime Transpose. - """ - current = _maybe_quant_input(layer, prefix, current, nodes, initializers, quant_fn) - - # Transpose kernel to [out, in]; Gemm will use transB=1 so Y = X @ W^T = X @ kernel. - kernel_np = _np(layer._kernel).T # [out, in] - out_units = kernel_np.shape[0] - - if use_qonnx: - k_w, i_w, f_w = layer.weight_quantizer.get_quantization_bits() - wfp_name = f"{prefix}_weight_fp" - initializers.append(onh.from_array(kernel_np, name=wfp_name)) - w_nodes, q_weight = _quant_node( - f"{prefix}_weight", - wfp_name, - layer.weight_quantizer.round_mode, - _np(k_w), - _np(i_w), - _np(f_w), - initializers, - overflow_mode=getattr(layer.weight_quantizer, "overflow", "SAT"), - ) - nodes.extend(w_nodes) - elif store_integer_weights: - k_w, i_w, f_w = layer.weight_quantizer.get_quantization_bits() - f_np_w = _np(f_w) - f_for_onnx = _weight_f_for_onnx(f_np_w, out_units) - k_for_onnx = _weight_f_for_onnx(_np(k_w), out_units) - i_for_onnx = _weight_f_for_onnx(_np(i_w), out_units) - w_nodes, q_weight = _int_weight_node(f"{prefix}_weight", kernel_np, k_for_onnx, i_for_onnx, f_for_onnx, initializers) - nodes.extend(w_nodes) - else: - q_weight = f"{prefix}_weight" - initializers.append(onh.from_array(kernel_np, name=q_weight)) - - gemm_inputs = [current, q_weight] - - if layer._bias is not None: - bias_np = _np(layer._bias) - if use_qonnx: - k_b, i_b, f_b = layer.bias_quantizer.get_quantization_bits() - bfp_name = f"{prefix}_bias_fp" - initializers.append(onh.from_array(bias_np, name=bfp_name)) - b_nodes, q_bias = _quant_node( - f"{prefix}_bias", - bfp_name, - layer.bias_quantizer.round_mode, - _np(k_b), - _np(i_b), - _np(f_b), - initializers, - overflow_mode=getattr(layer.bias_quantizer, "overflow", "SAT"), - ) - nodes.extend(b_nodes) - elif store_integer_weights: - k_b, i_b, f_b = layer.bias_quantizer.get_quantization_bits() - b_nodes, q_bias = _int_weight_node(f"{prefix}_bias", bias_np, _np(k_b), _np(i_b), _np(f_b), initializers) - nodes.extend(b_nodes) - else: - q_bias = f"{prefix}_bias" - initializers.append(onh.from_array(bias_np, name=q_bias)) - gemm_inputs.append(q_bias) - - gemm_out = f"{prefix}_gemm" - nodes.append(oh.make_node("Gemm", inputs=gemm_inputs, outputs=[gemm_out], transB=1)) - current = gemm_out - - current = _maybe_quant_output(layer, prefix, current, nodes, initializers, quant_fn) - return current - - -def _add_conv(layer, prefix, current, nodes, initializers, ndim, quant_fn, use_qonnx, store_integer_weights): - """PQConv2d / PQConv1d. Keras kernel: [*kernel, in/g, out] → ONNX [out, in/g, *kernel].""" - cl = _channels_last(layer) - - if cl: - perm_to_nchw = [0, 3, 1, 2] if ndim == 2 else [0, 2, 1] - perm_to_nhwx = [0, 2, 3, 1] if ndim == 2 else [0, 2, 1] - current = _add_transpose(f"{prefix}_pre", current, perm_to_nchw, nodes) - - current = _maybe_quant_input(layer, prefix, current, nodes, initializers, quant_fn) - - kernel_np = _np(layer._kernel) - # Transpose kernel from Keras HWIO to ONNX OIHW - if ndim == 2: - kernel_onnx = np.transpose(kernel_np, (3, 2, 0, 1)) # [kH,kW,in,out] → [out,in,kH,kW] - else: - kernel_onnx = np.transpose(kernel_np, (2, 1, 0)) # [kL,in,out] → [out,in,kL] - - out_channels = kernel_onnx.shape[0] - - if use_qonnx: - k_w, i_w, f_w = layer.weight_quantizer.get_quantization_bits() - wfp_name = f"{prefix}_weight_fp" - initializers.append(onh.from_array(kernel_onnx, name=wfp_name)) - w_nodes, q_weight = _quant_node( - f"{prefix}_weight", - wfp_name, - layer.weight_quantizer.round_mode, - _np(k_w), - _np(i_w), - _np(f_w), - initializers, - overflow_mode=getattr(layer.weight_quantizer, "overflow", "SAT"), - ) - nodes.extend(w_nodes) - elif store_integer_weights: - k_w, i_w, f_w = layer.weight_quantizer.get_quantization_bits() - f_for_onnx = _weight_f_for_onnx(_np(f_w), out_channels) - k_for_onnx = _weight_f_for_onnx(_np(k_w), out_channels) - i_for_onnx = _weight_f_for_onnx(_np(i_w), out_channels) - w_nodes, q_weight = _int_weight_node( - f"{prefix}_weight", kernel_onnx, k_for_onnx, i_for_onnx, f_for_onnx, initializers - ) - nodes.extend(w_nodes) - else: - q_weight = f"{prefix}_weight" - initializers.append(onh.from_array(kernel_onnx, name=q_weight)) - - conv_inputs = [current, q_weight] - - if layer._bias is not None: - bias_np = _np(layer._bias) - if use_qonnx: - k_b, i_b, f_b = layer.bias_quantizer.get_quantization_bits() - bfp_name = f"{prefix}_bias_fp" - initializers.append(onh.from_array(bias_np, name=bfp_name)) - b_nodes, q_bias = _quant_node( - f"{prefix}_bias", - bfp_name, - layer.bias_quantizer.round_mode, - _np(k_b), - _np(i_b), - _np(f_b), - initializers, - overflow_mode=getattr(layer.bias_quantizer, "overflow", "SAT"), - ) - nodes.extend(b_nodes) - elif store_integer_weights: - k_b, i_b, f_b = layer.bias_quantizer.get_quantization_bits() - b_nodes, q_bias = _int_weight_node(f"{prefix}_bias", bias_np, _np(k_b), _np(i_b), _np(f_b), initializers) - nodes.extend(b_nodes) - else: - q_bias = f"{prefix}_bias" - initializers.append(onh.from_array(bias_np, name=q_bias)) - conv_inputs.append(q_bias) - - # Padding - padding = layer.padding - if isinstance(padding, str): - auto_pad = "SAME_UPPER" if padding == "same" else "VALID" - pads = None - else: - p = list(padding) if hasattr(padding, "__iter__") else [padding] * ndim - pads = p + p # ONNX format: [begin_0, begin_1, ..., end_0, end_1, ...] - auto_pad = "NOTSET" - - to_list = lambda v, n: list(v) if hasattr(v, "__iter__") else [v] * n # noqa: E731 - conv_attrs = dict( - kernel_shape=to_list(layer.kernel_size, ndim), - strides=to_list(layer.strides, ndim), - dilations=to_list(layer.dilation_rate, ndim), - group=getattr(layer, "groups", 1), - auto_pad=auto_pad, - ) - if pads is not None: - conv_attrs["pads"] = pads - - conv_out = f"{prefix}_conv" - nodes.append(oh.make_node("Conv", inputs=conv_inputs, outputs=[conv_out], **conv_attrs)) - current = conv_out - - current = _maybe_quant_output(layer, prefix, current, nodes, initializers, quant_fn) - - if cl: - current = _add_transpose(f"{prefix}_post", current, perm_to_nhwx, nodes) - return current - - -def _add_depthwise_conv(layer, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights): - """PQDepthwiseConv2d. - - Keras kernel: [kH, kW, in, depth_mult] - ONNX Conv with groups=in: weight [in*depth_mult, 1, kH, kW] - """ - cl = _channels_last(layer) - - if cl: - current = _add_transpose(f"{prefix}_pre", current, [0, 3, 1, 2], nodes) - - current = _maybe_quant_input(layer, prefix, current, nodes, initializers, quant_fn) - - kernel_np = _np(layer._kernel) # [kH, kW, in, depth_mult] - in_ch, depth_mult = kernel_np.shape[2], kernel_np.shape[3] - # Rearrange to [in*depth_mult, 1, kH, kW] - kernel_onnx = np.transpose(kernel_np, (2, 3, 0, 1)).reshape(in_ch * depth_mult, 1, *kernel_np.shape[:2]) - - out_channels = kernel_onnx.shape[0] - - if use_qonnx: - k_w, i_w, f_w = layer.weight_quantizer.get_quantization_bits() - wfp_name = f"{prefix}_weight_fp" - initializers.append(onh.from_array(kernel_onnx, name=wfp_name)) - w_nodes, q_weight = _quant_node( - f"{prefix}_weight", - wfp_name, - layer.weight_quantizer.round_mode, - _np(k_w), - _np(i_w), - _np(f_w), - initializers, - overflow_mode=getattr(layer.weight_quantizer, "overflow", "SAT"), - ) - nodes.extend(w_nodes) - elif store_integer_weights: - k_w, i_w, f_w = layer.weight_quantizer.get_quantization_bits() - f_for_onnx = _weight_f_for_onnx(_np(f_w), out_channels) - k_for_onnx = _weight_f_for_onnx(_np(k_w), out_channels) - i_for_onnx = _weight_f_for_onnx(_np(i_w), out_channels) - w_nodes, q_weight = _int_weight_node( - f"{prefix}_weight", kernel_onnx, k_for_onnx, i_for_onnx, f_for_onnx, initializers - ) - nodes.extend(w_nodes) - else: - q_weight = f"{prefix}_weight" - initializers.append(onh.from_array(kernel_onnx, name=q_weight)) - - conv_inputs = [current, q_weight] - - if layer._bias is not None: - bias_np = _np(layer._bias) - if use_qonnx: - k_b, i_b, f_b = layer.bias_quantizer.get_quantization_bits() - bfp_name = f"{prefix}_bias_fp" - initializers.append(onh.from_array(bias_np, name=bfp_name)) - b_nodes, q_bias = _quant_node( - f"{prefix}_bias", - bfp_name, - layer.bias_quantizer.round_mode, - _np(k_b), - _np(i_b), - _np(f_b), - initializers, - overflow_mode=getattr(layer.bias_quantizer, "overflow", "SAT"), - ) - nodes.extend(b_nodes) - elif store_integer_weights: - k_b, i_b, f_b = layer.bias_quantizer.get_quantization_bits() - b_nodes, q_bias = _int_weight_node(f"{prefix}_bias", bias_np, _np(k_b), _np(i_b), _np(f_b), initializers) - nodes.extend(b_nodes) - else: - q_bias = f"{prefix}_bias" - initializers.append(onh.from_array(bias_np, name=q_bias)) - conv_inputs.append(q_bias) - - padding = layer.padding - if isinstance(padding, str): - auto_pad = "SAME_UPPER" if padding == "same" else "VALID" - pads = None - else: - p = list(padding) if hasattr(padding, "__iter__") else [padding, padding] - pads = p + p - auto_pad = "NOTSET" - - to_list = lambda v, n: list(v) if hasattr(v, "__iter__") else [v] * n # noqa: E731 - conv_attrs = dict( - kernel_shape=to_list(layer.kernel_size, 2), - strides=to_list(layer.strides, 2), - dilations=to_list(layer.dilation_rate, 2), - group=in_ch, - auto_pad=auto_pad, - ) - if pads is not None: - conv_attrs["pads"] = pads - - conv_out = f"{prefix}_conv" - nodes.append(oh.make_node("Conv", inputs=conv_inputs, outputs=[conv_out], **conv_attrs)) - current = conv_out - - current = _maybe_quant_output(layer, prefix, current, nodes, initializers, quant_fn) - - if cl: - current = _add_transpose(f"{prefix}_post", current, [0, 2, 3, 1], nodes) - return current - - -def _add_batchnorm(layer, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights): - """PQBatchNormalization / standard BatchNormalization.""" - need_tr, perm_to_nchw, perm_to_nhwx = _bn_transpose_info(layer) - - if need_tr: - current = _add_transpose(f"{prefix}_pre", current, perm_to_nchw, nodes) - - current = _maybe_quant_input(layer, prefix, current, nodes, initializers, quant_fn) - - is_pq = isinstance(layer, PQBatchNormalization) - - gamma_np = _np(layer.gamma) if layer.gamma is not None else None - beta_np = _np(layer.beta) if layer.beta is not None else None - - if gamma_np is None: - # scale=False: use ones - n_ch = _np(layer.moving_mean).shape[0] - gamma_np = np.ones(n_ch, dtype=np.float32) - if beta_np is None: - # center=False: use zeros - n_ch = _np(layer.moving_mean).shape[0] - beta_np = np.zeros(n_ch, dtype=np.float32) - - if is_pq and use_qonnx: - k_w, i_w, f_w = layer.weight_quantizer.get_quantization_bits() - gfp = f"{prefix}_gamma_fp" - initializers.append(onh.from_array(gamma_np, name=gfp)) - g_nodes, q_gamma = _quant_node( - f"{prefix}_gamma", - gfp, - layer.weight_quantizer.round_mode, - _np(k_w), - _np(i_w), - _np(f_w), - initializers, - overflow_mode=getattr(layer.weight_quantizer, "overflow", "SAT"), - ) - nodes.extend(g_nodes) - - k_b, i_b, f_b = layer.bias_quantizer.get_quantization_bits() - bfp = f"{prefix}_beta_fp" - initializers.append(onh.from_array(beta_np, name=bfp)) - b_nodes, q_beta = _quant_node( - f"{prefix}_beta", - bfp, - layer.bias_quantizer.round_mode, - _np(k_b), - _np(i_b), - _np(f_b), - initializers, - overflow_mode=getattr(layer.bias_quantizer, "overflow", "SAT"), - ) - nodes.extend(b_nodes) - elif is_pq and store_integer_weights: - k_w, i_w, f_w = layer.weight_quantizer.get_quantization_bits() - g_nodes, q_gamma = _int_weight_node(f"{prefix}_gamma", gamma_np, _np(k_w), _np(i_w), _np(f_w), initializers) - nodes.extend(g_nodes) - k_b, i_b, f_b = layer.bias_quantizer.get_quantization_bits() - b_nodes, q_beta = _int_weight_node(f"{prefix}_beta", beta_np, _np(k_b), _np(i_b), _np(f_b), initializers) - nodes.extend(b_nodes) - else: - q_gamma = f"{prefix}_gamma" - q_beta = f"{prefix}_beta" - initializers.append(onh.from_array(gamma_np, name=q_gamma)) - initializers.append(onh.from_array(beta_np, name=q_beta)) - - mean_name = f"{prefix}_running_mean" - var_name = f"{prefix}_running_var" - initializers.append(onh.from_array(_np(layer.moving_mean), name=mean_name)) - initializers.append(onh.from_array(_np(layer.moving_variance), name=var_name)) - - bn_out = f"{prefix}_bn" - nodes.append( - oh.make_node( - "BatchNormalization", - inputs=[current, q_gamma, q_beta, mean_name, var_name], - outputs=[bn_out], - epsilon=float(layer.epsilon), - ) - ) - current = bn_out - - if need_tr: - current = _add_transpose(f"{prefix}_post", current, perm_to_nhwx, nodes) - return current - - -def _add_dense_nd(layer, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights): - """PQDense for rank-3 inputs (B, T, E). - - Uses MatMul + Add instead of Gemm so the op works for any rank ≥ 2. - The kernel is stored as [out, in] (same layout as _add_dense / _int_weight_node), - then transposed to [in, out] at runtime via a Transpose node so that - MatMul(input, kernel_t) broadcasts correctly over the sequence dimension. - """ - current = _maybe_quant_input(layer, prefix, current, nodes, initializers, quant_fn) - - kernel_np = _np(layer._kernel).T # [out, in] - out_units = kernel_np.shape[0] - - if use_qonnx: - k_w, i_w, f_w = layer.weight_quantizer.get_quantization_bits() - wfp_name = f"{prefix}_weight_fp" - initializers.append(onh.from_array(kernel_np, name=wfp_name)) - w_nodes, q_weight = _quant_node( - f"{prefix}_weight", - wfp_name, - layer.weight_quantizer.round_mode, - _np(k_w), - _np(i_w), - _np(f_w), - initializers, - overflow_mode=getattr(layer.weight_quantizer, "overflow", "SAT"), - ) - nodes.extend(w_nodes) - elif store_integer_weights: - k_w, i_w, f_w = layer.weight_quantizer.get_quantization_bits() - f_for_onnx = _weight_f_for_onnx(_np(f_w), out_units) - k_for_onnx = _weight_f_for_onnx(_np(k_w), out_units) - i_for_onnx = _weight_f_for_onnx(_np(i_w), out_units) - w_nodes, q_weight = _int_weight_node(f"{prefix}_weight", kernel_np, k_for_onnx, i_for_onnx, f_for_onnx, initializers) - nodes.extend(w_nodes) - else: - q_weight = f"{prefix}_weight" - initializers.append(onh.from_array(kernel_np, name=q_weight)) - - # Transpose [out, in] → [in, out] so MatMul(input[..., in], kernel_t[in, out]) works - kernel_t_name = f"{prefix}_weight_t" - nodes.append(oh.make_node("Transpose", inputs=[q_weight], outputs=[kernel_t_name], perm=[1, 0])) - - mm_out = f"{prefix}_mm" - nodes.append(oh.make_node("MatMul", inputs=[current, kernel_t_name], outputs=[mm_out])) - current = mm_out - - if layer._bias is not None: - bias_np = _np(layer._bias) - if use_qonnx: - k_b, i_b, f_b = layer.bias_quantizer.get_quantization_bits() - bfp_name = f"{prefix}_bias_fp" - initializers.append(onh.from_array(bias_np, name=bfp_name)) - b_nodes, q_bias = _quant_node( - f"{prefix}_bias", - bfp_name, - layer.bias_quantizer.round_mode, - _np(k_b), - _np(i_b), - _np(f_b), - initializers, - overflow_mode=getattr(layer.bias_quantizer, "overflow", "SAT"), - ) - nodes.extend(b_nodes) - elif store_integer_weights: - k_b, i_b, f_b = layer.bias_quantizer.get_quantization_bits() - b_nodes, q_bias = _int_weight_node(f"{prefix}_bias", bias_np, _np(k_b), _np(i_b), _np(f_b), initializers) - nodes.extend(b_nodes) - else: - q_bias = f"{prefix}_bias" - initializers.append(onh.from_array(bias_np, name=q_bias)) - add_out = f"{prefix}_bias_add" - nodes.append(oh.make_node("Add", inputs=[current, q_bias], outputs=[add_out])) - current = add_out - - current = _maybe_quant_output(layer, prefix, current, nodes, initializers, quant_fn) - return current - - -def _add_mha(layer, prefix, q_input, k_input, v_input, nodes, initializers, quant_fn, use_qonnx, store_integer_weights): - """Build ONNX nodes for PQMultiheadAttention (Keras version, always batch-first). - - Decomposes multi-head attention into primitive ONNX ops: - - Q/K/V MatMul projections (rank-3 MatMul via _add_dense_nd) - Reshape (B, L, E) → (B, H, L, head_dim) + Transpose - MatMul(Q, K^T) * scale → optional Quant - Softmax → optional Quant - MatMul(attn_weights, V) → optional context Quant - Transpose + Reshape → (B, T, E) - out_proj MatMul - - Returns (out_name, avg_attn_weights_name). - """ - H = layer.num_heads - head_dim = layer.head_dim - E = layer.embed_dim - scale_val = float(layer.scale) - - # --- Q / K / V projections: (B, L, E) → (B, L, E) --- - q_proj_out = _add_dense_nd( - layer.q_proj, f"{prefix}_q_proj", q_input, nodes, initializers, quant_fn, use_qonnx, store_integer_weights - ) - k_proj_out = _add_dense_nd( - layer.k_proj, f"{prefix}_k_proj", k_input, nodes, initializers, quant_fn, use_qonnx, store_integer_weights - ) - v_proj_out = _add_dense_nd( - layer.v_proj, f"{prefix}_v_proj", v_input, nodes, initializers, quant_fn, use_qonnx, store_integer_weights - ) - - # --- Helper: (B, L, E) → (B, H, L, head_dim) using dynamic shapes --- - def _split_heads(x_name, pfx): - shape_out = f"{pfx}_shape" - b_scalar = f"{pfx}_b_sc" - l_scalar = f"{pfx}_l_sc" - b_1d = f"{pfx}_b_1d" - l_1d = f"{pfx}_l_1d" - h_1d_const = f"{pfx}_H_1d" - hd_1d_const = f"{pfx}_hd_1d" - shape_4d = f"{pfx}_shape4d" - reshaped = f"{pfx}_reshaped" - transposed = f"{pfx}_transposed" - idx0 = f"{pfx}_gi0" - idx1 = f"{pfx}_gi1" - ax0 = f"{pfx}_ax0" - - nodes.append(oh.make_node("Shape", inputs=[x_name], outputs=[shape_out])) - initializers.extend( - [ - onh.from_array(np.array(0, dtype=np.int64), name=idx0), - onh.from_array(np.array(1, dtype=np.int64), name=idx1), - onh.from_array(np.array([0], dtype=np.int64), name=ax0), - onh.from_array(np.array([H], dtype=np.int64), name=h_1d_const), - onh.from_array(np.array([head_dim], dtype=np.int64), name=hd_1d_const), - ] - ) - nodes.append(oh.make_node("Gather", inputs=[shape_out, idx0], outputs=[b_scalar])) - nodes.append(oh.make_node("Gather", inputs=[shape_out, idx1], outputs=[l_scalar])) - nodes.append(oh.make_node("Unsqueeze", inputs=[b_scalar, ax0], outputs=[b_1d])) - nodes.append(oh.make_node("Unsqueeze", inputs=[l_scalar, ax0], outputs=[l_1d])) - nodes.append(oh.make_node("Concat", inputs=[b_1d, l_1d, h_1d_const, hd_1d_const], outputs=[shape_4d], axis=0)) - nodes.append(oh.make_node("Reshape", inputs=[x_name, shape_4d], outputs=[reshaped])) - # (B, L, H, head_dim) → (B, H, L, head_dim) - nodes.append(oh.make_node("Transpose", inputs=[reshaped], outputs=[transposed], perm=[0, 2, 1, 3])) - return transposed - - q_h = _split_heads(q_proj_out, f"{prefix}_q") - k_h = _split_heads(k_proj_out, f"{prefix}_k") - v_h = _split_heads(v_proj_out, f"{prefix}_v") - - k_t_name = f"{prefix}_k_T" - nodes.append(oh.make_node("Transpose", inputs=[k_h], outputs=[k_t_name], perm=[0, 1, 3, 2])) - - raw_scores = f"{prefix}_scores_raw" - scaled_scores = f"{prefix}_scores_scaled" - scale_cst = f"{prefix}_attn_scale" - nodes.append(oh.make_node("MatMul", inputs=[q_h, k_t_name], outputs=[raw_scores])) - initializers.append(onh.from_array(np.array(scale_val, dtype=np.float32), name=scale_cst)) - nodes.append(oh.make_node("Mul", inputs=[raw_scores, scale_cst], outputs=[scaled_scores])) - current = scaled_scores - - if layer.softmax.quantize_input and getattr(layer, "enable_quantization", True): - q = layer.softmax.input_quantizer - k_q, i_q, f_q = q.get_quantization_bits() - q_nodes, current = quant_fn( - f"{prefix}_attn_score_q", - current, - q.round_mode, - _np(k_q), - _np(i_q), - _np(f_q), - initializers, - overflow_mode=getattr(q, "overflow", "SAT"), - ) - nodes.extend(q_nodes) - - # --- Softmax (axis=-1); approximate_softmax falls back to standard Softmax in ONNX --- - attn_w_name = f"{prefix}_attn_weights" - nodes.append(oh.make_node("Softmax", inputs=[current], outputs=[attn_w_name], axis=-1)) - current = attn_w_name - - # --- Softmax output quantization (the MHA enables the softmax's output quantizer) --- - if layer.softmax.quantize_output and getattr(layer, "enable_quantization", True): - q = layer.softmax.output_quantizer - k_q, i_q, f_q = q.get_quantization_bits() - q_nodes, current = quant_fn( - f"{prefix}_attn_weight_q", - current, - q.round_mode, - _np(k_q), - _np(i_q), - _np(f_q), - initializers, - overflow_mode=getattr(q, "overflow", "SAT"), - ) - nodes.extend(q_nodes) - - ctx_raw = f"{prefix}_ctx_raw" - nodes.append(oh.make_node("MatMul", inputs=[current, v_h], outputs=[ctx_raw])) - current_ctx = ctx_raw - - ctx_t = f"{prefix}_ctx_t" - ctx_shape = f"{prefix}_ctx_shape" - ctx_b_sc = f"{prefix}_ctx_b_sc" - ctx_t_sc = f"{prefix}_ctx_t_sc" - ctx_b_1d = f"{prefix}_ctx_b_1d" - ctx_t_1d = f"{prefix}_ctx_t_1d" - ctx_E_1d = f"{prefix}_ctx_E_1d" - ctx_ax0 = f"{prefix}_ctx_ax0" - ctx_gi0 = f"{prefix}_ctx_gi0" - ctx_gi1 = f"{prefix}_ctx_gi1" - ctx_3d = f"{prefix}_ctx_shape3d" - ctx_merged = f"{prefix}_ctx_merged" - - nodes.append(oh.make_node("Transpose", inputs=[current_ctx], outputs=[ctx_t], perm=[0, 2, 1, 3])) - nodes.append(oh.make_node("Shape", inputs=[ctx_t], outputs=[ctx_shape])) - initializers += [ - onh.from_array(np.array(0, dtype=np.int64), name=ctx_gi0), - onh.from_array(np.array(1, dtype=np.int64), name=ctx_gi1), - onh.from_array(np.array([0], dtype=np.int64), name=ctx_ax0), - onh.from_array(np.array([E], dtype=np.int64), name=ctx_E_1d), - ] - nodes.append(oh.make_node("Gather", inputs=[ctx_shape, ctx_gi0], outputs=[ctx_b_sc])) - nodes.append(oh.make_node("Gather", inputs=[ctx_shape, ctx_gi1], outputs=[ctx_t_sc])) - nodes.append(oh.make_node("Unsqueeze", inputs=[ctx_b_sc, ctx_ax0], outputs=[ctx_b_1d])) - nodes.append(oh.make_node("Unsqueeze", inputs=[ctx_t_sc, ctx_ax0], outputs=[ctx_t_1d])) - nodes.append(oh.make_node("Concat", inputs=[ctx_b_1d, ctx_t_1d, ctx_E_1d], outputs=[ctx_3d], axis=0)) - nodes.append(oh.make_node("Reshape", inputs=[ctx_t, ctx_3d], outputs=[ctx_merged])) - - # --- Output projection: (B, T, E) → (B, T, E) --- - out = _add_dense_nd( - layer.out_proj, f"{prefix}_out_proj", ctx_merged, nodes, initializers, quant_fn, use_qonnx, store_integer_weights - ) - - # --- Average attention weights over heads: (B, H, T, S) → (B, T, S) --- - avg_attn = f"{prefix}_avg_attn_weights" - nodes.append(oh.make_node("ReduceMean", inputs=[attn_w_name], outputs=[avg_attn], axes=[1], keepdims=0)) - - return out, avg_attn - - -def _add_avgpool(layer, prefix, current, nodes, initializers, ndim, quant_fn): - cl = _channels_last(layer) - - if cl: - perm_to_nchw = [0, 3, 1, 2] if ndim == 2 else [0, 2, 1] - perm_to_nhwx = [0, 2, 3, 1] if ndim == 2 else [0, 2, 1] - current = _add_transpose(f"{prefix}_pre", current, perm_to_nchw, nodes) - - current = _maybe_quant_input(layer, prefix, current, nodes, initializers, quant_fn) - - to_list = lambda v, n: list(v) if hasattr(v, "__iter__") else [v] * n # noqa: E731 - - pool_out = f"{prefix}_pool" - nodes.append( - oh.make_node( - "AveragePool", - inputs=[current], - outputs=[pool_out], - kernel_shape=to_list(layer.pool_size, ndim), - strides=to_list(layer.strides, ndim), - pads=[0] * (ndim * 2), - count_include_pad=0, - ) - ) - current = pool_out - - current = _maybe_quant_output(layer, prefix, current, nodes, initializers, quant_fn) - - if cl: - current = _add_transpose(f"{prefix}_post", current, perm_to_nhwx, nodes) - return current - - -def _add_global_avgpool(layer, prefix, current, nodes, ndim): - cl = _channels_last(layer) - - if cl: - perm_to_nchw = [0, 3, 1, 2] if ndim == 2 else [0, 2, 1] - current = _add_transpose(f"{prefix}_pre", current, perm_to_nchw, nodes) - - pool_out = f"{prefix}_global_pool" - nodes.append(oh.make_node("GlobalAveragePool", inputs=[current], outputs=[pool_out])) - current = pool_out - - if cl: - # GlobalAveragePool returns [N, C, 1, 1]; emit Flatten to [N, C]. - # Actually after GlobalAveragePool output is [N, C, 1, 1]; transpose back would give - # [N, 1, 1, C] which then needs squeezing — that's the same as just squeezing [N, C]. - # Emit Flatten to [N, C] instead of bothering with transpose. - flatten_name = f"{prefix}_flatten" - nodes.append(oh.make_node("Flatten", inputs=[pool_out], outputs=[flatten_name], axis=1)) - current = flatten_name - - return current - - -def _add_pq_activation(layer, prefix, current, nodes, initializers, quant_fn): - """PQActivation: [input QDQ] → [multiplier scale] → activation → [output QDQ]. - - Supported activations: relu, tanh, hard_tanh (= Clip(-1, 1)). - - The optional relu multiplier is baked to a constant: 2^round(m). - """ - # --- optional input quantization --- - current = _maybe_quant_input(layer, prefix, current, nodes, initializers, quant_fn) - - # --- optional learnable multiplier (relu only) --- - if ( - getattr(layer, "use_multiplier", False) - and getattr(layer, "activation_name", "") == "relu" - and hasattr(layer, "multiplier") - ): - m_val = float(np.array(layer.multiplier).ravel()[0]) - scale = float(2.0 ** round(m_val)) - scale_name = f"{prefix}_mul_scale" - scaled_out = f"{prefix}_scaled" - initializers.append(onh.from_array(np.array(scale, dtype=np.float32), name=scale_name)) - nodes.append(oh.make_node("Mul", inputs=[current, scale_name], outputs=[scaled_out])) - current = scaled_out - - # --- activation --- - act = getattr(layer, "activation_name", "relu") - act_out = f"{prefix}_act" - if act == "relu": - nodes.append(oh.make_node("Relu", inputs=[current], outputs=[act_out])) - elif act == "tanh": - nodes.append(oh.make_node("Tanh", inputs=[current], outputs=[act_out])) - elif act == "hard_tanh": - # hard_tanh(x) = clip(x, -1, 1) - cmin_name = f"{prefix}_htanh_min" - cmax_name = f"{prefix}_htanh_max" - initializers += [ - onh.from_array(np.array(-1.0, dtype=np.float32), name=cmin_name), - onh.from_array(np.array(1.0, dtype=np.float32), name=cmax_name), - ] - nodes.append(oh.make_node("Clip", inputs=[current, cmin_name, cmax_name], outputs=[act_out])) - else: - raise TypeError(f"PQActivation: unsupported activation {act!r} for ONNX export") - current = act_out - - # --- optional output quantization --- - current = _maybe_quant_output(layer, prefix, current, nodes, initializers, quant_fn) - return current - - -# --------------------------------------------------------------------------- -# shared layer dispatcher -# --------------------------------------------------------------------------- - - -def _emit_layer( - layer, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights, input_onnx_names=None -): - """Emit ONNX nodes for a single Keras layer. Returns the ONNX output name.""" - - # --- PQuant layers --- - if isinstance(layer, PQMultiheadAttention): - # input_onnx_names = [query, key, value] or [single_input] for self-attention - if len(input_onnx_names) >= 3: - q_in, k_in, v_in = input_onnx_names[0], input_onnx_names[1], input_onnx_names[2] - elif len(input_onnx_names) == 2: - q_in, k_in, v_in = input_onnx_names[0], input_onnx_names[1], input_onnx_names[1] - else: - q_in = k_in = v_in = input_onnx_names[0] - return _add_mha(layer, prefix, q_in, k_in, v_in, nodes, initializers, quant_fn, use_qonnx, store_integer_weights) - - if isinstance(layer, PQActivation): - return _add_pq_activation(layer, prefix, current, nodes, initializers, quant_fn) - - if isinstance(layer, PQDense): - return _add_dense(layer, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights) - - if isinstance(layer, PQDepthwiseConv2d): - return _add_depthwise_conv(layer, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights) - - if isinstance(layer, PQConv2d): - return _add_conv( - layer, - prefix, - current, - nodes, - initializers, - ndim=2, - quant_fn=quant_fn, - use_qonnx=use_qonnx, - store_integer_weights=store_integer_weights, - ) - - if isinstance(layer, PQConv1d): - return _add_conv( - layer, - prefix, - current, - nodes, - initializers, - ndim=1, - quant_fn=quant_fn, - use_qonnx=use_qonnx, - store_integer_weights=store_integer_weights, - ) - - if isinstance(layer, PQBatchNormalization): - return _add_batchnorm(layer, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights) - - # --- Standard Keras layers --- - if isinstance(layer, keras.layers.BatchNormalization): - return _add_batchnorm( - layer, prefix, current, nodes, initializers, quant_fn=quant_fn, use_qonnx=False, store_integer_weights=False - ) - - if isinstance(layer, keras.layers.Conv2D): - # Plain Conv2D (non-PQ): wrap in a minimal shim - _layer = layer - _layer._bias = layer.bias - return _add_conv_plain(layer, prefix, current, nodes, initializers) - - if isinstance(layer, keras.layers.Dense): - out = f"{prefix}_gemm" - w_name = f"{prefix}_weight" - initializers.append(onh.from_array(_np(layer.kernel).T, name=w_name)) # store [out, in] - gemm_inputs = [current, w_name] - if layer.bias is not None: - b_name = f"{prefix}_bias" - initializers.append(onh.from_array(_np(layer.bias), name=b_name)) - gemm_inputs.append(b_name) - nodes.append(oh.make_node("Gemm", inputs=gemm_inputs, outputs=[out], transB=1)) - return out - - if isinstance(layer, (keras.layers.ReLU, keras.layers.Activation)): - activation = ( - layer.activation.__name__ - if isinstance(layer, keras.layers.Activation) and callable(layer.activation) - else getattr(layer, "activation", "relu") - ) - act_name = activation if isinstance(activation, str) else "relu" - out = f"{prefix}_act" - if "relu" in act_name.lower(): - nodes.append(oh.make_node("Relu", inputs=[current], outputs=[out])) - elif "sigmoid" in act_name.lower(): - nodes.append(oh.make_node("Sigmoid", inputs=[current], outputs=[out])) - elif "tanh" in act_name.lower(): - nodes.append(oh.make_node("Tanh", inputs=[current], outputs=[out])) - else: - raise TypeError(f"Unsupported Activation for ONNX export: {act_name!r}") - return out - - if isinstance(layer, keras.layers.Flatten): - out = f"{prefix}_flatten" - nodes.append(oh.make_node("Flatten", inputs=[current], outputs=[out], axis=1)) - return out - - if isinstance(layer, keras.layers.Reshape): - target_shape = list(layer.target_shape) - # Prepend batch dim (-1 means dynamic) - full_shape = [-1] + target_shape - shape_name = f"{prefix}_shape" - out = f"{prefix}_reshape" - initializers.append(onh.from_array(np.array(full_shape, dtype=np.int64), name=shape_name)) - nodes.append(oh.make_node("Reshape", inputs=[current, shape_name], outputs=[out])) - return out - - if isinstance(layer, keras.layers.Add): - assert input_onnx_names is not None and len(input_onnx_names) == 2 - out = f"{prefix}_add" - nodes.append(oh.make_node("Add", inputs=input_onnx_names, outputs=[out])) - return out - - if isinstance(layer, keras.layers.Concatenate): - assert input_onnx_names is not None - axis = layer.axis - # Negative axis: leave as-is; onnx Concat supports negative axes - out = f"{prefix}_concat" - nodes.append(oh.make_node("Concat", inputs=input_onnx_names, outputs=[out], axis=axis)) - return out - - if isinstance(layer, keras.layers.Multiply): - assert input_onnx_names is not None and len(input_onnx_names) == 2 - out = f"{prefix}_mul" - nodes.append(oh.make_node("Mul", inputs=input_onnx_names, outputs=[out])) - return out - - if isinstance(layer, keras.layers.AveragePooling2D): - return _add_avgpool(layer, prefix, current, nodes, initializers, ndim=2, quant_fn=quant_fn) - - if isinstance(layer, keras.layers.AveragePooling1D): - return _add_avgpool(layer, prefix, current, nodes, initializers, ndim=1, quant_fn=quant_fn) - - if isinstance(layer, keras.layers.GlobalAveragePooling2D): - return _add_global_avgpool(layer, prefix, current, nodes, ndim=2) - - if isinstance(layer, keras.layers.GlobalAveragePooling1D): - return _add_global_avgpool(layer, prefix, current, nodes, ndim=1) - - if isinstance(layer, (keras.layers.Dropout,)): - return current # identity at inference - - raise TypeError(f"Unsupported Keras layer type for ONNX export: {type(layer).__name__!r}") - - -def _add_conv_plain(layer, prefix, current, nodes, initializers): - """Emit a plain (non-PQ) Conv2D layer.""" - cl = _channels_last(layer) - if cl: - current = _add_transpose(f"{prefix}_pre", current, [0, 3, 1, 2], nodes) - - kernel_np = _np(layer.kernel) - kernel_onnx = np.transpose(kernel_np, (3, 2, 0, 1)) - w_name = f"{prefix}_weight" - initializers.append(onh.from_array(kernel_onnx, name=w_name)) - conv_inputs = [current, w_name] - - if layer.bias is not None: - b_name = f"{prefix}_bias" - initializers.append(onh.from_array(_np(layer.bias), name=b_name)) - conv_inputs.append(b_name) - - padding = layer.padding - auto_pad = "SAME_UPPER" if padding == "same" else "VALID" - to_list = lambda v, n: list(v) if hasattr(v, "__iter__") else [v] * n # noqa: E731 - conv_attrs = dict( - kernel_shape=to_list(layer.kernel_size, 2), - strides=to_list(layer.strides, 2), - dilations=to_list(layer.dilation_rate, 2), - group=layer.groups, - auto_pad=auto_pad, - ) - conv_out = f"{prefix}_conv" - nodes.append(oh.make_node("Conv", inputs=conv_inputs, outputs=[conv_out], **conv_attrs)) - current = conv_out - - if cl: - current = _add_transpose(f"{prefix}_post", current, [0, 2, 3, 1], nodes) - return current - - -# --------------------------------------------------------------------------- -# Keras functional model graph traversal -# --------------------------------------------------------------------------- - - -def _build_tensor_onnx_map(model): - """ - Return a dict mapping id(KerasTensor) → ONNX tensor name for model.inputs. - Multi-input models are supported; inputs are named "input_0", "input_1", etc. - (or just "input" for single-input models). - """ - tensor_to_onnx = {} - for i, inp in enumerate(model.inputs): - name = "input" if len(model.inputs) == 1 else f"input_{i}" - tensor_to_onnx[id(inp)] = name - return tensor_to_onnx - - -def _inbound_input_names(layer, tensor_to_onnx): - """Return the list of ONNX input names for this layer based on its inbound node.""" - if not layer._inbound_nodes: - return [] - node = layer._inbound_nodes[0] - input_tensors = node.input_tensors - if not isinstance(input_tensors, (list, tuple)): - input_tensors = [input_tensors] - result = [] - for t in input_tensors: - key = id(t) - if key not in tensor_to_onnx: - raise RuntimeError( - f"Layer {layer.name!r}: input tensor not found in tensor_to_onnx map. " - "Ensure model.layers is in topological order." - ) - result.append(tensor_to_onnx[key]) - return result - - -def _register_layer_output(layer, onnx_name, tensor_to_onnx): - """Register the ONNX output name for a layer's output tensor(s). - - onnx_name may be a plain string (single-output layer) or a tuple of strings - (multi-output layer, e.g. PQMultiheadAttention returns (out, avg_attn_weights)). - """ - if not layer._inbound_nodes: - return - node = layer._inbound_nodes[0] - out_tensors = node.output_tensors - if not isinstance(out_tensors, (list, tuple)): - out_tensors = [out_tensors] - if isinstance(onnx_name, (list, tuple)): - for tensor, name in zip(out_tensors, onnx_name): - tensor_to_onnx[id(tensor)] = name - else: - tensor_to_onnx[id(out_tensors[0])] = onnx_name - - -# --------------------------------------------------------------------------- -# main conversion -# --------------------------------------------------------------------------- - - -def convert_to_onnx( - model: keras.Model, - input_shape: tuple, - output_path: str = "model.onnx", - opset: int = 13, - use_qonnx: bool = False, - store_integer_weights: bool = False, - include_clip: bool = True, - batch_size: int | None = None, -) -> onnx.ModelProto: - """ - Convert a Keras functional model of PQuant layers to ONNX or QONNX. - - The model must have apply_final_compression() called on all PQ layers - before calling this function. Only inference-mode semantics are exported. - - Args: - model: Trained keras.Model. Must be a functional model - (built with the Keras functional API or subclassed - models whose layers are accessible via model.layers). - input_shape: Shape of a single sample excluding batch, e.g. (3, 32, 32). - For channels_last Conv models use e.g. (32, 32, 3). - output_path: Where to save the .onnx file. - opset: ONNX opset version (≥13 required for per-channel - DequantizeLinear). - use_qonnx: Emit QONNX Quant custom nodes if True. - store_integer_weights: Store weight initializers as int8/uint8 + - DequantizeLinear instead of float32 (ignored when - use_qonnx=True). - include_clip: Prepend a Clip node before each QuantizeLinear when - True (default). Set to False to emit bare - QuantizeLinear+DequantizeLinear pairs — safe when - values are guaranteed in-range at inference time since - QuantizeLinear saturates naturally. Ignored when - use_qonnx=True. - batch_size: If not None, fix the batch dimension of all graph - inputs and outputs to this value. If None (default), - the batch dimension is left dynamic. - - Returns: - The constructed onnx.ModelProto. - """ - quant_fn = _quant_node if use_qonnx else functools.partial(_qdq_node, include_clip=include_clip) - - onnx_nodes: list[onnx.NodeProto] = [] - initializers: list[onnx.TensorProto] = [] - - tensor_to_onnx = _build_tensor_onnx_map(model) - last_output_name: str = "" - - for layer in model.layers: - # Skip InputLayer — already seeded in tensor_to_onnx - if isinstance(layer, keras.layers.InputLayer): - continue - - input_onnx_names = _inbound_input_names(layer, tensor_to_onnx) - if not input_onnx_names: - continue - - current = input_onnx_names[0] # primary input (used by single-input layers) - prefix = layer.name.replace("/", "_").replace(":", "_") - - output_name = _emit_layer( - layer, - prefix, - current, - onnx_nodes, - initializers, - quant_fn, - use_qonnx, - store_integer_weights, - input_onnx_names=input_onnx_names, - ) - - _register_layer_output(layer, output_name, tensor_to_onnx) - # For multi-output layers (e.g. MHA returns (out, avg_attn)), track only the - # primary output as the graph's last output name. - last_output_name = output_name[0] if isinstance(output_name, tuple) else output_name - - # Determine output shape via a forward pass - dummy = np.zeros((1, *input_shape), dtype=np.float32) - dummy_out = model(dummy, training=False) - dummy_out_np = np.array(ops.convert_to_numpy(dummy_out)) - batch_dim = batch_size # None → dynamic, int → fixed - output_shape = [batch_dim] + list(dummy_out_np.shape[1:]) - - # Build ONNX graph - if len(model.inputs) == 1: - input_vis = [oh.make_tensor_value_info("input", TensorProto.FLOAT, [batch_dim, *input_shape])] - else: - input_vis = [ - oh.make_tensor_value_info(f"input_{i}", TensorProto.FLOAT, [batch_dim, *input_shape]) - for i in range(len(model.inputs)) - ] - output_vi = oh.make_tensor_value_info(last_output_name, TensorProto.FLOAT, output_shape) - - graph = oh.make_graph( - nodes=onnx_nodes, - name="pquant_keras_onnx", - inputs=input_vis, - outputs=[output_vi], - initializer=initializers, - ) - - opset_imports = [oh.make_opsetid("", opset)] - if use_qonnx: - opset_imports.append(oh.make_opsetid("qonnx.custom_op.general", 1)) - model_proto = oh.make_model(graph, opset_imports=opset_imports) - model_proto.ir_version = 6 - - # ONNX opset >= 9: initializers are implicit constants and must NOT appear in - # graph.input — otherwise tools treat weight tensors as runtime inputs. - # Some onnx library versions add them automatically for backward compatibility; - # strip them here so only the actual data inputs remain. - _init_names = {t.name for t in model_proto.graph.initializer} - _data_inputs = [vi for vi in model_proto.graph.input if vi.name not in _init_names] - del model_proto.graph.input[:] - model_proto.graph.input.extend(_data_inputs) - - onnx.checker.check_model(model_proto) - onnx.save(model_proto, output_path) - fmt = "QONNX" if use_qonnx else "ONNX (QDQ)" - logging.info("Saved %s Keras model → %s", fmt, output_path) - return model_proto - - -# --------------------------------------------------------------------------- -# usage example -# --------------------------------------------------------------------------- - -if __name__ == "__main__": - import pquant - from pquant import apply_final_compression - - cfg = pquant.pdp_config() - - inp = keras.Input(shape=(3, 32, 32)) - x = PQConv2d(cfg, filters=16, kernel_size=3, padding="same")(inp) - x = PQBatchNormalization(cfg)(x) - x = keras.layers.ReLU()(x) - x = PQConv2d(cfg, filters=32, kernel_size=3, padding="same")(x) - x = keras.layers.ReLU()(x) - x = keras.layers.Flatten()(x) - x = PQDense(cfg, units=10)(x) - model = keras.Model(inp, x) - - apply_final_compression(model) - - convert_to_onnx(model, input_shape=(3, 32, 32), output_path="model_keras.onnx") - - import onnxruntime as ort - - sess = ort.InferenceSession("model_keras.onnx") - out = sess.run(None, {"input": np.random.randn(2, 3, 32, 32).astype(np.float32)}) - print("Output shape:", out[0].shape) # noqa: T201 diff --git a/src/pquant/core/keras/layers.py b/src/pquant/core/keras/layers.py index 4a058c0..ed61973 100644 --- a/src/pquant/core/keras/layers.py +++ b/src/pquant/core/keras/layers.py @@ -1955,6 +1955,32 @@ def hgq_loss(self): return ops.convert_to_tensor(0.0) return ops.convert_to_tensor(self.hgq_beta * self._attention_ebops() + self.softmax.hgq_loss()) + @staticmethod + def _split_qkv_shapes(input_shape): + if isinstance(input_shape, (list, tuple)) and len(input_shape) > 0 and isinstance(input_shape[0], (list, tuple)): + q_shape = input_shape[0] + k_shape = input_shape[1] if len(input_shape) > 1 else q_shape + v_shape = input_shape[2] if len(input_shape) > 2 else k_shape + else: + q_shape = k_shape = v_shape = input_shape + return q_shape, k_shape, v_shape + + def compute_output_shape(self, input_shape): + q_shape, k_shape, _ = self._split_qkv_shapes(input_shape) + batch, tgt_len = q_shape[0], q_shape[1] + src_len = k_shape[1] + return (batch, tgt_len, self.embed_dim), (batch, tgt_len, src_len) + + def build(self, input_shape): + q_shape, k_shape, v_shape = self._split_qkv_shapes(input_shape) + self.q_proj.build(q_shape) + self.k_proj.build(k_shape) + self.v_proj.build(v_shape) + self.out_proj.build(tuple(q_shape[:-1]) + (self.embed_dim,)) + # Softmax operates on the per-head attention scores (B, H, Tq, Tk). + self.softmax.build((q_shape[0], self.num_heads, q_shape[1], k_shape[1])) + super().build(input_shape) + def call( self, inputs, diff --git a/src/pquant/core/keras/onnx/__init__.py b/src/pquant/core/keras/onnx/__init__.py new file mode 100644 index 0000000..6a27e0e --- /dev/null +++ b/src/pquant/core/keras/onnx/__init__.py @@ -0,0 +1,5 @@ +from pquant.core.keras.onnx.convert_to_onnx import ( + convert_to_onnx, +) + +__all__ = ["convert_to_onnx"] diff --git a/src/pquant/core/keras/onnx/convert_to_onnx.py b/src/pquant/core/keras/onnx/convert_to_onnx.py new file mode 100644 index 0000000..c0b7a6f --- /dev/null +++ b/src/pquant/core/keras/onnx/convert_to_onnx.py @@ -0,0 +1,402 @@ +""" +Convert a PQuant Keras functional model to ONNX or QONNX format. + +Pass ``use_qonnx=True`` to emit QONNX ``Quant`` custom nodes (requires the +qonnx runtime). Pass ``use_qonnx=False`` (default) to emit standard +``Clip + QuantizeLinear + DequantizeLinear`` nodes runnable with plain +onnxruntime. + + +Keras weight layout (kernel always stored as HWIO regardless of data_format): + Conv2D kernel: [kH, kW, in/g, out] → [out, in/g, kH, kW] for ONNX + Conv1D kernel: [kL, in/g, out] → [out, in/g, kL] for ONNX + DepthwiseConv2D kernel: [kH, kW, in, dm] → [in*dm, 1, kH, kW] for ONNX + Dense kernel: [in, out] → stored as [out, in] for Gemm (transB=1) +""" + +import functools +import logging + +import keras +import numpy as np +import onnx +import onnx.helper as oh +from keras import ops +from onnx import TensorProto + +from pquant.core.keras.activations import PQActivation +from pquant.core.keras.layers import ( + PQBatchNormalization, + PQConv1d, + PQConv2d, + PQDense, + PQDepthwiseConv2d, + PQMultiheadAttention, +) +from pquant.core.keras.onnx.helpers import _keras_dtype_to_tp +from pquant.core.keras.onnx.layer_builders import ( + _add_avgpool, + _add_batchnorm, + _add_conv, + _add_dense, + _add_dense_nd, + _add_depthwise_conv, + _add_global_avgpool, + _add_mha, + _add_pq_activation, +) +from pquant.core.onnx_common import ( + add_initializer, + add_int64_array, + emit_getitem, + emit_squeeze, + emit_unsqueeze, + qdq_node, + quant_node, + save_model, + to_np, +) + +_ACTIVATION_OPS = {"relu": "Relu", "sigmoid": "Sigmoid", "tanh": "Tanh"} + + +def _resolve_mask_arg(mask, prefix, kind, tensor_to_onnx, initializers): + """Resolve an MHA mask call argument to an ONNX name (constant masks become initializers).""" + if mask is None: + return None + if tensor_to_onnx is not None and id(mask) in tensor_to_onnx: + return tensor_to_onnx[id(mask)] + return add_initializer(initializers, f"{prefix}_{kind}_const", np.asarray(to_np(mask))) + + +def _call_arguments(layer): + """The recorded call arguments of the layer's first inbound node.""" + return layer._inbound_nodes[0].arguments if layer._inbound_nodes else None + + +def _input_rank(layer): + """Rank (batch dim included) of the layer's first symbolic input tensor.""" + tensors = layer._inbound_nodes[0].input_tensors + tensor = tensors[0] if isinstance(tensors, (list, tuple)) else tensors + return len(tensor.shape) + + +def _add_mha_layer(layer, prefix, input_onnx_names, nodes, initializers, quant_fn, use_qonnx, store_int, tensor_to_onnx): + if len(input_onnx_names) >= 3: + q_in, k_in, v_in = input_onnx_names[:3] + elif len(input_onnx_names) == 2: + q_in, k_in, v_in = input_onnx_names[0], input_onnx_names[1], input_onnx_names[1] + else: + q_in = k_in = v_in = input_onnx_names[0] + + arguments = _call_arguments(layer) + kwargs = arguments.kwargs if arguments else {} + kpm = _resolve_mask_arg(kwargs.get("key_padding_mask"), prefix, "kpm", tensor_to_onnx, initializers) + attn_mask = _resolve_mask_arg(kwargs.get("attn_mask"), prefix, "attn_mask", tensor_to_onnx, initializers) + return _add_mha( + layer, + prefix, + q_in, + k_in, + v_in, + nodes, + initializers, + quant_fn, + use_qonnx, + store_int, + key_padding_mask=kpm, + attn_mask=attn_mask, + ) + + +def _add_getitem_op(layer, prefix, current, nodes, initializers): + """keras.ops GetItem operation recorded by ``x[...]`` KerasTensor syntax.""" + arguments = _call_arguments(layer) + spec = arguments.args[1] if len(arguments.args) > 1 else arguments.kwargs["key"] + rank = len(arguments.args[0].shape) + return emit_getitem(prefix, current, spec, rank, nodes, initializers) + + +def _add_expand_dims_op(layer, prefix, current, nodes, initializers): + """keras.ops.expand_dims operation; the axis is stored on the op.""" + rank = len(_call_arguments(layer).args[0].shape) + return emit_unsqueeze(prefix, current, [int(layer.axis) % (rank + 1)], nodes, initializers) + + +def _add_squeeze_op(layer, prefix, current, nodes, initializers): + """keras.ops.squeeze operation; axis=None squeezes every size-1 axis + (the batch axis is None in the symbolic shape, so it is never squeezed).""" + in_shape = _call_arguments(layer).args[0].shape + axis = layer.axis + if axis is None: + axes = [i for i, s in enumerate(in_shape) if s == 1] + else: + axis = axis if isinstance(axis, (list, tuple)) else (axis,) + axes = [a for a in (int(a) % len(in_shape) for a in axis) if in_shape[a] == 1] + return emit_squeeze(prefix, current, axes, nodes, initializers) + + +def _add_standard_activation(layer, prefix, current, nodes): + """keras.layers.ReLU or keras.layers.Activation with a supported activation.""" + activation = ( + layer.activation.__name__ + if isinstance(layer, keras.layers.Activation) and callable(layer.activation) + else getattr(layer, "activation", "relu") + ) + act_name = activation if isinstance(activation, str) else "relu" + op_type = next((op for key, op in _ACTIVATION_OPS.items() if key in act_name.lower()), None) + if op_type is None: + raise TypeError(f"Unsupported Activation for ONNX export: {act_name!r}") + out = f"{prefix}_act" + nodes.append(oh.make_node(op_type, inputs=[current], outputs=[out])) + return out + + +def _emit_layer( + layer, + prefix, + current, + nodes, + initializers, + quant_fn, + use_qonnx, + store_integer_weights, + input_onnx_names=None, + tensor_to_onnx=None, +): + """Emit ONNX nodes for a single Keras layer. Returns the ONNX output name.""" + + if isinstance(layer, PQMultiheadAttention): + return _add_mha_layer( + layer, prefix, input_onnx_names, nodes, initializers, quant_fn, use_qonnx, store_integer_weights, tensor_to_onnx + ) + + if isinstance(layer, PQActivation): + return _add_pq_activation(layer, prefix, current, nodes, initializers, quant_fn) + + if isinstance(layer, PQDense): + # Gemm only accepts rank-2 inputs; higher ranks (e.g. [batch, seq, dim]) go through MatMul + Add. + add_fn = _add_dense_nd if _input_rank(layer) > 2 else _add_dense + return add_fn(layer, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights) + + if isinstance(layer, PQDepthwiseConv2d): + return _add_depthwise_conv(layer, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights) + + if isinstance(layer, (PQConv2d, PQConv1d)): + ndim = 2 if isinstance(layer, PQConv2d) else 1 + return _add_conv( + layer, + prefix, + current, + nodes, + initializers, + ndim=ndim, + quant_fn=quant_fn, + use_qonnx=use_qonnx, + store_integer_weights=store_integer_weights, + ) + + if isinstance(layer, PQBatchNormalization): + return _add_batchnorm(layer, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights) + + if type(layer).__name__ == "GetItem": + return _add_getitem_op(layer, prefix, current, nodes, initializers) + + if type(layer).__name__ == "ExpandDims": + return _add_expand_dims_op(layer, prefix, current, nodes, initializers) + + if type(layer).__name__ == "Squeeze": + return _add_squeeze_op(layer, prefix, current, nodes, initializers) + + if isinstance(layer, (keras.layers.ReLU, keras.layers.Activation)): + return _add_standard_activation(layer, prefix, current, nodes) + + if isinstance(layer, keras.layers.Flatten): + out = f"{prefix}_flatten" + nodes.append(oh.make_node("Flatten", inputs=[current], outputs=[out], axis=1)) + return out + + if isinstance(layer, keras.layers.Reshape): + full_shape = [-1] + list(layer.target_shape) # -1 keeps the batch dim dynamic + shape_name = add_int64_array(initializers, f"{prefix}_shape", full_shape) + out = f"{prefix}_reshape" + nodes.append(oh.make_node("Reshape", inputs=[current, shape_name], outputs=[out])) + return out + + if isinstance(layer, keras.layers.Add): + assert input_onnx_names is not None and len(input_onnx_names) == 2 + out = f"{prefix}_add" + nodes.append(oh.make_node("Add", inputs=input_onnx_names, outputs=[out])) + return out + + if isinstance(layer, keras.layers.Multiply): + assert input_onnx_names is not None and len(input_onnx_names) == 2 + out = f"{prefix}_mul" + nodes.append(oh.make_node("Mul", inputs=input_onnx_names, outputs=[out])) + return out + + if isinstance(layer, keras.layers.Concatenate): + assert input_onnx_names is not None + out = f"{prefix}_concat" + # Negative axes are fine: ONNX Concat supports them. + nodes.append(oh.make_node("Concat", inputs=input_onnx_names, outputs=[out], axis=layer.axis)) + return out + + if isinstance(layer, (keras.layers.AveragePooling2D, keras.layers.AveragePooling1D)): + ndim = 2 if isinstance(layer, keras.layers.AveragePooling2D) else 1 + return _add_avgpool(layer, prefix, current, nodes, initializers, ndim=ndim, quant_fn=quant_fn) + + if isinstance(layer, (keras.layers.GlobalAveragePooling2D, keras.layers.GlobalAveragePooling1D)): + ndim = 2 if isinstance(layer, keras.layers.GlobalAveragePooling2D) else 1 + return _add_global_avgpool(layer, prefix, current, nodes, ndim=ndim) + + if isinstance(layer, keras.layers.Dropout): + return current # identity at inference + + raise TypeError(f"Unsupported Keras layer type for ONNX export: {type(layer).__name__!r}") + + +def _build_tensor_onnx_map(model): + """Seed the KerasTensor-id → ONNX-name map with the model inputs.""" + return {id(inp): name for inp, name in zip(model.inputs, _model_input_names(model))} + + +def _model_input_names(model): + if len(model.inputs) == 1: + return ["input"] + return [f"input_{i}" for i in range(len(model.inputs))] + + +def _inbound_input_names(layer, tensor_to_onnx): + """Return the list of ONNX input names for this layer based on its inbound node.""" + if not layer._inbound_nodes: + return [] + input_tensors = layer._inbound_nodes[0].input_tensors + if not isinstance(input_tensors, (list, tuple)): + input_tensors = [input_tensors] + result = [] + for t in input_tensors: + if id(t) not in tensor_to_onnx: + raise RuntimeError( + f"Layer {layer.name!r}: input tensor not found in tensor_to_onnx map. " + "Ensure model.layers is in topological order." + ) + result.append(tensor_to_onnx[id(t)]) + return result + + +def _register_layer_output(layer, onnx_name, tensor_to_onnx): + if not layer._inbound_nodes: + return + out_tensors = layer._inbound_nodes[0].output_tensors + if not isinstance(out_tensors, (list, tuple)): + out_tensors = [out_tensors] + if isinstance(onnx_name, (list, tuple)): + for tensor, name in zip(out_tensors, onnx_name): + tensor_to_onnx[id(tensor)] = name + else: + tensor_to_onnx[id(out_tensors[0])] = onnx_name + + +def convert_to_onnx( + model: keras.Model, + input_shape: tuple, + output_path: str = "model.onnx", + opset: int = 13, + use_qonnx: bool = False, + store_integer_weights: bool = False, + include_clip: bool = True, + batch_size: int | None = None, +) -> onnx.ModelProto: + """ + Convert a Keras functional model of PQuant layers to ONNX or QONNX. + + The model must have apply_final_compression() called on all PQ layers + before calling this function. Only inference-mode semantics are exported. + + Args: + model: Trained keras.Model. Must be a functional model + (built with the Keras functional API or subclassed + models whose layers are accessible via model.layers). + input_shape: Shape of a single sample excluding batch, e.g. (3, 32, 32). + For channels_last Conv models use e.g. (32, 32, 3). + output_path: Where to save the .onnx file. + opset: ONNX opset version (≥13 required for per-channel + DequantizeLinear). + use_qonnx: Emit QONNX Quant custom nodes if True. + store_integer_weights: Store weight initializers as int8/uint8 + + DequantizeLinear instead of float32 (ignored when + use_qonnx=True). + include_clip: Prepend a Clip node before each QuantizeLinear when + True (default). Set to False to emit bare + QuantizeLinear+DequantizeLinear pairs — safe when + values are guaranteed in-range at inference time since + QuantizeLinear saturates naturally. Ignored when + use_qonnx=True. + batch_size: If not None, fix the batch dimension of all graph + inputs and outputs to this value. If None (default), + the batch dimension is left dynamic. + + Returns: + The constructed onnx.ModelProto. + """ + quant_fn = quant_node if use_qonnx else functools.partial(qdq_node, include_clip=include_clip) + + onnx_nodes: list[onnx.NodeProto] = [] + initializers: list[onnx.TensorProto] = [] + tensor_to_onnx = _build_tensor_onnx_map(model) + last_output_name: str = "" + + for layer in getattr(model, "operations", None) or model.layers: + if isinstance(layer, keras.layers.InputLayer): + continue + + input_onnx_names = _inbound_input_names(layer, tensor_to_onnx) + if not input_onnx_names: + continue + + prefix = layer.name.replace("/", "_").replace(":", "_") + output_name = _emit_layer( + layer, + prefix, + input_onnx_names[0], + onnx_nodes, + initializers, + quant_fn, + use_qonnx, + store_integer_weights, + input_onnx_names=input_onnx_names, + tensor_to_onnx=tensor_to_onnx, + ) + + _register_layer_output(layer, output_name, tensor_to_onnx) + last_output_name = output_name[0] if isinstance(output_name, tuple) else output_name + + input_names = _model_input_names(model) + if len(model.inputs) == 1: + input_shapes = [tuple(input_shape)] + else: + input_shapes = [tuple(t.shape[1:]) for t in model.inputs] + np_dtypes = [np.dtype(str(t.dtype)) for t in model.inputs] + tp_dtypes = [_keras_dtype_to_tp(t.dtype) for t in model.inputs] + + dummies = [np.zeros((1, *shp), dtype=dt) for shp, dt in zip(input_shapes, np_dtypes)] + dummy_out = model(dummies[0] if len(dummies) == 1 else dummies, training=False) + dummy_out_np = np.array(ops.convert_to_numpy(dummy_out)) + + batch_dim = batch_size # None → dynamic, int → fixed + input_vis = [ + oh.make_tensor_value_info(name, tp, [batch_dim, *shp]) for name, shp, tp in zip(input_names, input_shapes, tp_dtypes) + ] + output_vi = oh.make_tensor_value_info(last_output_name, TensorProto.FLOAT, [batch_dim] + list(dummy_out_np.shape[1:])) + + graph = oh.make_graph( + nodes=onnx_nodes, + name="pquant_keras_onnx", + inputs=input_vis, + outputs=[output_vi], + initializer=initializers, + ) + model_proto = save_model(graph, output_path, opset, use_qonnx=use_qonnx) + logging.info("Saved %s Keras model → %s", "QONNX" if use_qonnx else "ONNX (QDQ)", output_path) + return model_proto diff --git a/src/pquant/core/keras/onnx/helpers.py b/src/pquant/core/keras/onnx/helpers.py new file mode 100644 index 0000000..9f43c0f --- /dev/null +++ b/src/pquant/core/keras/onnx/helpers.py @@ -0,0 +1,60 @@ +"""Keras-specific utilities for the PQuant Keras → ONNX converter. + +The backend-agnostic node emitters live in ``pquant.core.onnx_common``. +""" + +import keras +from onnx import TensorProto + + +def _keras_dtype_to_tp(dtype): + """Map a Keras/numpy dtype string to an ONNX TensorProto dtype (default float32).""" + return { + "float32": TensorProto.FLOAT, + "float64": TensorProto.DOUBLE, + "float16": TensorProto.FLOAT16, + "bool": TensorProto.BOOL, + "int64": TensorProto.INT64, + "int32": TensorProto.INT32, + }.get(str(dtype), TensorProto.FLOAT) + + +def _channels_last(layer): + return getattr(layer, "data_format", keras.config.image_data_format()) == "channels_last" + + +def _nchw_perms(ndim): + """Permutations between the Keras channels_last and ONNX channels_first layouts.""" + if ndim == 2: + return [0, 3, 1, 2], [0, 2, 3, 1] + return [0, 2, 1], [0, 2, 1] + + +def _bn_transpose_info(layer): + """ + Return (need_transpose, perm_fwd, perm_bwd) for a BatchNormalization layer. + + ONNX BN always normalises on axis 1 (NCHW; true in every opset), so + channels_last inputs need Transpose nodes around the BN op. + """ + axis = getattr(layer, "axis", 1) + stored = getattr(layer, "input_shape", None) + ndim = len(stored) if stored is not None else len(layer.input.shape) + eff_axis = axis if axis >= 0 else (ndim + axis) + + if eff_axis == 1 or ndim <= 2: + # channels already at position 1, or 2-D input — no transpose needed + return False, None, None + + if ndim == 4 and eff_axis == 3: + return True, [0, 3, 1, 2], [0, 2, 3, 1] + + if ndim == 3 and eff_axis == 2: + return True, [0, 2, 1], [0, 2, 1] + + # Fallback: general permutation that moves eff_axis to position 1 + perm_fwd = [0, eff_axis] + [i for i in range(1, ndim) if i != eff_axis] + perm_bwd = [0] * ndim + for i, p in enumerate(perm_fwd): + perm_bwd[p] = i + return True, perm_fwd, perm_bwd diff --git a/src/pquant/core/keras/onnx/layer_builders.py b/src/pquant/core/keras/onnx/layer_builders.py new file mode 100644 index 0000000..5e07932 --- /dev/null +++ b/src/pquant/core/keras/onnx/layer_builders.py @@ -0,0 +1,295 @@ +"""Per-layer ONNX graph builders (Dense/Conv/BN/Pool/Activation/MHA) for the PQuant Keras converter.""" + +import numpy as np +import onnx.helper as oh + +from pquant.core.keras.layers import PQBatchNormalization +from pquant.core.keras.onnx.helpers import ( + _bn_transpose_info, + _channels_last, + _nchw_perms, +) +from pquant.core.onnx_common import ( + add_float_scalar, + add_initializer, + add_transpose, + conv_padding_attrs, + emit_mha_core, + emit_param, + maybe_quant_input, + maybe_quant_output, + to_list, + to_np, +) + + +def _add_dense(layer, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights): + current = maybe_quant_input(layer, prefix, current, nodes, initializers, quant_fn) + + kernel_np = to_np(layer._kernel).T # [in, out] → [out, in] for Gemm (transB=1) + q_weight = emit_param( + prefix, "weight", kernel_np, layer.weight_quantizer, nodes, initializers, use_qonnx, store_integer_weights + ) + gemm_inputs = [current, q_weight] + + if layer._bias is not None: + gemm_inputs.append( + emit_param( + prefix, + "bias", + to_np(layer._bias), + layer.bias_quantizer, + nodes, + initializers, + use_qonnx, + store_integer_weights, + ) + ) + + gemm_out = f"{prefix}_gemm" + nodes.append(oh.make_node("Gemm", inputs=gemm_inputs, outputs=[gemm_out], transB=1)) + + return maybe_quant_output(layer, prefix, gemm_out, nodes, initializers, quant_fn) + + +def _add_dense_nd(layer, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights): + """Dense layer as MatMul + Add, for inputs of rank > 2 (Gemm only takes rank-2).""" + current = maybe_quant_input(layer, prefix, current, nodes, initializers, quant_fn) + + kernel_np = to_np(layer._kernel).T # [out, in] + if use_qonnx or store_integer_weights: + # Quantized/int-stored weight is emitted in native [out, in] layout, then transposed. + q_weight_native = emit_param( + prefix, "weight", kernel_np, layer.weight_quantizer, nodes, initializers, use_qonnx, store_integer_weights + ) + q_weight = f"{prefix}_weight_t" + nodes.append(oh.make_node("Transpose", inputs=[q_weight_native], outputs=[q_weight], perm=[1, 0])) + else: + q_weight = f"{prefix}_weight_t" + add_initializer(initializers, q_weight, kernel_np.T) # pre-transposed [in, out] + + mm_out = f"{prefix}_mm" + nodes.append(oh.make_node("MatMul", inputs=[current, q_weight], outputs=[mm_out])) + current = mm_out + + if layer._bias is not None: + q_bias = emit_param( + prefix, "bias", to_np(layer._bias), layer.bias_quantizer, nodes, initializers, use_qonnx, store_integer_weights + ) + add_out = f"{prefix}_bias_add" + nodes.append(oh.make_node("Add", inputs=[current, q_bias], outputs=[add_out])) + current = add_out + + return maybe_quant_output(layer, prefix, current, nodes, initializers, quant_fn) + + +def _add_conv_node(layer, prefix, conv_inputs, groups, ndim, nodes): + """Emit the Conv node shared by the regular and depthwise builders.""" + auto_pad, pads = conv_padding_attrs(layer.padding, ndim) + conv_attrs = dict( + kernel_shape=to_list(layer.kernel_size, ndim), + strides=to_list(layer.strides, ndim), + dilations=to_list(layer.dilation_rate, ndim), + group=groups, + auto_pad=auto_pad, + ) + if pads is not None: + conv_attrs["pads"] = pads + + conv_out = f"{prefix}_conv" + nodes.append(oh.make_node("Conv", inputs=conv_inputs, outputs=[conv_out], **conv_attrs)) + return conv_out + + +def _add_conv_common(layer, prefix, current, kernel_onnx, groups, ndim, nodes, initializers, quant_fn, use_qonnx, store_int): + """Shared body of the conv builders: layout transposes, param emission, Conv, quantization.""" + is_channels_last = _channels_last(layer) + if is_channels_last: + perm_to_nchw, perm_to_nhwx = _nchw_perms(ndim) + current = add_transpose(f"{prefix}_pre", current, perm_to_nchw, nodes) + + current = maybe_quant_input(layer, prefix, current, nodes, initializers, quant_fn) + + q_weight = emit_param(prefix, "weight", kernel_onnx, layer.weight_quantizer, nodes, initializers, use_qonnx, store_int) + conv_inputs = [current, q_weight] + if layer._bias is not None: + conv_inputs.append( + emit_param(prefix, "bias", to_np(layer._bias), layer.bias_quantizer, nodes, initializers, use_qonnx, store_int) + ) + + current = _add_conv_node(layer, prefix, conv_inputs, groups, ndim, nodes) + current = maybe_quant_output(layer, prefix, current, nodes, initializers, quant_fn) + + if is_channels_last: + current = add_transpose(f"{prefix}_post", current, perm_to_nhwx, nodes) + return current + + +def _add_conv(layer, prefix, current, nodes, initializers, ndim, quant_fn, use_qonnx, store_integer_weights): + kernel_np = to_np(layer._kernel) + # Transpose kernel from Keras HWIO to ONNX OIHW + if ndim == 2: + kernel_onnx = np.transpose(kernel_np, (3, 2, 0, 1)) # [kH,kW,in,out] → [out,in,kH,kW] + else: + kernel_onnx = np.transpose(kernel_np, (2, 1, 0)) # [kL,in,out] → [out,in,kL] + groups = getattr(layer, "groups", 1) + return _add_conv_common( + layer, prefix, current, kernel_onnx, groups, ndim, nodes, initializers, quant_fn, use_qonnx, store_integer_weights + ) + + +def _add_depthwise_conv(layer, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights): + kernel_np = to_np(layer._kernel) # [kH, kW, in, depth_mult] + in_ch, depth_mult = kernel_np.shape[2], kernel_np.shape[3] + # ONNX depthwise = Conv with groups=in and weight [in*depth_mult, 1, kH, kW] + kernel_onnx = np.transpose(kernel_np, (2, 3, 0, 1)).reshape(in_ch * depth_mult, 1, *kernel_np.shape[:2]) + return _add_conv_common( + layer, prefix, current, kernel_onnx, in_ch, 2, nodes, initializers, quant_fn, use_qonnx, store_integer_weights + ) + + +def _add_batchnorm(layer, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights): + """PQBatchNormalization (also handles plain keras BatchNormalization, + but _emit_layer currently only dispatches the PQ variant here).""" + need_transpose, perm_to_nchw, perm_to_nhwx = _bn_transpose_info(layer) + if need_transpose: + current = add_transpose(f"{prefix}_pre", current, perm_to_nchw, nodes) + + current = maybe_quant_input(layer, prefix, current, nodes, initializers, quant_fn) + + n_ch = to_np(layer.moving_mean).shape[0] + gamma_np = to_np(layer.gamma) if layer.gamma is not None else np.ones(n_ch, dtype=np.float32) # scale=False + beta_np = to_np(layer.beta) if layer.beta is not None else np.zeros(n_ch, dtype=np.float32) # center=False + + # Plain (non-PQ) BatchNormalization has no quantizers; emit plain float parameters. + is_pq = isinstance(layer, PQBatchNormalization) + use_qonnx = use_qonnx and is_pq + store_integer_weights = store_integer_weights and is_pq + weight_quantizer = layer.weight_quantizer if is_pq else None + bias_quantizer = layer.bias_quantizer if is_pq else None + + q_gamma = emit_param(prefix, "gamma", gamma_np, weight_quantizer, nodes, initializers, use_qonnx, store_integer_weights) + q_beta = emit_param(prefix, "beta", beta_np, bias_quantizer, nodes, initializers, use_qonnx, store_integer_weights) + mean_name = add_initializer(initializers, f"{prefix}_running_mean", to_np(layer.moving_mean)) + var_name = add_initializer(initializers, f"{prefix}_running_var", to_np(layer.moving_variance)) + + bn_out = f"{prefix}_bn" + nodes.append( + oh.make_node( + "BatchNormalization", + inputs=[current, q_gamma, q_beta, mean_name, var_name], + outputs=[bn_out], + epsilon=float(layer.epsilon), + ) + ) + current = bn_out + + if need_transpose: + current = add_transpose(f"{prefix}_post", current, perm_to_nhwx, nodes) + return current + + +def _add_mha( + layer, + prefix, + q_input, + k_input, + v_input, + nodes, + initializers, + quant_fn, + use_qonnx, + store_integer_weights, + key_padding_mask=None, + attn_mask=None, +): + # Q / K / V projections: (B, L, E) → (B, L, E) + q_proj_out = _add_dense_nd( + layer.q_proj, f"{prefix}_q_proj", q_input, nodes, initializers, quant_fn, use_qonnx, store_integer_weights + ) + k_proj_out = _add_dense_nd( + layer.k_proj, f"{prefix}_k_proj", k_input, nodes, initializers, quant_fn, use_qonnx, store_integer_weights + ) + v_proj_out = _add_dense_nd( + layer.v_proj, f"{prefix}_v_proj", v_input, nodes, initializers, quant_fn, use_qonnx, store_integer_weights + ) + + context, avg_attn = emit_mha_core( + layer, prefix, q_proj_out, k_proj_out, v_proj_out, nodes, initializers, quant_fn, key_padding_mask, attn_mask + ) + + # Output projection: (B, T, E) → (B, T, E) + out = _add_dense_nd( + layer.out_proj, f"{prefix}_out_proj", context, nodes, initializers, quant_fn, use_qonnx, store_integer_weights + ) + return out, avg_attn + + +def _add_avgpool(layer, prefix, current, nodes, initializers, ndim, quant_fn): + is_channels_last = _channels_last(layer) + if is_channels_last: + perm_to_nchw, perm_to_nhwx = _nchw_perms(ndim) + current = add_transpose(f"{prefix}_pre", current, perm_to_nchw, nodes) + + current = maybe_quant_input(layer, prefix, current, nodes, initializers, quant_fn) + + pool_out = f"{prefix}_pool" + nodes.append( + oh.make_node( + "AveragePool", + inputs=[current], + outputs=[pool_out], + kernel_shape=to_list(layer.pool_size, ndim), + strides=to_list(layer.strides, ndim), + pads=[0] * (ndim * 2), + count_include_pad=0, + ) + ) + current = maybe_quant_output(layer, prefix, pool_out, nodes, initializers, quant_fn) + + if is_channels_last: + current = add_transpose(f"{prefix}_post", current, perm_to_nhwx, nodes) + return current + + +def _add_global_avgpool(layer, prefix, current, nodes, ndim): + is_channels_last = _channels_last(layer) + if is_channels_last: + perm_to_nchw, _ = _nchw_perms(ndim) + current = add_transpose(f"{prefix}_pre", current, perm_to_nchw, nodes) + + pool_out = f"{prefix}_global_pool" + nodes.append(oh.make_node("GlobalAveragePool", inputs=[current], outputs=[pool_out])) + current = pool_out + + if is_channels_last: + flatten_name = f"{prefix}_flatten" + nodes.append(oh.make_node("Flatten", inputs=[pool_out], outputs=[flatten_name], axis=1)) + current = flatten_name + return current + + +def _add_pq_activation(layer, prefix, current, nodes, initializers, quant_fn): + current = maybe_quant_input(layer, prefix, current, nodes, initializers, quant_fn) + + if layer.use_multiplier and layer.activation_name == "relu" and hasattr(layer, "multiplier"): + multiplier = float(np.array(layer.multiplier).ravel()[0]) + scale_name = add_float_scalar(initializers, f"{prefix}_mul_scale", 2.0 ** round(multiplier)) + scaled_out = f"{prefix}_scaled" + nodes.append(oh.make_node("Mul", inputs=[current, scale_name], outputs=[scaled_out])) + current = scaled_out + + activation = layer.activation_name + act_out = f"{prefix}_act" + if activation == "relu": + nodes.append(oh.make_node("Relu", inputs=[current], outputs=[act_out])) + elif activation == "tanh": + nodes.append(oh.make_node("Tanh", inputs=[current], outputs=[act_out])) + elif activation == "hard_tanh": + cmin_name = add_float_scalar(initializers, f"{prefix}_htanh_min", -1.0) + cmax_name = add_float_scalar(initializers, f"{prefix}_htanh_max", 1.0) + nodes.append(oh.make_node("Clip", inputs=[current, cmin_name, cmax_name], outputs=[act_out])) + else: + raise TypeError(f"PQActivation: unsupported activation {activation!r} for ONNX export") + + return maybe_quant_output(layer, prefix, act_out, nodes, initializers, quant_fn) diff --git a/src/pquant/core/onnx_common.py b/src/pquant/core/onnx_common.py new file mode 100644 index 0000000..9f36b75 --- /dev/null +++ b/src/pquant/core/onnx_common.py @@ -0,0 +1,519 @@ +""" +Backend-agnostic ONNX node emitters shared by the PQuant Keras and torch +ONNX converters. + +Fixed-point (k, i, f) mapping +------------------------------ +QONNX: + scale = 2^(-f) + zero_point = 0 + bit_width = k + i + f + signed = int(k) + +Standard ONNX (QDQ): + scale = 2^(-f) + zero_point = 0 (int8 signed, uint8 unsigned) + clip range = [-2^i, 2^i - 2^(-f)] signed + = [0, 2^i - 2^(-f)] unsigned + Rounding is always nearest-even (QuantizeLinear behaviour). + +All quantization parameters (k, i, f) are accepted as anything ``to_np`` can +convert: torch tensors, Keras/TF tensors, numpy arrays, or Python scalars. +""" + +import numpy as np +import onnx +import onnx.helper as oh +import onnx.numpy_helper as onh + +ROUND_MODE_MAP = { + "TRN": "FLOOR", + "RND": "ROUND", + "RND_CONV": "ROUND", + "TRN_ZERO": "TRUNCATE", + "RND_ZERO": "ROUND", + "RND_MIN_INF": "FLOOR", + "RND_INF": "ROUND", +} + + +def to_np(tensor): + """Convert a torch/Keras/TF tensor (or scalar) to a float32 numpy array.""" + if hasattr(tensor, "detach"): # torch tensor, possibly on GPU + tensor = tensor.detach().cpu() + return np.asarray(tensor, dtype=np.float32) + + +def add_initializer(initializers, name, array): + """Register a constant tensor and return its name.""" + initializers.append(onh.from_array(array, name=name)) + return name + + +def add_float_scalar(initializers, name, value): + return add_initializer(initializers, name, np.array(value, dtype=np.float32)) + + +def add_int64_array(initializers, name, values): + return add_initializer(initializers, name, np.array(values, dtype=np.int64)) + + +def add_transpose(name, input_name, perm, nodes): + """Emit a Transpose node and return the output name.""" + out = f"{name}_transpose_{''.join(str(p) for p in perm)}" + nodes.append(oh.make_node("Transpose", inputs=[input_name], outputs=[out], perm=list(perm))) + return out + + +def to_list(v, n): + """Normalize a scalar-or-sequence layer attribute (kernel/stride/...) to an n-length list.""" + return list(v) if hasattr(v, "__iter__") else [v] * n + + +def symmetric_pads(padding, ndim): + """Expand a symmetric padding spec to the ONNX [begin_0, ..., end_0, ...] form.""" + per_axis = to_list(padding, ndim) + return per_axis + per_axis + + +def conv_padding_attrs(padding, ndim): + """Map a Keras/torch conv padding spec to ONNX (auto_pad, pads) attributes.""" + if isinstance(padding, str): + return ("SAME_UPPER" if padding == "same" else "VALID"), None + return "NOTSET", symmetric_pads(padding, ndim) + + +def fixed_point_clip_range(signed, i_val, f_val, overflow_mode): + """Representable [min, max] of a fixed-point grid with i integer and f fractional bits.""" + clip_max = float(2.0**i_val - 2.0 ** (-f_val)) + if not signed: + return 0.0, clip_max + if overflow_mode == "SAT_SYM": + return -clip_max, clip_max # symmetric: excludes -2^i + return float(-(2.0**i_val)), clip_max # SAT: -2^i + + +def quant_node(name_prefix, input_name, rounding_mode, k, i, f, initializers, overflow_mode="SAT"): + """Build a QONNX Quant node. Returns ([node], output_name).""" + k_np, i_np, f_np = to_np(k), to_np(i), to_np(f) + k_val = int(k_np.ravel()[0]) + if f_np.size > 1: + # The Quant node holds a single scale/bit_width; widen to cover every element. + i_np = i_np.ravel().max() + f_np = f_np.ravel().min() + i_val = float(i_np) + f_val = float(f_np) + scale = float(2.0 ** (-f_val)) + bit_width = float(k_val + i_val + f_val) + qonnx_rnd = ROUND_MODE_MAP.get(rounding_mode, "ROUND") + # SAT_SYM excludes the most-negative value → QONNX narrow=1 + narrow = 1 if (k_val == 1 and overflow_mode == "SAT_SYM") else 0 + + scale_name = add_float_scalar(initializers, f"{name_prefix}_scale", scale) + zp_name = add_float_scalar(initializers, f"{name_prefix}_zero_point", 0.0) + bw_name = add_float_scalar(initializers, f"{name_prefix}_bit_width", bit_width) + out_name = f"{name_prefix}_quantized" + + node = oh.make_node( + op_type="Quant", + inputs=[input_name, scale_name, zp_name, bw_name], + outputs=[out_name], + domain="qonnx.custom_op.general", + signed=k_val, + narrow=narrow, + rounding_mode=qonnx_rnd, + ) + return [node], out_name + + +def qdq_node( + name_prefix, input_name, rounding_mode, k, i, f, initializers, overflow_mode="SAT", include_clip=True +): # noqa: ARG001 (rounding_mode kept for the shared quant_fn signature) + """Build QuantizeLinear+DequantizeLinear nodes, optionally preceded by a Clip. + + Returns ([nodes], output_name). Set include_clip=False to skip the Clip node + (safe when values are guaranteed to be in-range at inference time). + """ + k_val = int(to_np(k).ravel()[0]) + i_val = float(to_np(i).ravel()[0]) + f_val = float(to_np(f).ravel()[0]) + signed = k_val == 1 + clip_min, clip_max = fixed_point_clip_range(signed, i_val, f_val, overflow_mode) + zp_val = np.int8(0) if signed else np.uint8(0) + + scale_name = add_float_scalar(initializers, f"{name_prefix}_scale", 2.0 ** (-f_val)) + zp_name = add_initializer(initializers, f"{name_prefix}_zero_point", np.array(zp_val)) + quantized_name = f"{name_prefix}_quantized" + out_name = f"{name_prefix}_dequantized" + + nodes = [] + quantize_input = input_name + if include_clip: + clip_min_name = add_float_scalar(initializers, f"{name_prefix}_clip_min", clip_min) + clip_max_name = add_float_scalar(initializers, f"{name_prefix}_clip_max", clip_max) + quantize_input = f"{name_prefix}_clipped" + nodes.append(oh.make_node("Clip", inputs=[input_name, clip_min_name, clip_max_name], outputs=[quantize_input])) + nodes.append(oh.make_node("QuantizeLinear", inputs=[quantize_input, scale_name, zp_name], outputs=[quantized_name])) + nodes.append(oh.make_node("DequantizeLinear", inputs=[quantized_name, scale_name, zp_name], outputs=[out_name])) + return nodes, out_name + + +def per_channel_scale(f_np, out_channels): + """Return the (out_channels,) scale vector, or None when f varies within a channel.""" + if f_np.size % out_channels != 0: + return None + f_per_channel = f_np.reshape(out_channels, -1) + if not np.allclose(f_per_channel, f_per_channel[:, :1]): + return None + return (2.0 ** (-f_per_channel[:, 0])).astype(np.float32) + + +def int_weight_node(name_prefix, weight_np, k, f, initializers): + """ + Store a weight tensor as int8/uint8 + DequantizeLinear. + + weight_np must already be in ONNX layout and on the fixed-point grid + (guaranteed after apply_final_compression). Converts by dividing by the + scale and casting — no re-rounding needed. + + Granularity handling: + - per-tensor (f has one element): single scale, standard DequantizeLinear. + - per-channel (f constant within each output channel): 1D scale with axis=0. + - per-weight (f fully per-element): ONNX has no per-weight quantization; + falls back to float32 storage (no DequantizeLinear node). + + Returns ([nodes], output_name). + """ + k_val = int(to_np(k).ravel()[0]) + dtype = np.int8 if k_val == 1 else np.uint8 + out_channels = weight_np.shape[0] + f_np = to_np(f) + + if f_np.size == 1: + scale_np = np.array(2.0 ** (-float(f_np.ravel()[0])), dtype=np.float32) + int_weights = np.round(weight_np / float(scale_np)).astype(dtype) + per_channel = False + else: + scale_np = per_channel_scale(f_np, out_channels) + if scale_np is None: + float_name = add_initializer(initializers, f"{name_prefix}_float", weight_np) + return [], float_name + broadcast_scale = scale_np.reshape((out_channels,) + (1,) * (weight_np.ndim - 1)) + int_weights = np.round(weight_np / broadcast_scale).astype(dtype) + per_channel = True + + int_name = add_initializer(initializers, f"{name_prefix}_int", int_weights) + scale_name = add_initializer(initializers, f"{name_prefix}_dq_scale", scale_np) + zp_np = np.zeros(out_channels, dtype=dtype) if per_channel else np.array(dtype(0)) + zp_name = add_initializer(initializers, f"{name_prefix}_dq_zp", zp_np) + + out_name = f"{name_prefix}_dequantized" + node_kwargs = {"axis": 0} if per_channel else {} + node = oh.make_node("DequantizeLinear", inputs=[int_name, scale_name, zp_name], outputs=[out_name], **node_kwargs) + return [node], out_name + + +def emit_param(prefix, name, arr, quantizer, nodes, initializers, use_qonnx, store_integer_weights): + """Emit the ONNX value for a learnable parameter (kernel/bias/gamma/beta) and return its name. + + arr must already be in ONNX layout (e.g. OIHW for conv kernels). + """ + if use_qonnx: + fp_name = add_initializer(initializers, f"{prefix}_{name}_fp", arr) + k, i, f = quantizer.get_quantization_bits() + q_nodes, out = quant_node( + f"{prefix}_{name}", fp_name, quantizer.round_mode, k, i, f, initializers, overflow_mode=quantizer.overflow + ) + nodes.extend(q_nodes) + return out + if store_integer_weights: + k, _, f = quantizer.get_quantization_bits() + q_nodes, out = int_weight_node(f"{prefix}_{name}", arr, k, f, initializers) + nodes.extend(q_nodes) + return out + return add_initializer(initializers, f"{prefix}_{name}", arr) + + +def apply_quantizer(quantizer, prefix, current, nodes, initializers, quant_fn): + """Emit quant_fn nodes for one Quantizer and return the new tensor name.""" + k, i, f = quantizer.get_quantization_bits() + new_nodes, out = quant_fn(prefix, current, quantizer.round_mode, k, i, f, initializers, overflow_mode=quantizer.overflow) + nodes.extend(new_nodes) + return out + + +def maybe_quant_input(layer, prefix, current, nodes, initializers, quant_fn): + # input_quantizer is created conditionally, so guard it; the bool flags are always present. + if getattr(layer, "input_quantizer", None) is not None and layer.quantize_input and layer.enable_quantization: + current = apply_quantizer(layer.input_quantizer, f"{prefix}_in", current, nodes, initializers, quant_fn) + return current + + +def maybe_quant_output(layer, prefix, current, nodes, initializers, quant_fn): + if getattr(layer, "output_quantizer", None) is not None and layer.quantize_output and layer.enable_quantization: + current = apply_quantizer(layer.output_quantizer, f"{prefix}_out", current, nodes, initializers, quant_fn) + return current + + +def emit_getitem(prefix, input_name, spec, rank, nodes, initializers): + """Translate a constant Python indexing spec into ONNX Slice (+ Squeeze).""" + if not isinstance(spec, tuple): + spec = (spec,) + n_ellipsis = sum(1 for s in spec if s is Ellipsis) + if n_ellipsis > 1: + raise TypeError("indexing with more than one Ellipsis is not supported in ONNX export") + if n_ellipsis: + pos = spec.index(Ellipsis) + fill = rank - (len(spec) - 1) + spec = spec[:pos] + (slice(None),) * fill + spec[pos + 1 :] + if len(spec) > rank: + raise TypeError(f"indexing spec has {len(spec)} dims but tensor rank is {rank}") + + int64_max = np.iinfo(np.int64).max + starts, ends, axes, steps, squeeze_axes = [], [], [], [], [] + for axis, s in enumerate(spec): + if isinstance(s, slice): + if s.start is None and s.stop is None and s.step in (None, 1): + continue # full slice: no-op on this axis + step = 1 if s.step is None else int(s.step) + if step < 1: + raise TypeError("slice steps < 1 are not supported in ONNX export") + starts.append(0 if s.start is None else int(s.start)) + ends.append(int64_max if s.stop is None else int(s.stop)) + axes.append(axis) + steps.append(step) + elif isinstance(s, int): + starts.append(s) + ends.append(int64_max if s == -1 else s + 1) + axes.append(axis) + steps.append(1) + squeeze_axes.append(axis) + else: + raise TypeError(f"unsupported index element {s!r} for ONNX export (constant int/slice/Ellipsis only)") + + current = input_name + if axes: + slice_inputs = [current] + for part, vals in (("starts", starts), ("ends", ends), ("axes", axes), ("steps", steps)): + slice_inputs.append(add_int64_array(initializers, f"{prefix}_slice_{part}", vals)) + current = f"{prefix}_slice" + nodes.append(oh.make_node("Slice", inputs=slice_inputs, outputs=[current])) + if squeeze_axes: + current = emit_squeeze(prefix, current, squeeze_axes, nodes, initializers) + return current + + +def emit_squeeze(prefix, input_name, axes, nodes, initializers): + """Emit an ONNX Squeeze removing the given size-1 axes (no-op if axes is empty). + + Squeeze takes axes as an input tensor from opset 13 on (the converter minimum). + """ + if not axes: + return input_name + ax_name = add_int64_array(initializers, f"{prefix}_squeeze_axes", sorted(axes)) + out = f"{prefix}_squeeze" + nodes.append(oh.make_node("Squeeze", inputs=[input_name, ax_name], outputs=[out])) + return out + + +def emit_unsqueeze(prefix, input_name, axes, nodes, initializers): + """Emit an ONNX Unsqueeze inserting size-1 dims at the given axes.""" + ax_name = add_int64_array(initializers, f"{prefix}_unsqueeze_axes", axes) + out = f"{prefix}_unsqueeze" + nodes.append(oh.make_node("Unsqueeze", inputs=[input_name, ax_name], outputs=[out])) + return out + + +def add_quantized_softmax(sm, prefix, current, nodes, initializers, quant_fn, kpm_mask=None): + """Emit the PQuant quantized-softmax decomposition over the last axis.""" + enable = sm.enable_quantization + scaler = float(sm.input_scaler) + stable = bool(sm.stable) + + # 1) Softmax input quantizer. + if sm.quantize_input and enable: + current = apply_quantizer(sm.input_quantizer, f"{prefix}_sm_in_q", current, nodes, initializers, quant_fn) + + # 2) Stable max-subtract over the last axis (ReduceMax keeps axes as an attribute). + if stable: + m_name = f"{prefix}_sm_max" + nodes.append(oh.make_node("ReduceMax", inputs=[current], outputs=[m_name], axes=[-1], keepdims=1)) + exp_in = f"{prefix}_sm_sub" + nodes.append(oh.make_node("Sub", inputs=[m_name, current], outputs=[exp_in])) + else: + exp_in = current + + # 3) Quantized exp table: optional input QDQ, Exp of (-scaler * x) for the + # stable branch (+scaler otherwise), optional output QDQ. + exp_table = sm.exp_table + if exp_table.quantize_input and enable: + exp_in = apply_quantizer(exp_table.input_quantizer, f"{prefix}_sm_exp_in_q", exp_in, nodes, initializers, quant_fn) + coeff = -scaler if stable else scaler + exp_arg = exp_in + if coeff != 1.0: + coeff_name = add_float_scalar(initializers, f"{prefix}_sm_exp_coeff", coeff) + exp_arg = f"{prefix}_sm_exp_arg" + nodes.append(oh.make_node("Mul", inputs=[exp_in, coeff_name], outputs=[exp_arg])) + numerator = f"{prefix}_sm_exp" + nodes.append(oh.make_node("Exp", inputs=[exp_arg], outputs=[numerator])) + if exp_table.quantize_output and enable: + numerator = apply_quantizer( + exp_table.output_quantizer, f"{prefix}_sm_exp_out_q", numerator, nodes, initializers, quant_fn + ) + + # 3b) Optional key-padding mask: zero the exp-numerator at masked positions. + if kpm_mask is not None: + kpm_f = f"{prefix}_sm_mask_f" + nodes.append(oh.make_node("Cast", inputs=[kpm_mask], outputs=[kpm_f], to=onnx.TensorProto.FLOAT)) + masked = f"{prefix}_sm_masked" + nodes.append(oh.make_node("Mul", inputs=[kpm_f, numerator], outputs=[masked])) + numerator = masked + + # 4) Sum over the last axis (ReduceSum takes axes as an input from opset 13). + sum_axes = add_int64_array(initializers, f"{prefix}_sm_sum_axes", [-1]) + sums = f"{prefix}_sm_sum" + nodes.append(oh.make_node("ReduceSum", inputs=[numerator, sum_axes], outputs=[sums], keepdims=1)) + + # 5) Quantized reciprocal table: input QDQ, 1/(x+eps), output QDQ. + inv_table = sm.inv_table + inv_in = sums + if inv_table.quantize_input and enable: + inv_in = apply_quantizer(inv_table.input_quantizer, f"{prefix}_sm_inv_in_q", inv_in, nodes, initializers, quant_fn) + eps_name = add_float_scalar(initializers, f"{prefix}_sm_eps", float(sm.epsilon)) + inv_add = f"{prefix}_sm_inv_add" + nodes.append(oh.make_node("Add", inputs=[inv_in, eps_name], outputs=[inv_add])) + divisor = f"{prefix}_sm_inv" + nodes.append(oh.make_node("Reciprocal", inputs=[inv_add], outputs=[divisor])) + if inv_table.quantize_output and enable: + divisor = apply_quantizer( + inv_table.output_quantizer, f"{prefix}_sm_inv_out_q", divisor, nodes, initializers, quant_fn + ) + + # 6) Multiply numerator by reciprocal. + out = f"{prefix}_sm_out" + nodes.append(oh.make_node("Mul", inputs=[numerator, divisor], outputs=[out])) + current = out + + # 7) Softmax output quantizer. + if sm.quantize_output and enable: + current = apply_quantizer(sm.output_quantizer, f"{prefix}_sm_out_q", current, nodes, initializers, quant_fn) + return current + + +def split_heads(x_name, pfx, num_heads, head_dim, nodes, initializers): + """(B, L, E) → (B, num_heads, L, head_dim) using runtime Shape ops so B and L stay dynamic.""" + shape_out = f"{pfx}_shape" + nodes.append(oh.make_node("Shape", inputs=[x_name], outputs=[shape_out])) + + idx0 = add_int64_array(initializers, f"{pfx}_gi0", 0) + idx1 = add_int64_array(initializers, f"{pfx}_gi1", 1) + ax0 = add_int64_array(initializers, f"{pfx}_ax0", [0]) + heads_1d = add_int64_array(initializers, f"{pfx}_H_1d", [num_heads]) + head_dim_1d = add_int64_array(initializers, f"{pfx}_hd_1d", [head_dim]) + + batch_scalar = f"{pfx}_b_sc" + length_scalar = f"{pfx}_l_sc" + batch_1d = f"{pfx}_b_1d" + length_1d = f"{pfx}_l_1d" + shape_4d = f"{pfx}_shape4d" + reshaped = f"{pfx}_reshaped" + transposed = f"{pfx}_transposed" + + nodes.append(oh.make_node("Gather", inputs=[shape_out, idx0], outputs=[batch_scalar])) + nodes.append(oh.make_node("Gather", inputs=[shape_out, idx1], outputs=[length_scalar])) + nodes.append(oh.make_node("Unsqueeze", inputs=[batch_scalar, ax0], outputs=[batch_1d])) + nodes.append(oh.make_node("Unsqueeze", inputs=[length_scalar, ax0], outputs=[length_1d])) + nodes.append(oh.make_node("Concat", inputs=[batch_1d, length_1d, heads_1d, head_dim_1d], outputs=[shape_4d], axis=0)) + nodes.append(oh.make_node("Reshape", inputs=[x_name, shape_4d], outputs=[reshaped])) + # (B, L, H, head_dim) → (B, H, L, head_dim) + nodes.append(oh.make_node("Transpose", inputs=[reshaped], outputs=[transposed], perm=[0, 2, 1, 3])) + return transposed + + +def merge_heads(x_name, pfx, embed_dim, nodes, initializers): + """(B, H, T, head_dim) → (B, T, embed_dim): the inverse of split_heads.""" + transposed = f"{pfx}_t" # (B, T, H, head_dim) + shape_out = f"{pfx}_shape" + nodes.append(oh.make_node("Transpose", inputs=[x_name], outputs=[transposed], perm=[0, 2, 1, 3])) + nodes.append(oh.make_node("Shape", inputs=[transposed], outputs=[shape_out])) + + idx0 = add_int64_array(initializers, f"{pfx}_gi0", 0) + idx1 = add_int64_array(initializers, f"{pfx}_gi1", 1) + ax0 = add_int64_array(initializers, f"{pfx}_ax0", [0]) + embed_1d = add_int64_array(initializers, f"{pfx}_E_1d", [embed_dim]) + + batch_scalar = f"{pfx}_b_sc" + length_scalar = f"{pfx}_t_sc" + batch_1d = f"{pfx}_b_1d" + length_1d = f"{pfx}_t_1d" + shape_3d = f"{pfx}_shape3d" + merged = f"{pfx}_merged" + + nodes.append(oh.make_node("Gather", inputs=[shape_out, idx0], outputs=[batch_scalar])) + nodes.append(oh.make_node("Gather", inputs=[shape_out, idx1], outputs=[length_scalar])) + nodes.append(oh.make_node("Unsqueeze", inputs=[batch_scalar, ax0], outputs=[batch_1d])) + nodes.append(oh.make_node("Unsqueeze", inputs=[length_scalar, ax0], outputs=[length_1d])) + nodes.append(oh.make_node("Concat", inputs=[batch_1d, length_1d, embed_1d], outputs=[shape_3d], axis=0)) + nodes.append(oh.make_node("Reshape", inputs=[transposed, shape_3d], outputs=[merged])) + return merged + + +def emit_mha_core( + mha, prefix, q_proj_out, k_proj_out, v_proj_out, nodes, initializers, quant_fn, key_padding_mask, attn_mask +): + """Scaled-dot-product attention between projected Q/K/V, quantized softmax included. + + Returns (context_name, avg_attn_name): the merged (B, T, E) context ready for + the output projection, and the attention weights averaged over heads. + """ + q_heads = split_heads(q_proj_out, f"{prefix}_q", mha.num_heads, mha.head_dim, nodes, initializers) + k_heads = split_heads(k_proj_out, f"{prefix}_k", mha.num_heads, mha.head_dim, nodes, initializers) + v_heads = split_heads(v_proj_out, f"{prefix}_v", mha.num_heads, mha.head_dim, nodes, initializers) + + k_transposed = f"{prefix}_k_T" + nodes.append(oh.make_node("Transpose", inputs=[k_heads], outputs=[k_transposed], perm=[0, 1, 3, 2])) + + raw_scores = f"{prefix}_scores_raw" + scaled_scores = f"{prefix}_scores_scaled" + scale_name = add_float_scalar(initializers, f"{prefix}_attn_scale", float(mha.scale)) + nodes.append(oh.make_node("MatMul", inputs=[q_heads, k_transposed], outputs=[raw_scores])) + nodes.append(oh.make_node("Mul", inputs=[raw_scores, scale_name], outputs=[scaled_scores])) + current = scaled_scores + + if attn_mask is not None: + masked_scores = f"{prefix}_scores_masked" + nodes.append(oh.make_node("Add", inputs=[current, attn_mask], outputs=[masked_scores])) + current = masked_scores + + kpm_mult = None + if key_padding_mask is not None: + kpm_not = f"{prefix}_kpm_not" + nodes.append(oh.make_node("Not", inputs=[key_padding_mask], outputs=[kpm_not])) + kpm_axes = add_int64_array(initializers, f"{prefix}_kpm_axes", [1, 2]) + kpm_mult = f"{prefix}_kpm_mask" # (B, 1, 1, S) bool, cast to float inside the softmax + nodes.append(oh.make_node("Unsqueeze", inputs=[kpm_not, kpm_axes], outputs=[kpm_mult])) + + attn_weights = add_quantized_softmax( + mha.softmax, f"{prefix}_attn", current, nodes, initializers, quant_fn, kpm_mask=kpm_mult + ) + + context = f"{prefix}_ctx_raw" + nodes.append(oh.make_node("MatMul", inputs=[attn_weights, v_heads], outputs=[context])) + merged = merge_heads(context, f"{prefix}_ctx", mha.embed_dim, nodes, initializers) + + # Average attention weights over heads: (B, H, T, S) → (B, T, S) + avg_attn = f"{prefix}_avg_attn_weights" + nodes.append(oh.make_node("ReduceMean", inputs=[attn_weights], outputs=[avg_attn], axes=[1], keepdims=0)) + return merged, avg_attn + + +def save_model(graph, output_path, opset, use_qonnx=False, ir_version=6): + """Assemble the ModelProto for a finished graph, validate it, and save it.""" + opset_imports = [oh.make_opsetid("", opset)] + if use_qonnx: + opset_imports.append(oh.make_opsetid("qonnx.custom_op.general", 1)) + model_proto = oh.make_model(graph, opset_imports=opset_imports) + model_proto.ir_version = ir_version + onnx.checker.check_model(model_proto) + onnx.save(model_proto, output_path) + return model_proto diff --git a/src/pquant/core/torch/convert_to_onnx.py b/src/pquant/core/torch/convert_to_onnx.py deleted file mode 100644 index a61aeaa..0000000 --- a/src/pquant/core/torch/convert_to_onnx.py +++ /dev/null @@ -1,1859 +0,0 @@ -""" -Convert a PQuant model to ONNX or QONNX format. - -Pass ``use_qonnx=True`` to emit QONNX ``Quant`` custom nodes (requires the -qonnx runtime). Pass ``use_qonnx=False`` (default) to emit standard -``Clip + QuantizeLinear + DequantizeLinear`` nodes runnable with plain -onnxruntime. - -Fixed-point (k, i, f) mapping ------------------------------- -QONNX: - scale = 2^(-f) - zero_point = 0 - bit_width = k + i + f - signed = int(k) - -Standard ONNX (QDQ): - scale = 2^(-f) - zero_point = 0 (int8 signed, uint8 unsigned) - clip range = [-2^i, 2^i - 2^(-f)] signed - = [0, 2^i - 2^(-f)] unsigned - Rounding is always nearest-even (QuantizeLinear behaviour). - Weights are stored as plain float32 initializers — after - apply_final_compression() they are already on the fixed-point grid. -""" - -import functools -import logging -import operator as _operator -import os - -import numpy as np -import onnx -import onnx.helper as oh -import onnx.numpy_helper as onh -import torch -import torch.fx as _fx -import torch.nn as nn -import torch.nn.functional as _F -from onnx import TensorProto - -os.environ["KERAS_BACKEND"] = "torch" # must be set before any keras/pquant import - -from pquant.core.torch.activations import PQActivation # noqa: E402 -from pquant.core.torch.layers import ( # noqa: E402 - PQAvgPool1d, - PQAvgPool2d, - PQBatchNorm1d, - PQBatchNorm2d, - PQConv1d, - PQConv2d, - PQDense, - PQLayerNorm, - PQMultiheadAttention, -) -from pquant.core.torch.quantizer import Quantizer # noqa: E402 - -# --------------------------------------------------------------------------- -# QONNX Quant node -# --------------------------------------------------------------------------- - -ROUND_MODE_MAP = { - "TRN": "FLOOR", - "RND": "ROUND", - "RND_CONV": "ROUND", - "TRN_ZERO": "TRUNCATE", - "RND_ZERO": "ROUND", - "RND_MIN_INF": "FLOOR", - "RND_INF": "ROUND", -} - - -def _quant_node(name_prefix, input_name, rounding_mode, k, i, f, initializers, overflow_mode="SAT"): - """Build a QONNX Quant node. Returns ([node], output_name). - - QONNX Quant is per-tensor only. If i/f are per-channel or per-weight tensors - (non-scalar), collapse to the broadest range: min(f) / max(i) ensures no channel - overflows at the cost of slightly coarser quantization for small-value channels. - """ - k_val = int(k.item()) - if hasattr(f, "numel") and f.numel() > 1: - i = i.reshape(-1).max() - f = f.reshape(-1).min() - i_val = float(i.item()) - f_val = float(f.item()) - scale = float(2.0 ** (-f_val)) - bit_width = float(k_val + i_val + f_val) - qonnx_rnd = ROUND_MODE_MAP.get(rounding_mode, "ROUND") - narrow = 1 if (k_val == 1 and overflow_mode == "SAT_SYM") else 0 - - scale_name = f"{name_prefix}_scale" - zp_name = f"{name_prefix}_zero_point" - bw_name = f"{name_prefix}_bit_width" - out_name = f"{name_prefix}_quantized" - - initializers.append(onh.from_array(np.array(scale, dtype=np.float32), name=scale_name)) - initializers.append(onh.from_array(np.array(0.0, dtype=np.float32), name=zp_name)) - initializers.append(onh.from_array(np.array(bit_width, dtype=np.float32), name=bw_name)) - - node = oh.make_node( - op_type="Quant", - inputs=[input_name, scale_name, zp_name, bw_name], - outputs=[out_name], - domain="qonnx.custom_op.general", - signed=k_val, - narrow=narrow, - rounding_mode=qonnx_rnd, - ) - return [node], out_name - - -# --------------------------------------------------------------------------- -# Standard ONNX QDQ triple -# --------------------------------------------------------------------------- - - -def _qdq_node( - name_prefix, input_name, rounding_mode, k, i, f, initializers, overflow_mode="SAT", include_clip=True -): # noqa: ARG001 - """Build QuantizeLinear+DequantizeLinear nodes, optionally preceded by a Clip. - - Returns ([nodes], output_name). Set include_clip=False to skip the Clip node - (safe when values are guaranteed in-range at inference time, since - QuantizeLinear saturates naturally). - """ - k_val = int(k.item()) - i_val = float(i.item()) - f_val = float(f.item()) - scale = float(2.0 ** (-f_val)) - signed = k_val == 1 - - clip_max = float(2.0**i_val - 2.0 ** (-f_val)) - if not signed: - clip_min = 0.0 - elif overflow_mode == "SAT_SYM": - clip_min = -clip_max - else: - clip_min = float(-(2.0**i_val)) - zp_val = np.int8(0) if signed else np.uint8(0) - - scale_name = f"{name_prefix}_scale" - zp_name = f"{name_prefix}_zero_point" - quantized_name = f"{name_prefix}_quantized" - out_name = f"{name_prefix}_dequantized" - - initializers += [ - onh.from_array(np.array(scale, dtype=np.float32), name=scale_name), - onh.from_array(np.array(zp_val), name=zp_name), - ] - - if include_clip: - clip_min_name = f"{name_prefix}_clip_min" - clip_max_name = f"{name_prefix}_clip_max" - clipped_name = f"{name_prefix}_clipped" - initializers += [ - onh.from_array(np.array(clip_min, dtype=np.float32), name=clip_min_name), - onh.from_array(np.array(clip_max, dtype=np.float32), name=clip_max_name), - ] - nodes = [ - oh.make_node("Clip", inputs=[input_name, clip_min_name, clip_max_name], outputs=[clipped_name]), - oh.make_node("QuantizeLinear", inputs=[clipped_name, scale_name, zp_name], outputs=[quantized_name]), - ] - else: - nodes = [ - oh.make_node("QuantizeLinear", inputs=[input_name, scale_name, zp_name], outputs=[quantized_name]), - ] - nodes.append(oh.make_node("DequantizeLinear", inputs=[quantized_name, scale_name, zp_name], outputs=[out_name])) - return nodes, out_name - - -# --------------------------------------------------------------------------- -# helpers -# --------------------------------------------------------------------------- - - -def _int_weight_node(name_prefix, weight_np, k, i, f, initializers): # noqa: ARG001 (i unused) - """ - Store a weight tensor as int8/uint8 + DequantizeLinear. - - weight_np must already be on the fixed-point grid (guaranteed after - apply_final_compression). Converts by dividing by the scale and casting — - no re-rounding needed. - - Granularity handling: - - per-tensor (f is scalar): single scale, standard DequantizeLinear. - - per-channel (f has shape [out, 1, ...]): 1D scale with axis=0. - All weights in a channel share the same f so the conversion is exact. - - per-weight (f is fully per-element): ONNX has no per-weight quantization; - falls back to float32 storage (no DequantizeLinear node). - - Returns ([node], output_name). - """ - k_val = int(k.item()) if hasattr(k, "item") else int(k) - dtype = np.int8 if k_val == 1 else np.uint8 - out_channels = weight_np.shape[0] - out_name = f"{name_prefix}_dequantized" - - f_t = f.detach().cpu() if hasattr(f, "detach") else torch.as_tensor(f) - - if f_t.numel() == 1: - # per-tensor - scale_np = np.array(float(2.0 ** (-f_t.item())), dtype=np.float32) - int_weights = np.round(weight_np / float(scale_np)).astype(dtype) - per_channel = False - else: - f_np = f_t.float().numpy().reshape(out_channels, -1) - if np.allclose(f_np, f_np[:, :1]): - # per-channel: all elements within an output channel share one f - f_1d = f_np[:, 0] - scale_np = (2.0 ** (-f_1d)).astype(np.float32) - bcast = scale_np.reshape((out_channels,) + (1,) * (weight_np.ndim - 1)) - int_weights = np.round(weight_np / bcast).astype(dtype) - per_channel = True - else: - # per-weight: ONNX cannot represent this; store as float32 - float_name = f"{name_prefix}_float" - initializers.append(onh.from_array(weight_np, name=float_name)) - return [], float_name - - int_name = f"{name_prefix}_int" - scale_name = f"{name_prefix}_dq_scale" - zp_name = f"{name_prefix}_dq_zp" - - zp_np = np.zeros(out_channels if per_channel else 1, dtype=dtype) - initializers += [ - onh.from_array(int_weights, name=int_name), - onh.from_array(scale_np, name=scale_name), - onh.from_array(zp_np if per_channel else np.array(dtype(0)), name=zp_name), - ] - node_kwargs = {"axis": 0} if per_channel else {} - node = oh.make_node("DequantizeLinear", inputs=[int_name, scale_name, zp_name], outputs=[out_name], **node_kwargs) - return [node], out_name - - -def _torch_padding_to_onnx(padding, ndim): - if isinstance(padding, int): - padding = (padding,) * ndim - return list(padding) + list(padding) - - -def _maybe_quant_input(module, prefix, current, nodes, initializers, quant_fn): - if ( - getattr(module, "input_quantizer", None) is not None - and getattr(module, "quantize_input", True) - and getattr(module, "enable_quantization", True) - ): - q = module.input_quantizer - k, i, f = q.get_quantization_bits() - new_nodes, current = quant_fn( - f"{prefix}_in", current, q.round_mode, k, i, f, initializers, overflow_mode=getattr(q, "overflow", "SAT") - ) - nodes.extend(new_nodes) - return current - - -def _maybe_quant_output(module, prefix, current, nodes, initializers, quant_fn): - if ( - getattr(module, "output_quantizer", None) is not None - and getattr(module, "quantize_output", False) - and getattr(module, "enable_quantization", True) - ): - q = module.output_quantizer - k, i, f = q.get_quantization_bits() - new_nodes, current = quant_fn( - f"{prefix}_out", current, q.round_mode, k, i, f, initializers, overflow_mode=getattr(q, "overflow", "SAT") - ) - nodes.extend(new_nodes) - return current - - -# --------------------------------------------------------------------------- -# per-layer graph builders -# --------------------------------------------------------------------------- - - -def _add_dense_integer(module, prefix, current, nodes, initializers): - """Dense layer using MatMulInteger for true integer arithmetic. - - Flow: - float → Clip+QuantizeLinear → int8 ─┐ - ├─ MatMulInteger → int32 - int8 weights (pre-transposed) ───────┘ - → Add int32 bias - → DequantizeLinear(scale = s_x * s_w) → float - - The inner product accumulates in int32; there is no float Gemm. - A single DequantizeLinear at the end converts back to float for activations. - Per-channel weights use axis=1 on the output DequantizeLinear. - """ - if not (getattr(module, "input_quantizer", None) and getattr(module, "quantize_input", True)): - raise ValueError(f"{prefix}: integer_ops requires quantize_input=True on the layer") - - # --- Input: Clip + QuantizeLinear → int8 (stop before DequantizeLinear) --- - k_x, i_x, f_x = module.input_quantizer.get_quantization_bits() - k_x_val = int(k_x.item()) - i_x_val = float(i_x.item()) - f_x_val = float(f_x.item()) - s_x = float(2.0 ** (-f_x_val)) - signed_x = k_x_val == 1 - - clip_min_x = float(-(2.0**i_x_val)) if signed_x else 0.0 - clip_max_x = float(2.0**i_x_val - 2.0 ** (-f_x_val)) - zp_x_np = np.int8(0) if signed_x else np.uint8(0) - - clip_min_name = f"{prefix}_in_clip_min" - clip_max_name = f"{prefix}_in_clip_max" - scale_x_name = f"{prefix}_in_scale" - zp_x_name = f"{prefix}_in_zp" - x_int_name = f"{prefix}_in_int" - - initializers += [ - onh.from_array(np.array(clip_min_x, dtype=np.float32), name=clip_min_name), - onh.from_array(np.array(clip_max_x, dtype=np.float32), name=clip_max_name), - onh.from_array(np.array(s_x, dtype=np.float32), name=scale_x_name), - onh.from_array(np.array(zp_x_np), name=zp_x_name), - ] - nodes += [ - oh.make_node("Clip", inputs=[current, clip_min_name, clip_max_name], outputs=[f"{prefix}_in_clipped"]), - oh.make_node("QuantizeLinear", inputs=[f"{prefix}_in_clipped", scale_x_name, zp_x_name], outputs=[x_int_name]), - ] - - # --- Weights: stored pre-transposed as int8 so MatMulInteger needs no Transpose --- - # PyTorch weight shape: [out, in]. MatMulInteger(A, B) = A @ B, so we need [in, out]. - weight_np = module._weight.detach().cpu().numpy().astype(np.float32) - k_w, _, f_w = module.weight_quantizer.get_quantization_bits() - k_w_val = int(k_w.item()) if hasattr(k_w, "item") else int(k_w) - dtype_w = np.int8 if k_w_val == 1 else np.uint8 - out_ch = weight_np.shape[0] - - f_w_t = f_w.detach().cpu() if hasattr(f_w, "detach") else torch.as_tensor(f_w) - if f_w_t.numel() == 1: - f_w_1d = np.array([float(f_w_t.item())]) - per_channel_w = False - else: - f_w_2d = f_w_t.float().numpy().reshape(out_ch, -1) - f_w_1d = f_w_2d.min(axis=1) # min f → max scale → covers all values - per_channel_w = True - - s_w_1d = (2.0 ** (-f_w_1d)).astype(np.float32) # shape [1] or [out] - bcast_s_w = s_w_1d.reshape((out_ch,) + (1,) * (weight_np.ndim - 1)) if per_channel_w else float(s_w_1d[0]) - # Transpose before storing so MatMulInteger can use it without a runtime Transpose node - int_weights_T = np.round(weight_np / bcast_s_w).astype(dtype_w).T # [in, out] - - zp_w_np = np.array(dtype_w(0)) # scalar zero-point; zero for symmetric quantization - w_int_name = f"{prefix}_weight_int" - w_zp_name = f"{prefix}_weight_zp" - initializers += [ - onh.from_array(int_weights_T, name=w_int_name), - onh.from_array(zp_w_np, name=w_zp_name), - ] - - # --- MatMulInteger([batch, in], [in, out]) → int32 [batch, out] --- - y_int_name = f"{prefix}_matmul_int" - nodes.append( - oh.make_node( - "MatMulInteger", - inputs=[x_int_name, w_int_name, zp_x_name, w_zp_name], - outputs=[y_int_name], - ) - ) - - # --- Bias added in int32 domain: bias_int[c] = round(bias[c] / (s_x * s_w[c])) --- - current_int32 = y_int_name - if module._bias is not None: - bias_np = module._bias.detach().cpu().numpy().astype(np.float32) - combined_s = s_x * s_w_1d # shape [1] or [out] - bias_int32 = np.round(bias_np / (combined_s if per_channel_w else float(combined_s[0]))).astype(np.int32) - bias_int_name = f"{prefix}_bias_int" - y_biased_name = f"{prefix}_matmul_biased" - initializers.append(onh.from_array(bias_int32, name=bias_int_name)) - nodes.append(oh.make_node("Add", inputs=[current_int32, bias_int_name], outputs=[y_biased_name])) - current_int32 = y_biased_name - - # --- DequantizeLinear: int32 → float32 using combined scale s_x * s_w --- - # Per-channel: axis=1 because the output tensor is [batch, out] and out is axis 1. - combined_scale_name = f"{prefix}_combined_scale" - combined_zp_name = f"{prefix}_combined_zp" - - if per_channel_w: - combined_scale_np = (s_x * s_w_1d).astype(np.float32) # [out] - combined_zp_np = np.zeros(out_ch, dtype=np.int32) - dql_kwargs = {"axis": 1} - else: - combined_scale_np = np.array(float(s_x * s_w_1d[0]), dtype=np.float32) - combined_zp_np = np.array(np.int32(0)) - dql_kwargs = {} - - initializers += [ - onh.from_array(combined_scale_np, name=combined_scale_name), - onh.from_array(combined_zp_np, name=combined_zp_name), - ] - y_float_name = f"{prefix}_dequantized" - nodes.append( - oh.make_node( - "DequantizeLinear", - inputs=[current_int32, combined_scale_name, combined_zp_name], - outputs=[y_float_name], - **dql_kwargs, - ) - ) - current = y_float_name - - # Optional output quantization (e.g. last layer with quantize_output=True) - current = _maybe_quant_output(module, prefix, current, nodes, initializers, _qdq_node) - return current - - -def _add_dense_nd(module, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights): - """Dense (linear) projection via MatMul, supporting input of any rank ≥ 2. - - Identical logic to _add_dense but emits ``MatMul(input, W_T)`` instead of - ``Gemm(input, W, transB=1)`` so it accepts (B, T, E) inputs (e.g. from MHA - projections) as well as the usual 2-D (batch, features) inputs. - Weight is stored pre-transposed as [in, out] to avoid a runtime Transpose node. - """ - current = _maybe_quant_input(module, prefix, current, nodes, initializers, quant_fn) - - weight_np = module._weight.detach().cpu().numpy().astype(np.float32) # [out, in] - if use_qonnx: - weight_fp_name = f"{prefix}_weight_fp" - initializers.append(onh.from_array(weight_np, name=weight_fp_name)) - k_w, i_w, f_w = module.weight_quantizer.get_quantization_bits() - w_nodes, q_weight_raw = _quant_node( - f"{prefix}_weight", - weight_fp_name, - module.weight_quantizer.round_mode, - k_w, - i_w, - f_w, - initializers, - overflow_mode=getattr(module.weight_quantizer, "overflow", "SAT"), - ) - nodes.extend(w_nodes) - q_weight_t = f"{prefix}_weight_T" - nodes.append(oh.make_node("Transpose", inputs=[q_weight_raw], outputs=[q_weight_t], perm=[1, 0])) - q_weight = q_weight_t - elif store_integer_weights: - k_w, i_w, f_w = module.weight_quantizer.get_quantization_bits() - w_nodes, q_weight_stored = _int_weight_node(f"{prefix}_weight", weight_np, k_w, i_w, f_w, initializers) - nodes.extend(w_nodes) - q_weight_t = f"{prefix}_weight_T" - nodes.append(oh.make_node("Transpose", inputs=[q_weight_stored], outputs=[q_weight_t], perm=[1, 0])) - q_weight = q_weight_t - else: - q_weight = f"{prefix}_weight_T" - initializers.append(onh.from_array(weight_np.T, name=q_weight)) # pre-transposed [in, out] - - matmul_out = f"{prefix}_matmul" - nodes.append(oh.make_node("MatMul", inputs=[current, q_weight], outputs=[matmul_out])) - current = matmul_out - - if module._bias is not None: - bias_np = module._bias.detach().cpu().numpy().astype(np.float32) - if use_qonnx: - bias_fp_name = f"{prefix}_bias_fp" - initializers.append(onh.from_array(bias_np, name=bias_fp_name)) - k_b, i_b, f_b = module.bias_quantizer.get_quantization_bits() - b_nodes, q_bias = _quant_node( - f"{prefix}_bias", - bias_fp_name, - module.bias_quantizer.round_mode, - k_b, - i_b, - f_b, - initializers, - overflow_mode=getattr(module.bias_quantizer, "overflow", "SAT"), - ) - nodes.extend(b_nodes) - elif store_integer_weights: - k_b, i_b, f_b = module.bias_quantizer.get_quantization_bits() - b_nodes, q_bias = _int_weight_node(f"{prefix}_bias", bias_np, k_b, i_b, f_b, initializers) - nodes.extend(b_nodes) - else: - q_bias = f"{prefix}_bias" - initializers.append(onh.from_array(bias_np, name=q_bias)) - biased_out = f"{prefix}_biased" - nodes.append(oh.make_node("Add", inputs=[matmul_out, q_bias], outputs=[biased_out])) - current = biased_out - - current = _maybe_quant_output(module, prefix, current, nodes, initializers, quant_fn) - return current - - -def _add_dense(module, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights, integer_ops=False): - if integer_ops and not use_qonnx: - return _add_dense_integer(module, prefix, current, nodes, initializers) - current = _maybe_quant_input(module, prefix, current, nodes, initializers, quant_fn) - - weight_np = module._weight.detach().cpu().numpy().astype(np.float32) - if use_qonnx: - weight_fp_name = f"{prefix}_weight_fp" - initializers.append(onh.from_array(weight_np, name=weight_fp_name)) - k_w, i_w, f_w = module.weight_quantizer.get_quantization_bits() - w_nodes, q_weight = _quant_node( - f"{prefix}_weight", - weight_fp_name, - module.weight_quantizer.round_mode, - k_w, - i_w, - f_w, - initializers, - overflow_mode=getattr(module.weight_quantizer, "overflow", "SAT"), - ) - nodes.extend(w_nodes) - elif store_integer_weights: - k_w, i_w, f_w = module.weight_quantizer.get_quantization_bits() - w_nodes, q_weight = _int_weight_node(f"{prefix}_weight", weight_np, k_w, i_w, f_w, initializers) - nodes.extend(w_nodes) - else: - q_weight = f"{prefix}_weight" - initializers.append(onh.from_array(weight_np, name=q_weight)) - - # Use Gemm with transB=1 — weight stays in its native [out, in] layout, - # no Transpose node needed. Bias (if any) is fused as the third Gemm input. - gemm_inputs = [current, q_weight] - - if module._bias is not None: - bias_np = module._bias.detach().cpu().numpy().astype(np.float32) - if use_qonnx: - bias_fp_name = f"{prefix}_bias_fp" - initializers.append(onh.from_array(bias_np, name=bias_fp_name)) - k_b, i_b, f_b = module.bias_quantizer.get_quantization_bits() - b_nodes, q_bias = _quant_node( - f"{prefix}_bias", - bias_fp_name, - module.bias_quantizer.round_mode, - k_b, - i_b, - f_b, - initializers, - overflow_mode=getattr(module.bias_quantizer, "overflow", "SAT"), - ) - nodes.extend(b_nodes) - elif store_integer_weights: - k_b, i_b, f_b = module.bias_quantizer.get_quantization_bits() - b_nodes, q_bias = _int_weight_node(f"{prefix}_bias", bias_np, k_b, i_b, f_b, initializers) - nodes.extend(b_nodes) - else: - q_bias = f"{prefix}_bias" - initializers.append(onh.from_array(bias_np, name=q_bias)) - gemm_inputs.append(q_bias) - - gemm_out = f"{prefix}_gemm" - nodes.append(oh.make_node("Gemm", inputs=gemm_inputs, outputs=[gemm_out], transB=1)) - current = gemm_out - - current = _maybe_quant_output(module, prefix, current, nodes, initializers, quant_fn) - return current - - -def _add_conv(module, prefix, current, nodes, initializers, ndim, quant_fn, use_qonnx, store_integer_weights): - current = _maybe_quant_input(module, prefix, current, nodes, initializers, quant_fn) - - weight_np = module._weight.detach().cpu().numpy().astype(np.float32) - if use_qonnx: - weight_fp_name = f"{prefix}_weight_fp" - initializers.append(onh.from_array(weight_np, name=weight_fp_name)) - k_w, i_w, f_w = module.weight_quantizer.get_quantization_bits() - w_nodes, q_weight = _quant_node( - f"{prefix}_weight", - weight_fp_name, - module.weight_quantizer.round_mode, - k_w, - i_w, - f_w, - initializers, - overflow_mode=getattr(module.weight_quantizer, "overflow", "SAT"), - ) - nodes.extend(w_nodes) - elif store_integer_weights: - k_w, i_w, f_w = module.weight_quantizer.get_quantization_bits() - w_nodes, q_weight = _int_weight_node(f"{prefix}_weight", weight_np, k_w, i_w, f_w, initializers) - nodes.extend(w_nodes) - else: - q_weight = f"{prefix}_weight" - initializers.append(onh.from_array(weight_np, name=q_weight)) - - conv_inputs = [current, q_weight] - - if module._bias is not None: - bias_np = module._bias.detach().cpu().numpy().astype(np.float32) - if use_qonnx: - bias_fp_name = f"{prefix}_bias_fp" - initializers.append(onh.from_array(bias_np, name=bias_fp_name)) - k_b, i_b, f_b = module.bias_quantizer.get_quantization_bits() - b_nodes, q_bias = _quant_node( - f"{prefix}_bias", - bias_fp_name, - module.bias_quantizer.round_mode, - k_b, - i_b, - f_b, - initializers, - overflow_mode=getattr(module.bias_quantizer, "overflow", "SAT"), - ) - nodes.extend(b_nodes) - elif store_integer_weights: - k_b, i_b, f_b = module.bias_quantizer.get_quantization_bits() - b_nodes, q_bias = _int_weight_node(f"{prefix}_bias", bias_np, k_b, i_b, f_b, initializers) - nodes.extend(b_nodes) - else: - q_bias = f"{prefix}_bias" - initializers.append(onh.from_array(bias_np, name=q_bias)) - conv_inputs.append(q_bias) - - padding = module.padding - if isinstance(padding, str): - auto_pad = "SAME_UPPER" if padding == "same" else "VALID" - pads = None - else: - auto_pad = "NOTSET" - pads = _torch_padding_to_onnx(padding, ndim) - - to_list = lambda v, n: list(v) if hasattr(v, "__iter__") else [v] * n # noqa: E731 - conv_attrs = dict( - kernel_shape=to_list(module.kernel_size, ndim), - strides=to_list(module.stride, ndim), - dilations=to_list(module.dilation, ndim), - group=module.groups, - auto_pad=auto_pad, - ) - if pads is not None: - conv_attrs["pads"] = pads - - conv_out = f"{prefix}_conv" - nodes.append(oh.make_node("Conv", inputs=conv_inputs, outputs=[conv_out], **conv_attrs)) - current = conv_out - - current = _maybe_quant_output(module, prefix, current, nodes, initializers, quant_fn) - return current - - -def _add_batchnorm(module, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights): - current = _maybe_quant_input(module, prefix, current, nodes, initializers, quant_fn) - - gamma_np = module._weight.detach().cpu().numpy().astype(np.float32) - beta_np = module._bias.detach().cpu().numpy().astype(np.float32) - - if use_qonnx: - gamma_fp_name = f"{prefix}_gamma_fp" - initializers.append(onh.from_array(gamma_np, name=gamma_fp_name)) - k_w, i_w, f_w = module.weight_quantizer.get_quantization_bits() - g_nodes, q_gamma = _quant_node( - f"{prefix}_gamma", - gamma_fp_name, - module.weight_quantizer.round_mode, - k_w, - i_w, - f_w, - initializers, - overflow_mode=getattr(module.weight_quantizer, "overflow", "SAT"), - ) - nodes.extend(g_nodes) - - beta_fp_name = f"{prefix}_beta_fp" - initializers.append(onh.from_array(beta_np, name=beta_fp_name)) - k_b, i_b, f_b = module.bias_quantizer.get_quantization_bits() - b_nodes, q_beta = _quant_node( - f"{prefix}_beta", - beta_fp_name, - module.bias_quantizer.round_mode, - k_b, - i_b, - f_b, - initializers, - overflow_mode=getattr(module.bias_quantizer, "overflow", "SAT"), - ) - nodes.extend(b_nodes) - elif store_integer_weights: - k_w, i_w, f_w = module.weight_quantizer.get_quantization_bits() - g_nodes, q_gamma = _int_weight_node(f"{prefix}_gamma", gamma_np, k_w, i_w, f_w, initializers) - nodes.extend(g_nodes) - k_b, i_b, f_b = module.bias_quantizer.get_quantization_bits() - b_nodes, q_beta = _int_weight_node(f"{prefix}_beta", beta_np, k_b, i_b, f_b, initializers) - nodes.extend(b_nodes) - else: - q_gamma = f"{prefix}_gamma" - q_beta = f"{prefix}_beta" - initializers.append(onh.from_array(gamma_np, name=q_gamma)) - initializers.append(onh.from_array(beta_np, name=q_beta)) - - mean_name = f"{prefix}_running_mean" - var_name = f"{prefix}_running_var" - initializers.append(onh.from_array(module.running_mean.detach().cpu().numpy().astype(np.float32), name=mean_name)) - initializers.append(onh.from_array(module.running_var.detach().cpu().numpy().astype(np.float32), name=var_name)) - - bn_out = f"{prefix}_bn" - nodes.append( - oh.make_node( - "BatchNormalization", - inputs=[current, q_gamma, q_beta, mean_name, var_name], - outputs=[bn_out], - epsilon=float(module.eps), - ) - ) - return bn_out - - -def _add_layernorm(module, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights): - """PQLayerNorm. Emits LayerNormalization (opset >= 17 required).""" - current = _maybe_quant_input(module, prefix, current, nodes, initializers, quant_fn) - - ns = ( - tuple(int(d) for d in module.normalized_shape) - if hasattr(module.normalized_shape, "__iter__") - else (int(module.normalized_shape),) - ) - axis = -len(ns) - - has_weight = module._weight is not None - has_bias = module._bias is not None - - gamma_np = module._weight.detach().cpu().numpy().astype(np.float32) if has_weight else np.ones(ns, dtype=np.float32) - beta_np = module._bias.detach().cpu().numpy().astype(np.float32) if has_bias else None - - if use_qonnx and has_weight: - gamma_fp_name = f"{prefix}_gamma_fp" - initializers.append(onh.from_array(gamma_np, name=gamma_fp_name)) - k_w, i_w, f_w = module.weight_quantizer.get_quantization_bits() - g_nodes, q_gamma = _quant_node( - f"{prefix}_gamma", - gamma_fp_name, - module.weight_quantizer.round_mode, - k_w, - i_w, - f_w, - initializers, - overflow_mode=getattr(module.weight_quantizer, "overflow", "SAT"), - ) - nodes.extend(g_nodes) - if has_bias: - beta_fp_name = f"{prefix}_beta_fp" - initializers.append(onh.from_array(beta_np, name=beta_fp_name)) - k_b, i_b, f_b = module.bias_quantizer.get_quantization_bits() - b_nodes, q_beta = _quant_node( - f"{prefix}_beta", - beta_fp_name, - module.bias_quantizer.round_mode, - k_b, - i_b, - f_b, - initializers, - overflow_mode=getattr(module.bias_quantizer, "overflow", "SAT"), - ) - nodes.extend(b_nodes) - elif store_integer_weights and has_weight: - k_w, i_w, f_w = module.weight_quantizer.get_quantization_bits() - g_nodes, q_gamma = _int_weight_node(f"{prefix}_gamma", gamma_np, k_w, i_w, f_w, initializers) - nodes.extend(g_nodes) - if has_bias: - k_b, i_b, f_b = module.bias_quantizer.get_quantization_bits() - b_nodes, q_beta = _int_weight_node(f"{prefix}_beta", beta_np, k_b, i_b, f_b, initializers) - nodes.extend(b_nodes) - else: - q_gamma = f"{prefix}_gamma" - initializers.append(onh.from_array(gamma_np, name=q_gamma)) - if has_bias: - q_beta = f"{prefix}_beta" - initializers.append(onh.from_array(beta_np, name=q_beta)) - - ln_inputs = [current, q_gamma] - if has_bias: - ln_inputs.append(q_beta) - ln_out = f"{prefix}_ln" - nodes.append( - oh.make_node( - "LayerNormalization", - inputs=ln_inputs, - outputs=[ln_out], - axis=axis, - epsilon=float(module.eps), - ) - ) - current = ln_out - current = _maybe_quant_output(module, prefix, current, nodes, initializers, quant_fn) - return current - - -def _add_avgpool(module, prefix, current, nodes, initializers, ndim, quant_fn): - current = _maybe_quant_input(module, prefix, current, nodes, initializers, quant_fn) - - to_list = lambda v, n: list(v) if hasattr(v, "__iter__") else [v] * n # noqa: E731 - pool_out = f"{prefix}_pool" - nodes.append( - oh.make_node( - "AveragePool", - inputs=[current], - outputs=[pool_out], - kernel_shape=to_list(module.kernel_size, ndim), - strides=to_list(module.stride, ndim), - pads=_torch_padding_to_onnx(module.padding, ndim), - ceil_mode=int(module.ceil_mode), - count_include_pad=int(module.count_include_pad), - ) - ) - current = pool_out - - current = _maybe_quant_output(module, prefix, current, nodes, initializers, quant_fn) - return current - - -# --------------------------------------------------------------------------- -# multi-head attention graph builder -# --------------------------------------------------------------------------- - - -def _add_mha(module, prefix, q_input, k_input, v_input, nodes, initializers, quant_fn, use_qonnx, store_integer_weights): - """Build ONNX nodes for PQMultiheadAttention. - - Decomposes multi-head attention into primitive ONNX ops: - - [optional transpose if not batch_first] - Q/K/V Gemm projections - Reshape (B, L, E) → (B, H, L, head_dim) + Transpose - MatMul(Q, K^T) * scale → optional Quant - Softmax → optional Quant - MatMul(attn_weights, V) → optional Quant - Transpose + Reshape (B, T, E) - out_proj Gemm - [optional transpose back if not batch_first] - - Returns (out_name, avg_attn_weights_name): the projected output and the - attention weights averaged over heads (B, T, S). Both names are valid ONNX - value names so downstream getitem(mha, 0) / getitem(mha, 1) work in the FX - converter. - - Note: if ``approximate_softmax=True`` the module uses a polynomial - approximation in PyTorch, but ONNX has no equivalent standard op — a plain - ``Softmax`` node is emitted instead. - """ - H = module.num_heads - head_dim = module.head_dim - E = module.embed_dim - scale_val = float(module.scale) - - # --- Optional transpose for seq-first inputs (T, B, E) → (B, T, E) --- - if not module.batch_first: - q_t = f"{prefix}_q_in_t" - k_t = f"{prefix}_k_in_t" - v_t = f"{prefix}_v_in_t" - nodes.append(oh.make_node("Transpose", inputs=[q_input], outputs=[q_t], perm=[1, 0, 2])) - nodes.append(oh.make_node("Transpose", inputs=[k_input], outputs=[k_t], perm=[1, 0, 2])) - nodes.append(oh.make_node("Transpose", inputs=[v_input], outputs=[v_t], perm=[1, 0, 2])) - q_input, k_input, v_input = q_t, k_t, v_t - - # --- Q / K / V projections: (B, L, E) → (B, L, E) via MatMul (input is rank-3) --- - q_proj_out = _add_dense_nd( - module.q_proj, f"{prefix}_q_proj", q_input, nodes, initializers, quant_fn, use_qonnx, store_integer_weights - ) - k_proj_out = _add_dense_nd( - module.k_proj, f"{prefix}_k_proj", k_input, nodes, initializers, quant_fn, use_qonnx, store_integer_weights - ) - v_proj_out = _add_dense_nd( - module.v_proj, f"{prefix}_v_proj", v_input, nodes, initializers, quant_fn, use_qonnx, store_integer_weights - ) - - # --- Helper: (B, L, E) → (B, H, L, head_dim) using dynamic shapes --- - def _split_heads(x_name, pfx): - shape_out = f"{pfx}_shape" - b_scalar = f"{pfx}_b_sc" - l_scalar = f"{pfx}_l_sc" - b_1d = f"{pfx}_b_1d" - l_1d = f"{pfx}_l_1d" - h_1d_const = f"{pfx}_H_1d" - hd_1d_const = f"{pfx}_hd_1d" - shape_4d = f"{pfx}_shape4d" - reshaped = f"{pfx}_reshaped" - transposed = f"{pfx}_transposed" - idx0 = f"{pfx}_gi0" - idx1 = f"{pfx}_gi1" - ax0 = f"{pfx}_ax0" - - nodes.append(oh.make_node("Shape", inputs=[x_name], outputs=[shape_out])) - initializers.extend( - [ - onh.from_array(np.array(0, dtype=np.int64), name=idx0), - onh.from_array(np.array(1, dtype=np.int64), name=idx1), - onh.from_array(np.array([0], dtype=np.int64), name=ax0), - onh.from_array(np.array([H], dtype=np.int64), name=h_1d_const), - onh.from_array(np.array([head_dim], dtype=np.int64), name=hd_1d_const), - ] - ) - nodes.append(oh.make_node("Gather", inputs=[shape_out, idx0], outputs=[b_scalar])) - nodes.append(oh.make_node("Gather", inputs=[shape_out, idx1], outputs=[l_scalar])) - nodes.append(oh.make_node("Unsqueeze", inputs=[b_scalar, ax0], outputs=[b_1d])) - nodes.append(oh.make_node("Unsqueeze", inputs=[l_scalar, ax0], outputs=[l_1d])) - nodes.append(oh.make_node("Concat", inputs=[b_1d, l_1d, h_1d_const, hd_1d_const], outputs=[shape_4d], axis=0)) - nodes.append(oh.make_node("Reshape", inputs=[x_name, shape_4d], outputs=[reshaped])) - # (B, L, H, head_dim) → (B, H, L, head_dim) - nodes.append(oh.make_node("Transpose", inputs=[reshaped], outputs=[transposed], perm=[0, 2, 1, 3])) - return transposed - - q_h = _split_heads(q_proj_out, f"{prefix}_q") - k_h = _split_heads(k_proj_out, f"{prefix}_k") - v_h = _split_heads(v_proj_out, f"{prefix}_v") - - # --- k^T: (B, H, S, head_dim) → (B, H, head_dim, S) --- - k_t_name = f"{prefix}_k_T" - nodes.append(oh.make_node("Transpose", inputs=[k_h], outputs=[k_t_name], perm=[0, 1, 3, 2])) - - # --- Scaled dot-product scores: (B, H, T, head_dim) @ (B, H, head_dim, S) → (B, H, T, S) --- - raw_scores = f"{prefix}_scores_raw" - scaled_scores = f"{prefix}_scores_scaled" - scale_cst = f"{prefix}_attn_scale" - nodes.append(oh.make_node("MatMul", inputs=[q_h, k_t_name], outputs=[raw_scores])) - initializers.append(onh.from_array(np.array(scale_val, dtype=np.float32), name=scale_cst)) - nodes.append(oh.make_node("Mul", inputs=[raw_scores, scale_cst], outputs=[scaled_scores])) - current = scaled_scores - - # --- Softmax input quantization (the MHA enables the softmax's input quantizer) --- - if module.softmax.quantize_input and getattr(module, "enable_quantization", True): - q = module.softmax.input_quantizer - k_q, i_q, f_q = q.get_quantization_bits() - q_nodes, current = quant_fn( - f"{prefix}_attn_score_q", - current, - q.round_mode, - k_q, - i_q, - f_q, - initializers, - overflow_mode=getattr(q, "overflow", "SAT"), - ) - nodes.extend(q_nodes) - - # --- Softmax (dim=-1); approximate_softmax falls back to standard Softmax in ONNX --- - attn_w_name = f"{prefix}_attn_weights" - nodes.append(oh.make_node("Softmax", inputs=[current], outputs=[attn_w_name], axis=-1)) - current = attn_w_name - - # --- Softmax output quantization (the MHA enables the softmax's output quantizer) --- - if module.softmax.quantize_output and getattr(module, "enable_quantization", True): - q = module.softmax.output_quantizer - k_q, i_q, f_q = q.get_quantization_bits() - q_nodes, current = quant_fn( - f"{prefix}_attn_weight_q", - current, - q.round_mode, - k_q, - i_q, - f_q, - initializers, - overflow_mode=getattr(q, "overflow", "SAT"), - ) - nodes.extend(q_nodes) - - # --- Context: (B, H, T, S) @ (B, H, S, head_dim) → (B, H, T, head_dim) --- - # No dedicated quantizer: out_proj's input quantizer (exported with the dense) covers it. - ctx_raw = f"{prefix}_ctx_raw" - nodes.append(oh.make_node("MatMul", inputs=[current, v_h], outputs=[ctx_raw])) - current_ctx = ctx_raw - - # --- Merge heads: (B, H, T, head_dim) → (B, T, E) using dynamic shapes --- - ctx_t = f"{prefix}_ctx_t" # after Transpose → (B, T, H, head_dim) - ctx_shape = f"{prefix}_ctx_shape" - ctx_b_sc = f"{prefix}_ctx_b_sc" - ctx_t_sc = f"{prefix}_ctx_t_sc" - ctx_b_1d = f"{prefix}_ctx_b_1d" - ctx_t_1d = f"{prefix}_ctx_t_1d" - ctx_E_1d = f"{prefix}_ctx_E_1d" - ctx_ax0 = f"{prefix}_ctx_ax0" - ctx_gi0 = f"{prefix}_ctx_gi0" - ctx_gi1 = f"{prefix}_ctx_gi1" - ctx_3d = f"{prefix}_ctx_shape3d" - ctx_merged = f"{prefix}_ctx_merged" - - nodes.append(oh.make_node("Transpose", inputs=[current_ctx], outputs=[ctx_t], perm=[0, 2, 1, 3])) - nodes.append(oh.make_node("Shape", inputs=[ctx_t], outputs=[ctx_shape])) - initializers += [ - onh.from_array(np.array(0, dtype=np.int64), name=ctx_gi0), - onh.from_array(np.array(1, dtype=np.int64), name=ctx_gi1), - onh.from_array(np.array([0], dtype=np.int64), name=ctx_ax0), - onh.from_array(np.array([E], dtype=np.int64), name=ctx_E_1d), - ] - nodes.append(oh.make_node("Gather", inputs=[ctx_shape, ctx_gi0], outputs=[ctx_b_sc])) - nodes.append(oh.make_node("Gather", inputs=[ctx_shape, ctx_gi1], outputs=[ctx_t_sc])) - nodes.append(oh.make_node("Unsqueeze", inputs=[ctx_b_sc, ctx_ax0], outputs=[ctx_b_1d])) - nodes.append(oh.make_node("Unsqueeze", inputs=[ctx_t_sc, ctx_ax0], outputs=[ctx_t_1d])) - nodes.append(oh.make_node("Concat", inputs=[ctx_b_1d, ctx_t_1d, ctx_E_1d], outputs=[ctx_3d], axis=0)) - nodes.append(oh.make_node("Reshape", inputs=[ctx_t, ctx_3d], outputs=[ctx_merged])) - - # --- Output projection (rank-3 input: (B, T, E)) --- - out = _add_dense_nd( - module.out_proj, f"{prefix}_out_proj", ctx_merged, nodes, initializers, quant_fn, use_qonnx, store_integer_weights - ) - - # --- Average attention weights over heads: (B, H, T, S) → (B, T, S) --- - # Emitted so that getitem(mha, 1) has a valid ONNX value name. - avg_attn = f"{prefix}_avg_attn_weights" - nodes.append(oh.make_node("ReduceMean", inputs=[attn_w_name], outputs=[avg_attn], axes=[1], keepdims=0)) - - # --- Optional transpose back for seq-first output --- - if not module.batch_first: - out_final = f"{prefix}_out_seq_first" - nodes.append(oh.make_node("Transpose", inputs=[out], outputs=[out_final], perm=[1, 0, 2])) - return out_final, avg_attn - - return out, avg_attn - - -# --------------------------------------------------------------------------- -# shared module dispatch (used by both sequential and FX converters) -# --------------------------------------------------------------------------- - - -def _emit_module( - module, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights, integer_ops=False -): - """Emit ONNX nodes for a single PQuant or standard torch.nn module.""" - if isinstance(module, PQDense): - return _add_dense( - module, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights, integer_ops - ) - if isinstance(module, PQConv2d): - return _add_conv( - module, - prefix, - current, - nodes, - initializers, - ndim=2, - quant_fn=quant_fn, - use_qonnx=use_qonnx, - store_integer_weights=store_integer_weights, - ) - if isinstance(module, PQConv1d): - return _add_conv( - module, - prefix, - current, - nodes, - initializers, - ndim=1, - quant_fn=quant_fn, - use_qonnx=use_qonnx, - store_integer_weights=store_integer_weights, - ) - if isinstance(module, (PQBatchNorm2d, PQBatchNorm1d)): - return _add_batchnorm(module, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights) - if isinstance(module, PQLayerNorm): - return _add_layernorm(module, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights) - if isinstance(module, PQAvgPool2d): - return _add_avgpool(module, prefix, current, nodes, initializers, ndim=2, quant_fn=quant_fn) - if isinstance(module, PQAvgPool1d): - return _add_avgpool(module, prefix, current, nodes, initializers, ndim=1, quant_fn=quant_fn) - if isinstance(module, nn.ReLU): - out = f"{prefix}_relu" - nodes.append(oh.make_node("Relu", inputs=[current], outputs=[out])) - return out - if isinstance(module, nn.Flatten): - out = f"{prefix}_flatten" - nodes.append(oh.make_node("Flatten", inputs=[current], outputs=[out], axis=module.start_dim)) - return out - if isinstance(module, (nn.BatchNorm1d, nn.BatchNorm2d)): - gamma_name = f"{prefix}_bn_gamma" - beta_name = f"{prefix}_bn_beta" - mean_name = f"{prefix}_bn_mean" - var_name = f"{prefix}_bn_var" - initializers += [ - onh.from_array(module.weight.detach().cpu().numpy().astype(np.float32), name=gamma_name), - onh.from_array(module.bias.detach().cpu().numpy().astype(np.float32), name=beta_name), - onh.from_array(module.running_mean.detach().cpu().numpy().astype(np.float32), name=mean_name), - onh.from_array(module.running_var.detach().cpu().numpy().astype(np.float32), name=var_name), - ] - out = f"{prefix}_bn" - nodes.append( - oh.make_node( - "BatchNormalization", - inputs=[current, gamma_name, beta_name, mean_name, var_name], - outputs=[out], - epsilon=float(module.eps), - ) - ) - return out - if isinstance(module, (nn.Dropout, nn.Dropout2d)): - return current # identity at inference - if isinstance(module, nn.LeakyReLU): - out = f"{prefix}_leakyrelu" - nodes.append(oh.make_node("LeakyRelu", inputs=[current], outputs=[out], alpha=module.negative_slope)) - return out - if isinstance(module, nn.MaxPool2d): - out = f"{prefix}_maxpool" - kernel = module.kernel_size if isinstance(module.kernel_size, (list, tuple)) else [module.kernel_size] * 2 - stride = module.stride if isinstance(module.stride, (list, tuple)) else [module.stride] * 2 - pad = module.padding if isinstance(module.padding, (list, tuple)) else [module.padding] * 2 - nodes.append( - oh.make_node( - "MaxPool", - inputs=[current], - outputs=[out], - kernel_shape=list(kernel), - strides=list(stride), - pads=[pad[0], pad[1], pad[0], pad[1]], - ) - ) - return out - if isinstance(module, nn.Upsample): - # Emit a Resize node with nearest/bilinear mode and scale factors. - roi_name = f"{prefix}_upsample_roi" - scales_name = f"{prefix}_upsample_scales" - initializers.append(onh.from_array(np.array([], dtype=np.float32), name=roi_name)) - scale_factor = module.scale_factor - if isinstance(scale_factor, (int, float)): - scale_factor = (scale_factor, scale_factor) - scales = np.array([1.0, 1.0, float(scale_factor[0]), float(scale_factor[1])], dtype=np.float32) - initializers.append(onh.from_array(scales, name=scales_name)) - mode = "nearest" if module.mode == "nearest" else "linear" - out = f"{prefix}_upsample" - nodes.append( - oh.make_node( - "Resize", - inputs=[current, roi_name, scales_name], - outputs=[out], - mode=mode, - coordinate_transformation_mode="asymmetric", - ) - ) - return out - if isinstance(module, PQActivation): - current = _maybe_quant_input(module, prefix, current, nodes, initializers, quant_fn) - act = module.activation_name - act_out = f"{prefix}_act" - if act == "relu": - nodes.append(oh.make_node("Relu", inputs=[current], outputs=[act_out])) - elif act == "tanh": - nodes.append(oh.make_node("Tanh", inputs=[current], outputs=[act_out])) - elif act == "hard_tanh": - cmin_name = f"{prefix}_htanh_min" - cmax_name = f"{prefix}_htanh_max" - initializers += [ - onh.from_array(np.array(-1.0, dtype=np.float32), name=cmin_name), - onh.from_array(np.array(1.0, dtype=np.float32), name=cmax_name), - ] - nodes.append(oh.make_node("Clip", inputs=[current, cmin_name, cmax_name], outputs=[act_out])) - elif act == "leaky_relu": - nodes.append( - oh.make_node( - "LeakyRelu", inputs=[current], outputs=[act_out], alpha=module.activation_function.negative_slope - ) - ) - elif act == "gelu": - # Decompose so the default opset (13) works; ONNX added a Gelu op only in opset 20. - approximate = getattr(module.activation_function, "approximate", "none") - half_name = f"{prefix}_gelu_half" - one_name = f"{prefix}_gelu_one" - initializers += [ - onh.from_array(np.array(0.5, dtype=np.float32), name=half_name), - onh.from_array(np.array(1.0, dtype=np.float32), name=one_name), - ] - if approximate == "tanh": - # 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))) - c0_name = f"{prefix}_gelu_sqrt2_over_pi" - c1_name = f"{prefix}_gelu_c1" - three_name = f"{prefix}_gelu_three" - initializers += [ - onh.from_array(np.array(np.sqrt(2.0 / np.pi), dtype=np.float32), name=c0_name), - onh.from_array(np.array(0.044715, dtype=np.float32), name=c1_name), - onh.from_array(np.array(3.0, dtype=np.float32), name=three_name), - ] - x3 = f"{prefix}_gelu_x3" - cx3 = f"{prefix}_gelu_cx3" - inner = f"{prefix}_gelu_inner" - scaled = f"{prefix}_gelu_scaled" - tanh_out = f"{prefix}_gelu_tanh" - plus_one = f"{prefix}_gelu_plus1" - x_times = f"{prefix}_gelu_xprod" - nodes += [ - oh.make_node("Pow", inputs=[current, three_name], outputs=[x3]), - oh.make_node("Mul", inputs=[x3, c1_name], outputs=[cx3]), - oh.make_node("Add", inputs=[current, cx3], outputs=[inner]), - oh.make_node("Mul", inputs=[inner, c0_name], outputs=[scaled]), - oh.make_node("Tanh", inputs=[scaled], outputs=[tanh_out]), - oh.make_node("Add", inputs=[tanh_out, one_name], outputs=[plus_one]), - oh.make_node("Mul", inputs=[current, plus_one], outputs=[x_times]), - oh.make_node("Mul", inputs=[x_times, half_name], outputs=[act_out]), - ] - else: - # Exact: 0.5 * x * (1 + erf(x / sqrt(2))) - inv_sqrt2_name = f"{prefix}_gelu_inv_sqrt2" - initializers.append(onh.from_array(np.array(1.0 / np.sqrt(2.0), dtype=np.float32), name=inv_sqrt2_name)) - scaled = f"{prefix}_gelu_scaled" - erf_out = f"{prefix}_gelu_erf" - plus_one = f"{prefix}_gelu_plus1" - x_times = f"{prefix}_gelu_xprod" - nodes += [ - oh.make_node("Mul", inputs=[current, inv_sqrt2_name], outputs=[scaled]), - oh.make_node("Erf", inputs=[scaled], outputs=[erf_out]), - oh.make_node("Add", inputs=[erf_out, one_name], outputs=[plus_one]), - oh.make_node("Mul", inputs=[current, plus_one], outputs=[x_times]), - oh.make_node("Mul", inputs=[x_times, half_name], outputs=[act_out]), - ] - else: - raise TypeError(f"PQActivation: unsupported activation {act!r} for ONNX export") - current = act_out - current = _maybe_quant_output(module, prefix, current, nodes, initializers, quant_fn) - return current - if isinstance(module, PQMultiheadAttention): - # Sequential converter: treat as self-attention (Q = K = V = current). - # Returns (out_name, avg_attn_name); expose only the attention output. - out, _ = _add_mha( - module, prefix, current, current, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights - ) - return out - if isinstance(module, Quantizer): - # Standalone quantizer (e.g. an auto-inserted missing quantizer or a - # constant-matrix quantizer): emit a single QDQ node from its k/i/f. - k, i, f = module.get_quantization_bits() - new_nodes, out = quant_fn( - prefix, current, module.round_mode, k, i, f, initializers, overflow_mode=getattr(module, "overflow", "SAT") - ) - nodes.extend(new_nodes) - return out - raise TypeError(f"Unsupported module type for ONNX export: {type(module).__name__}") - - -# --------------------------------------------------------------------------- -# main conversion -# --------------------------------------------------------------------------- - - -def convert_to_onnx( - model: nn.Sequential, - input_shape: tuple, - output_path: str = "model.onnx", - opset: int = 13, - use_qonnx: bool = False, - store_integer_weights: bool = False, - integer_ops: bool = False, - include_clip: bool = True, - batch_size: int | None = None, -) -> onnx.ModelProto: - """ - Convert a Sequential model of PQuant layers to ONNX or QONNX. - - Args: - model: Trained nn.Sequential. Call apply_final_compression() - on all PQ modules before passing here. - input_shape: Shape of a single sample (excluding batch), e.g. (3, 32, 32). - output_path: Where to save the .onnx file. - opset: ONNX opset version (≥13 required for per-channel DequantizeLinear). - use_qonnx: If True, emit QONNX Quant custom nodes (requires qonnx runtime). - If False (default), emit Clip+QuantizeLinear+DequantizeLinear - nodes runnable with plain onnxruntime. - store_integer_weights: If True (and use_qonnx=False), store weight/bias initializers - as int8/uint8 followed by DequantizeLinear instead of float32. - Ignored when use_qonnx=True or integer_ops=True. - integer_ops: If True (and use_qonnx=False), use MatMulInteger for Dense layers - so the inner product runs in int32 arithmetic. Weights are stored - as int8 (pre-transposed) and a single DequantizeLinear converts the - int32 accumulator back to float using the combined scale s_x * s_w. - Implies integer weight storage; store_integer_weights is ignored. - include_clip: Prepend a Clip node before each QuantizeLinear when True (default). - Set to False to emit bare QuantizeLinear+DequantizeLinear pairs — - safe when values are guaranteed in-range at inference time since - QuantizeLinear saturates naturally. Ignored when use_qonnx=True. - batch_size: If not None, fix the batch dimension of all graph inputs and - outputs to this value. If None (default), the batch dimension - is left dynamic. - - Returns: - The constructed onnx.ModelProto. - """ - model.eval() - quant_fn = _quant_node if use_qonnx else functools.partial(_qdq_node, include_clip=include_clip) - - nodes: list[onnx.NodeProto] = [] - initializers: list[onnx.TensorProto] = [] - current = "input" - - for layer_idx, module in enumerate(model): - prefix = f"layer{layer_idx}" - current = _emit_module( - module, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights, integer_ops - ) - - with torch.no_grad(): - dummy_out = model(torch.zeros(1, *input_shape)) - batch_dim = batch_size # None → dynamic, int → fixed - output_shape = [batch_dim] + list(dummy_out.shape[1:]) - - batch_dim_vi = oh.make_tensor_value_info("input", TensorProto.FLOAT, [batch_dim, *input_shape]) - output_vi = oh.make_tensor_value_info(current, TensorProto.FLOAT, output_shape) - - graph = oh.make_graph( - nodes=nodes, - name="pquant_onnx", - inputs=[batch_dim_vi], - outputs=[output_vi], - initializer=initializers, - ) - - opset_imports = [oh.make_opsetid("", opset)] - if use_qonnx: - opset_imports.append(oh.make_opsetid("qonnx.custom_op.general", 1)) - model_proto = oh.make_model(graph, opset_imports=opset_imports) - model_proto.ir_version = 6 - - onnx.checker.check_model(model_proto) - onnx.save(model_proto, output_path) - fmt = "QONNX" if use_qonnx else "ONNX (QDQ)" - logging.info("Saved %s model → %s", fmt, output_path) - return model_proto - - -# --------------------------------------------------------------------------- -# Hardware-targeted static-QDQ LayerNormalization graph -# --------------------------------------------------------------------------- - - -def _is_pow2(n: int) -> bool: - return n > 0 and (n & (n - 1)) == 0 - - -def export_qdq_layernorm( - output_path: str, - input_shape, - gamma: np.ndarray, - beta: np.ndarray, - input_scale_log2: int, - output_scale_log2: int, - eps_q0: int = 1, - opset: int = 17, -) -> onnx.ModelProto: - """Build and save a single-LayerNormalization ONNX graph using static QDQ quantization. - - Graph layout:: - - int8 input -> DequantizeLinear -> LayerNormalization -> QuantizeLinear -> DequantizeLinear -> output - - All quantization parameters are explicit float32 initializers (no dynamic - tensors). Per-tensor quantization only; activation zero-points are 0. - - Constraints (validated at build time, not in the graph): - * ``input_shape`` is rank-2 or rank-3 with no dynamic dims. - * Last dim ``D`` is a power of two AND a multiple of 32. - * ``gamma``/``beta`` are 1-D float arrays of length ``D``. - * ``gamma`` is exactly representable as int16 with scale ``2**-7`` (Q7). - * ``beta`` is exactly representable as int16 with scale ``2**-15`` (Q15). - * Input/output scales are exact powers of two, given as log2 exponents. - * ``epsilon = eps_q0 * input_scale**2`` with integer ``eps_q0 >= 1``. - * Normalization axis is the last axis. - - Args: - output_path: Where to save the .onnx file. - input_shape: Static shape of the int8 graph input, e.g. ``(4, 64)`` or ``(1, 4, 64)``. - gamma: Constant gamma initializer, shape ``(D,)``. - beta: Constant beta initializer, shape ``(D,)``. - input_scale_log2: Integer ``a`` with input scale ``= 2**a``. - output_scale_log2: Integer ``b`` with output scale ``= 2**b``. - eps_q0: Positive integer ``>= 1``; ``epsilon = eps_q0 * (2**a)**2``. - opset: ONNX opset version (must be ``>= 17`` for LayerNormalization). - - Returns: - The constructed ``onnx.ModelProto``. - """ - # ----- validate shape ----- - input_shape = tuple(int(d) for d in input_shape) - if len(input_shape) not in (2, 3): - raise ValueError(f"input_shape rank must be 2 or 3, got {len(input_shape)} ({input_shape})") - for d in input_shape: - if d <= 0: - raise ValueError(f"input_shape must be fully static and positive, got {input_shape}") - D = input_shape[-1] - if not _is_pow2(D): - raise ValueError(f"last dim must be a power of two, got {D}") - if D % 32 != 0: - raise ValueError(f"last dim must be a multiple of 32, got {D}") - - # ----- validate gamma / beta ----- - gamma = np.asarray(gamma, dtype=np.float32) - beta = np.asarray(beta, dtype=np.float32) - if gamma.shape != (D,): - raise ValueError(f"gamma must have shape ({D},), got {gamma.shape}") - if beta.shape != (D,): - raise ValueError(f"beta must have shape ({D},), got {beta.shape}") - - GAMMA_F = 7 # Q7 in int16 -> scale = 2**-7 - BETA_F = 15 # Q15 in int16 -> scale = 2**-15 - INT16_MIN, INT16_MAX = -(2**15), 2**15 - 1 - - def _check_q_int16(arr: np.ndarray, frac_bits: int, name: str) -> None: - scaled = arr.astype(np.float64) * (2**frac_bits) - rounded = np.round(scaled) - # Exactly representable: rounding is a no-op (within fp slack). - if not np.allclose(scaled, rounded, atol=1e-4): - raise ValueError( - f"{name} not exactly representable as int16 Q{frac_bits} " - f"(max abs round error = {np.max(np.abs(scaled - rounded)):.6g})" - ) - if rounded.min() < INT16_MIN or rounded.max() > INT16_MAX: - raise ValueError(f"{name} overflows int16 at Q{frac_bits} " f"(range [{rounded.min()}, {rounded.max()}])") - - _check_q_int16(gamma, GAMMA_F, "gamma") - _check_q_int16(beta, BETA_F, "beta") - - # ----- validate quant params ----- - input_scale_log2 = int(input_scale_log2) - output_scale_log2 = int(output_scale_log2) - eps_q0 = int(eps_q0) - if eps_q0 < 1: - raise ValueError(f"eps_q0 must be >= 1, got {eps_q0}") - - if opset < 17: - raise ValueError(f"opset must be >= 17 for LayerNormalization, got {opset}") - - input_scale = float(2.0**input_scale_log2) - output_scale = float(2.0**output_scale_log2) - epsilon = float(eps_q0) * input_scale * input_scale - - # ----- build initializers ----- - initializers = [ - onh.from_array(np.array(input_scale, dtype=np.float32), name="input_scale"), - onh.from_array(np.array(0, dtype=np.int8), name="input_zero_point"), - onh.from_array(np.array(output_scale, dtype=np.float32), name="output_scale"), - onh.from_array(np.array(0, dtype=np.int8), name="output_zero_point"), - onh.from_array(gamma.astype(np.float32), name="gamma"), - onh.from_array(beta.astype(np.float32), name="beta"), - ] - - # ----- build nodes ----- - nodes = [ - oh.make_node( - "DequantizeLinear", - inputs=["input_q", "input_scale", "input_zero_point"], - outputs=["x_dq"], - name="input_dq", - ), - oh.make_node( - "LayerNormalization", - inputs=["x_dq", "gamma", "beta"], - outputs=["ln_out"], - name="layernorm", - axis=-1, - epsilon=epsilon, - ), - oh.make_node( - "QuantizeLinear", - inputs=["ln_out", "output_scale", "output_zero_point"], - outputs=["y_q"], - name="output_q", - ), - oh.make_node( - "DequantizeLinear", - inputs=["y_q", "output_scale", "output_zero_point"], - outputs=["output"], - name="output_dq", - ), - ] - - # ----- build graph + model ----- - input_vi = oh.make_tensor_value_info("input_q", TensorProto.INT8, list(input_shape)) - output_vi = oh.make_tensor_value_info("output", TensorProto.FLOAT, list(input_shape)) - - graph = oh.make_graph( - nodes=nodes, - name="qdq_layernorm", - inputs=[input_vi], - outputs=[output_vi], - initializer=initializers, - ) - - model_proto = oh.make_model(graph, opset_imports=[oh.make_opsetid("", opset)]) - model_proto.ir_version = 8 - - # Strip any initializer names that the onnx library may have added to graph.input. - _init_names = {t.name for t in model_proto.graph.initializer} - _data_inputs = [vi for vi in model_proto.graph.input if vi.name not in _init_names] - del model_proto.graph.input[:] - model_proto.graph.input.extend(_data_inputs) - - onnx.checker.check_model(model_proto) - onnx.save(model_proto, output_path) - return model_proto - - -# --------------------------------------------------------------------------- -# FX-based conversion (supports arbitrary nn.Module topology / skip connections) -# --------------------------------------------------------------------------- - - -class _PQTracer(_fx.Tracer): - """Tracer that treats all PQuant layer types (and standard torch.nn leaves) as atomic.""" - - _LEAF_TYPES = ( - PQDense, - PQConv2d, - PQConv1d, - PQBatchNorm1d, - PQBatchNorm2d, - PQLayerNorm, - PQAvgPool1d, - PQAvgPool2d, - PQMultiheadAttention, - PQActivation, - Quantizer, - ) - - def is_leaf_module(self, m: nn.Module, qualname: str) -> bool: - return isinstance(m, self._LEAF_TYPES) or super().is_leaf_module(m, qualname) - - -def convert_to_onnx_fx( - model: nn.Module, - input_shape: tuple, - output_path: str = "model.onnx", - opset: int = 13, - use_qonnx: bool = False, - store_integer_weights: bool = False, - integer_ops: bool = False, - include_clip: bool = True, -) -> onnx.ModelProto: - """ - Convert any PQuant nn.Module to ONNX using torch.fx symbolic tracing. - - Unlike convert_to_onnx(), this function works with arbitrary model topologies - including residual/skip connections, branches, and concatenations. It requires - the model to be symbolically traceable (no data-dependent control flow). - - Args match convert_to_onnx() exactly; see that function for parameter docs. - """ - model.eval() - quant_fn = _quant_node if use_qonnx else functools.partial(_qdq_node, include_clip=include_clip) - - graph = _PQTracer().trace(model) - gm = _fx.GraphModule(model, graph) - - # ShapeProp populates node.meta["tensor_meta"], which transpose/permute - # need to expand torch's two-arg .transpose(d0, d1) into a full ONNX perm. - from torch.fx.passes.shape_prop import ShapeProp - - # Build the probe tensor on the model's own device so ShapeProp doesn't hit a - # device mismatch when a default device (e.g. CUDA) is set via torch.set_default_device. - device = next((p.device for p in model.parameters()), None) - with torch.no_grad(): - ShapeProp(gm).propagate(torch.zeros(1, *input_shape, device=device)) - - onnx_nodes: list[onnx.NodeProto] = [] - initializers: list[onnx.TensorProto] = [] - node_to_name: dict[_fx.Node, str] = {} - output_names: list[str] = [] - - def _res(arg) -> str: - if isinstance(arg, _fx.Node): - return node_to_name[arg] - raise TypeError(f"Expected fx.Node, got {type(arg)}") - - def _binop_inputs(node: _fx.Node) -> list[str]: - # Like _res for both args, but lifts scalar literals (int/float/bool) - # to float32 initializers so patterns like ``x / 2.0`` work. - names: list[str] = [] - for i, a in enumerate(node.args[:2]): - if isinstance(a, _fx.Node): - names.append(node_to_name[a]) - elif isinstance(a, (int, float, bool)): - cname = f"{node.name}_arg{i}_const" - initializers.append(onh.from_array(np.array(float(a), dtype=np.float32), name=cname)) - names.append(cname) - else: - raise TypeError(f"FX export: unsupported binary-op arg type {type(a).__name__}") - return names - - def _rank(n: _fx.Node) -> int: - meta = n.meta.get("tensor_meta") - if meta is None or not hasattr(meta, "shape"): - raise RuntimeError(f"FX export: ShapeProp did not produce tensor_meta for {n.name!r}") - return len(meta.shape) - - def _swap_perm(rank: int, d0: int, d1: int) -> list[int]: - perm = list(range(rank)) - a, b = d0 % rank, d1 % rank - perm[a], perm[b] = perm[b], perm[a] - return perm - - def _resolve_perm_dims(args, rank: int) -> list[int]: - # Accept both permute(d0, d1, ...) and permute([d0, d1, ...]) shapes. - if len(args) == 1 and isinstance(args[0], (list, tuple)): - dims = args[0] - else: - dims = args - return [int(d) % rank for d in dims] - - for node in gm.graph.nodes: - if node.op == "placeholder": - node_to_name[node] = "input" - - elif node.op == "get_attr": - # Constant tensor attributes — store as initializer on first use. - # Retrieve the actual tensor from the GraphModule. - obj = gm - for part in node.target.split("."): - obj = getattr(obj, part) - attr_name = node.name - if isinstance(obj, torch.Tensor): - initializers.append(onh.from_array(obj.detach().cpu().numpy(), name=attr_name)) - node_to_name[node] = attr_name - - elif node.op == "call_module": - mod = gm.get_submodule(node.target) - mod_prefix = node.name.replace(".", "_") - if isinstance(mod, PQMultiheadAttention): - # node.args = (query, key, value[, key_padding_mask, attn_mask, ...]) - q_name = node_to_name[node.args[0]] - k_name = node_to_name[node.args[1]] if len(node.args) > 1 else q_name - v_name = node_to_name[node.args[2]] if len(node.args) > 2 else q_name - out_name, avg_attn_name = _add_mha( - mod, - mod_prefix, - q_name, - k_name, - v_name, - onnx_nodes, - initializers, - quant_fn, - use_qonnx, - store_integer_weights, - ) - # Store tuple so operator.getitem(node, 0/1) resolves correctly. - node_to_name[node] = (out_name, avg_attn_name) - else: - current = _emit_module( - mod, - mod_prefix, - node_to_name[node.args[0]], - onnx_nodes, - initializers, - quant_fn, - use_qonnx, - store_integer_weights, - integer_ops, - ) - node_to_name[node] = current - - elif node.op == "call_function": - fn = node.target - - if fn is _operator.getitem: - # Unpack a tuple output (e.g. from PQMultiheadAttention). - container = node_to_name[node.args[0]] - if not isinstance(container, tuple): - raise TypeError( - f"operator.getitem on non-tuple node {node.args[0].name!r} " f"is not supported in FX ONNX export" - ) - node_to_name[node] = container[node.args[1]] - continue - - if fn in (torch.add, _operator.add, _operator.iadd): - out = f"{node.name}_add" - onnx_nodes.append(oh.make_node("Add", inputs=_binop_inputs(node), outputs=[out])) - node_to_name[node] = out - - elif fn in (torch.mul, _operator.mul): - out = f"{node.name}_mul" - onnx_nodes.append(oh.make_node("Mul", inputs=_binop_inputs(node), outputs=[out])) - node_to_name[node] = out - - elif fn in (torch.sub, _operator.sub, _operator.isub): - out = f"{node.name}_sub" - onnx_nodes.append(oh.make_node("Sub", inputs=_binop_inputs(node), outputs=[out])) - node_to_name[node] = out - - elif fn in (torch.div, _operator.truediv, _operator.itruediv): - out = f"{node.name}_div" - onnx_nodes.append(oh.make_node("Div", inputs=_binop_inputs(node), outputs=[out])) - node_to_name[node] = out - - elif fn in (torch.matmul, _operator.matmul): - out = f"{node.name}_matmul" - onnx_nodes.append(oh.make_node("MatMul", inputs=_binop_inputs(node), outputs=[out])) - node_to_name[node] = out - - elif fn is torch.transpose: - # torch.transpose(t, d0, d1) swaps two dims; ONNX needs a full perm. - rank = _rank(node.args[0]) - perm = _swap_perm(rank, int(node.args[1]), int(node.args[2])) - out = f"{node.name}_transpose" - onnx_nodes.append(oh.make_node("Transpose", inputs=[_res(node.args[0])], outputs=[out], perm=perm)) - node_to_name[node] = out - - elif fn is torch.permute: - rank = _rank(node.args[0]) - perm = _resolve_perm_dims(node.args[1:], rank) - out = f"{node.name}_permute" - onnx_nodes.append(oh.make_node("Transpose", inputs=[_res(node.args[0])], outputs=[out], perm=perm)) - node_to_name[node] = out - - elif fn is torch.cat: - tensors = [_res(a) for a in node.args[0]] - dim = node.args[1] if len(node.args) > 1 else node.kwargs.get("dim", 0) - out = f"{node.name}_concat" - onnx_nodes.append(oh.make_node("Concat", inputs=tensors, outputs=[out], axis=int(dim))) - node_to_name[node] = out - - elif fn in (_F.relu, torch.relu): - out = f"{node.name}_relu" - onnx_nodes.append(oh.make_node("Relu", inputs=[_res(node.args[0])], outputs=[out])) - node_to_name[node] = out - - elif fn in (_F.sigmoid, torch.sigmoid): - out = f"{node.name}_sigmoid" - onnx_nodes.append(oh.make_node("Sigmoid", inputs=[_res(node.args[0])], outputs=[out])) - node_to_name[node] = out - - elif fn is torch.flatten: - start_dim = node.args[1] if len(node.args) > 1 else node.kwargs.get("start_dim", 0) - out = f"{node.name}_flatten" - onnx_nodes.append(oh.make_node("Flatten", inputs=[_res(node.args[0])], outputs=[out], axis=int(start_dim))) - node_to_name[node] = out - - else: - raise TypeError(f"Unsupported call_function for FX ONNX export: {fn}") - - elif node.op == "call_method": - x = _res(node.args[0]) - - if node.target == "relu": - out = f"{node.name}_relu" - onnx_nodes.append(oh.make_node("Relu", inputs=[x], outputs=[out])) - node_to_name[node] = out - - elif node.target == "flatten": - start_dim = node.args[1] if len(node.args) > 1 else node.kwargs.get("start_dim", 1) - out = f"{node.name}_flatten" - onnx_nodes.append(oh.make_node("Flatten", inputs=[x], outputs=[out], axis=int(start_dim))) - node_to_name[node] = out - - elif node.target in ("view", "reshape"): - shape_vals = [] - for a in node.args[1:]: - if not isinstance(a, int): - raise TypeError("Dynamic reshape (non-constant shape) is not supported in FX ONNX export") - shape_vals.append(a) - shape_name = f"{node.name}_shape" - out = f"{node.name}_reshape" - initializers.append(onh.from_array(np.array(shape_vals, dtype=np.int64), name=shape_name)) - onnx_nodes.append(oh.make_node("Reshape", inputs=[x, shape_name], outputs=[out])) - node_to_name[node] = out - - elif node.target == "transpose": - rank = _rank(node.args[0]) - perm = _swap_perm(rank, int(node.args[1]), int(node.args[2])) - out = f"{node.name}_transpose" - onnx_nodes.append(oh.make_node("Transpose", inputs=[x], outputs=[out], perm=perm)) - node_to_name[node] = out - - elif node.target == "permute": - rank = _rank(node.args[0]) - perm = _resolve_perm_dims(node.args[1:], rank) - out = f"{node.name}_permute" - onnx_nodes.append(oh.make_node("Transpose", inputs=[x], outputs=[out], perm=perm)) - node_to_name[node] = out - - elif node.target == "matmul": - out = f"{node.name}_matmul" - onnx_nodes.append(oh.make_node("MatMul", inputs=[x, _res(node.args[1])], outputs=[out])) - node_to_name[node] = out - - else: - raise TypeError(f"Unsupported call_method for FX ONNX export: {node.target!r}") - - elif node.op == "output": - ret = node.args[0] - rets = list(ret) if isinstance(ret, (tuple, list)) else [ret] - for r in rets: - if not isinstance(r, _fx.Node): - raise TypeError("FX ONNX export: unsupported (non-tensor) model output") - val = node_to_name[r] - # MHA nodes store a tuple (out, avg_attn); expose the attention output. - output_names.append(val[0] if isinstance(val, tuple) else val) - - with torch.no_grad(): - dummy_out = model(torch.zeros(1, *input_shape, device=device)) - dummy_outs = list(dummy_out) if isinstance(dummy_out, (tuple, list)) else [dummy_out] - - batch_dim = oh.make_tensor_value_info("input", TensorProto.FLOAT, [None, *input_shape]) - output_vis = [ - oh.make_tensor_value_info(name, TensorProto.FLOAT, [None] + list(t.shape[1:])) - for name, t in zip(output_names, dummy_outs) - ] - - onnx_graph = oh.make_graph( - nodes=onnx_nodes, - name="pquant_onnx_fx", - inputs=[batch_dim], - outputs=output_vis, - initializer=initializers, - ) - - opset_imports = [oh.make_opsetid("", opset)] - if use_qonnx: - opset_imports.append(oh.make_opsetid("qonnx.custom_op.general", 1)) - model_proto = oh.make_model(onnx_graph, opset_imports=opset_imports) - model_proto.ir_version = 6 - - onnx.checker.check_model(model_proto) - onnx.save(model_proto, output_path) - fmt = "QONNX" if use_qonnx else "ONNX (QDQ)" - logging.info("Saved %s model (FX) → %s", fmt, output_path) - return model_proto - - -# --------------------------------------------------------------------------- -# usage example -# --------------------------------------------------------------------------- - -if __name__ == "__main__": - import onnxruntime as ort - - import pquant - - cfg = pquant.cs_config() - cfg.quantization_parameters.granularity = "per-channel" - - model = nn.Sequential( - PQConv2d(cfg, in_channels=3, out_channels=16, kernel_size=3, padding=1), - PQBatchNorm2d(cfg, num_features=16), - nn.ReLU(), - PQAvgPool2d(cfg, kernel_size=2, stride=2), - nn.Flatten(), - PQDense(cfg, in_features=16 * 16 * 16, out_features=64), - nn.ReLU(), - PQDense(cfg, in_features=64, out_features=10), - ) - - x = torch.randn(4, 3, 32, 32) - with torch.no_grad(): - model(x) - - for module in model.modules(): - if hasattr(module, "apply_final_compression"): - module.apply_final_compression() - - model.eval() - with torch.no_grad(): - torch_out = model(x).numpy() - - qonnx_path = "model_qonnx.onnx" - onnx_path = "model_qdq.onnx" - convert_to_onnx(model, input_shape=(3, 32, 32), output_path=qonnx_path, use_qonnx=True) - convert_to_onnx(model, input_shape=(3, 32, 32), output_path=onnx_path, use_qonnx=False) - - from qonnx.core.modelwrapper import ModelWrapper - from qonnx.core.onnx_exec import execute_onnx - from qonnx.transformation.infer_shapes import InferShapes - - qmodel = ModelWrapper(qonnx_path) - qmodel.graph.input[0].type.tensor_type.shape.dim[0].dim_value = x.shape[0] - qmodel = qmodel.transform(InferShapes()) - input_name = qmodel.graph.input[0].name - output_name = qmodel.graph.output[0].name - qonnx_out = execute_onnx(qmodel, {input_name: x.numpy()})[output_name] - - sess = ort.InferenceSession(onnx_path) - onnx_out = sess.run(None, {sess.get_inputs()[0].name: x.numpy()})[0] - - print(f"\n{'':=<55}") # noqa: T201 - print(f" max |torch - qonnx| : {np.abs(torch_out - qonnx_out).max():.6f}") # noqa: T201 - print(f" max |torch - onnx| : {np.abs(torch_out - onnx_out).max():.6f}") # noqa: T201 - print(f" max |qonnx - onnx| : {np.abs(qonnx_out - onnx_out).max():.6f}") # noqa: T201 - print(f"{'':=<55}") # noqa: T201 diff --git a/src/pquant/core/torch/onnx/__init__.py b/src/pquant/core/torch/onnx/__init__.py new file mode 100644 index 0000000..83612de --- /dev/null +++ b/src/pquant/core/torch/onnx/__init__.py @@ -0,0 +1,5 @@ +from pquant.core.torch.onnx.convert_to_onnx import ( + convert_to_onnx, +) + +__all__ = ["convert_to_onnx"] diff --git a/src/pquant/core/torch/onnx/convert_to_onnx.py b/src/pquant/core/torch/onnx/convert_to_onnx.py new file mode 100644 index 0000000..3d06a68 --- /dev/null +++ b/src/pquant/core/torch/onnx/convert_to_onnx.py @@ -0,0 +1,610 @@ +""" +Convert a PQuant model to ONNX or QONNX format. + +Pass ``use_qonnx=True`` to emit QONNX ``Quant`` custom nodes (requires the +qonnx runtime). Pass ``use_qonnx=False`` (default) to emit standard +``Clip + QuantizeLinear + DequantizeLinear`` nodes runnable with plain +onnxruntime. +""" + +import functools +import logging +import operator +import os + +import onnx +import onnx.helper as oh +import torch +import torch.fx as fx +import torch.nn as nn +import torch.nn.functional as F +from onnx import TensorProto +from torch.fx.passes.shape_prop import ShapeProp + +os.environ["KERAS_BACKEND"] = "torch" # must be set before any keras/pquant import + +from pquant.core.onnx_common import ( # noqa: E402 + add_float_scalar, + add_initializer, + add_int64_array, + apply_quantizer, + emit_getitem, + emit_squeeze, + emit_unsqueeze, + qdq_node, + quant_node, + save_model, +) +from pquant.core.torch.activations import PQActivation # noqa: E402 +from pquant.core.torch.layers import ( # noqa: E402 + PQAvgPool1d, + PQAvgPool2d, + PQBatchNorm1d, + PQBatchNorm2d, + PQConv1d, + PQConv2d, + PQDense, + PQLayerNorm, + PQMultiheadAttention, +) +from pquant.core.torch.onnx.layer_builders import ( # noqa: E402 + _add_activation, + _add_avgpool, + _add_batchnorm, + _add_conv, + _add_dense, + _add_dense_nd, + _add_layernorm, + _add_maxpool, + _add_mha, + _add_upsample, +) +from pquant.core.torch.quantizer import Quantizer # noqa: E402 + + +def _emit_module( + module, + prefix, + current, + nodes, + initializers, + quant_fn, + use_qonnx, + store_integer_weights, + integer_ops=False, + input_rank=None, +): + """Emit ONNX nodes for a single PQuant or standard torch.nn module.""" + if isinstance(module, PQDense): + # Gemm only accepts rank-2 inputs; higher ranks (e.g. [batch, seq, dim]) go through + # MatMul + Add. The integer_ops path is already MatMul-based and handles any rank. + if input_rank is not None and input_rank > 2 and not (integer_ops and not use_qonnx): + return _add_dense_nd(module, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights) + return _add_dense( + module, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights, integer_ops + ) + if isinstance(module, (PQConv2d, PQConv1d)): + ndim = 2 if isinstance(module, PQConv2d) else 1 + return _add_conv( + module, + prefix, + current, + nodes, + initializers, + ndim=ndim, + quant_fn=quant_fn, + use_qonnx=use_qonnx, + store_integer_weights=store_integer_weights, + ) + if isinstance(module, (PQBatchNorm2d, PQBatchNorm1d)): + return _add_batchnorm(module, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights) + if isinstance(module, PQLayerNorm): + return _add_layernorm(module, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights) + if isinstance(module, (PQAvgPool2d, PQAvgPool1d)): + ndim = 2 if isinstance(module, PQAvgPool2d) else 1 + return _add_avgpool(module, prefix, current, nodes, initializers, ndim=ndim, quant_fn=quant_fn) + if isinstance(module, PQActivation): + return _add_activation(module, prefix, current, nodes, initializers, quant_fn) + if isinstance(module, Quantizer): + return apply_quantizer(module, prefix, current, nodes, initializers, quant_fn) + if isinstance(module, nn.ReLU): + out = f"{prefix}_relu" + nodes.append(oh.make_node("Relu", inputs=[current], outputs=[out])) + return out + if isinstance(module, nn.LeakyReLU): + out = f"{prefix}_leakyrelu" + nodes.append(oh.make_node("LeakyRelu", inputs=[current], outputs=[out], alpha=module.negative_slope)) + return out + if isinstance(module, nn.Flatten): + out = f"{prefix}_flatten" + nodes.append(oh.make_node("Flatten", inputs=[current], outputs=[out], axis=module.start_dim)) + return out + if isinstance(module, nn.MaxPool2d): + return _add_maxpool(module, prefix, current, nodes) + if isinstance(module, nn.Upsample): + return _add_upsample(module, prefix, current, nodes, initializers) + if isinstance(module, (nn.Dropout, nn.Dropout2d)): + return current # identity at inference + raise TypeError(f"Unsupported module type for ONNX export: {type(module).__name__}") + + +class PQTracer(fx.Tracer): + _LEAF_TYPES = ( + PQDense, + PQConv2d, + PQConv1d, + PQBatchNorm1d, + PQBatchNorm2d, + PQLayerNorm, + PQAvgPool1d, + PQAvgPool2d, + PQMultiheadAttention, + PQActivation, + Quantizer, + ) + + def is_leaf_module(self, m: nn.Module, qualname: str) -> bool: + return isinstance(m, self._LEAF_TYPES) or super().is_leaf_module(m, qualname) + + +def _normalize_input_shapes(input_shape) -> list[tuple]: + seq = list(input_shape) + if len(seq) > 0 and all(isinstance(s, (list, tuple)) for s in seq): + return [tuple(int(d) for d in s) for s in seq] + return [tuple(int(d) for d in seq)] + + +def _normalize_input_dtypes(input_dtypes, n: int): + torch_map = { + "float32": torch.float32, + "float": torch.float32, + "bool": torch.bool, + "int64": torch.int64, + "int32": torch.int32, + } + tp_map = { + torch.float32: TensorProto.FLOAT, + torch.bool: TensorProto.BOOL, + torch.int64: TensorProto.INT64, + torch.int32: TensorProto.INT32, + } + + if input_dtypes is None: + items = [torch.float32] * n + elif isinstance(input_dtypes, (list, tuple)): + items = list(input_dtypes) + else: + items = [input_dtypes] * n + + if len(items) != n: + raise ValueError(f"input_dtypes has {len(items)} entries but there are {n} input(s)") + + torch_dtypes, tp_dtypes = [], [] + for d in items: + td = torch_map[d] if isinstance(d, str) else d + if td not in tp_map: + raise ValueError(f"Unsupported input dtype {d!r}; expected one of {list(torch_map)}") + torch_dtypes.append(td) + tp_dtypes.append(tp_map[td]) + return torch_dtypes, tp_dtypes + + +def _swap_perm(rank: int, d0: int, d1: int) -> list[int]: + perm = list(range(rank)) + a, b = d0 % rank, d1 % rank + perm[a], perm[b] = perm[b], perm[a] + return perm + + +def _resolve_perm_dims(args, rank: int) -> list[int]: + # Accept both permute(d0, d1, ...) and permute([d0, d1, ...]) shapes. + dims = args[0] if len(args) == 1 and isinstance(args[0], (list, tuple)) else args + return [int(d) % rank for d in dims] + + +class FxGraphEmitter: + """Translate a shape-propagated fx.Graph into ONNX nodes and initializers. + + ``node_to_name`` maps each fx.Node to the name of the ONNX tensor holding + its value — or, for modules with multiple outputs (PQMultiheadAttention), + a tuple of names. + """ + + _BINARY_OPS = { + torch.add: "Add", + operator.add: "Add", + operator.iadd: "Add", + torch.mul: "Mul", + operator.mul: "Mul", + torch.sub: "Sub", + operator.sub: "Sub", + operator.isub: "Sub", + torch.div: "Div", + operator.truediv: "Div", + operator.itruediv: "Div", + torch.matmul: "MatMul", + operator.matmul: "MatMul", + } + _UNARY_OPS = { + F.relu: "Relu", + torch.relu: "Relu", + F.sigmoid: "Sigmoid", + torch.sigmoid: "Sigmoid", + } + + def __init__(self, gm, ph_to_name, quant_fn, use_qonnx, store_integer_weights, integer_ops): + self.gm = gm + self.ph_to_name = ph_to_name + self.quant_fn = quant_fn + self.use_qonnx = use_qonnx + self.store_integer_weights = store_integer_weights + self.integer_ops = integer_ops + self.nodes: list[onnx.NodeProto] = [] + self.initializers: list[onnx.TensorProto] = [] + self.node_to_name: dict[fx.Node, str] = {} + self.output_names: list[str] = [] + + def _run(self) -> list[str]: + for node in self.gm.graph.nodes: + if node.op == "placeholder": + self.node_to_name[node] = self.ph_to_name[node] + elif node.op == "get_attr": + self._emit_get_attr(node) + elif node.op == "call_module": + self._emit_call_module(node) + elif node.op == "call_function": + self._emit_call_function(node) + elif node.op == "call_method": + self._emit_call_method(node) + elif node.op == "output": + self._collect_outputs(node) + return self.output_names + + def _name_of(self, arg) -> str: + if isinstance(arg, fx.Node): + return self.node_to_name[arg] + raise TypeError(f"Expected fx.Node, got {type(arg)}") + + def _binop_inputs(self, node: fx.Node) -> list[str]: + names: list[str] = [] + for idx, arg in enumerate(node.args[:2]): + if isinstance(arg, fx.Node): + names.append(self.node_to_name[arg]) + elif isinstance(arg, (int, float, bool)): + names.append(add_float_scalar(self.initializers, f"{node.name}_arg{idx}_const", float(arg))) + else: + raise TypeError(f"FX export: unsupported binary-op arg type {type(arg).__name__}") + return names + + def _node_shape(self, node: fx.Node) -> tuple: + meta = node.meta.get("tensor_meta") + if meta is None or not hasattr(meta, "shape"): + raise RuntimeError(f"FX export: ShapeProp did not produce tensor_meta for {node.name!r}") + return tuple(meta.shape) + + def _node_rank(self, node: fx.Node) -> int: + return len(self._node_shape(node)) + + def _emit_get_attr(self, node): + obj = self.gm + for part in node.target.split("."): + obj = getattr(obj, part) + if isinstance(obj, torch.Tensor): + add_initializer(self.initializers, node.name, obj.detach().cpu().numpy()) + self.node_to_name[node] = node.name + + def _emit_call_module(self, node): + module = self.gm.get_submodule(node.target) + prefix = node.name.replace(".", "_") + if isinstance(module, PQMultiheadAttention): + self.node_to_name[node] = self._emit_mha_module(module, node, prefix) + return + self.node_to_name[node] = _emit_module( + module, + prefix, + self._name_of(node.args[0]), + self.nodes, + self.initializers, + self.quant_fn, + self.use_qonnx, + self.store_integer_weights, + self.integer_ops, + input_rank=self._node_rank(node.args[0]), + ) + + def _emit_mha_module(self, module, node, prefix) -> tuple: + # forward(query, key, value, key_padding_mask=None, attn_mask=None, ...) + q_name = self._name_of(node.args[0]) + k_name = self._name_of(node.args[1]) if len(node.args) > 1 else q_name + v_name = self._name_of(node.args[2]) if len(node.args) > 2 else q_name + kpm_name = self._mask_name(node, 3, "key_padding_mask") + attn_mask_name = self._mask_name(node, 4, "attn_mask") + return _add_mha( + module, + prefix, + q_name, + k_name, + v_name, + self.nodes, + self.initializers, + self.quant_fn, + self.use_qonnx, + self.store_integer_weights, + key_padding_mask=kpm_name, + attn_mask=attn_mask_name, + ) + + def _mask_name(self, node, pos, kw): + arg = node.args[pos] if len(node.args) > pos else node.kwargs.get(kw) + if arg is None: + return None + if not isinstance(arg, fx.Node): + raise TypeError(f"FX ONNX export: MHA {kw} must be a tensor (constant or input), got {type(arg)}") + return self.node_to_name[arg] + + def _emit_call_function(self, node): + fn = node.target + if fn is torch._assert or getattr(fn, "__name__", "") == "_assert" or fn is operator.eq: + return # trace artifacts with no runtime effect + if fn is operator.getitem: + self._emit_getitem(node) + elif fn in self._BINARY_OPS: + self._add_simple_node(node, self._BINARY_OPS[fn], self._binop_inputs(node)) + elif fn in self._UNARY_OPS: + self._add_simple_node(node, self._UNARY_OPS[fn], [self._name_of(node.args[0])]) + elif fn is torch.transpose: + self._emit_transpose(node) + elif fn is torch.permute: + self._emit_permute(node) + elif fn is torch.cat: + tensors = [self._name_of(a) for a in node.args[0]] + dim = node.args[1] if len(node.args) > 1 else node.kwargs.get("dim", 0) + self._add_simple_node(node, "Concat", tensors, axis=int(dim)) + elif fn is torch.flatten: + self._emit_flatten(node, default_start_dim=0) + elif fn is torch.squeeze: + self._emit_squeeze(node) + elif fn is torch.unsqueeze: + self._emit_unsqueeze(node) + else: + raise TypeError(f"Unsupported call_function for FX ONNX export: {fn}") + + def _emit_call_method(self, node): + method = node.target + if method == "relu": + self._add_simple_node(node, "Relu", [self._name_of(node.args[0])]) + elif method == "flatten": + self._emit_flatten(node, default_start_dim=1) + elif method in ("view", "reshape"): + self._emit_reshape(node) + elif method == "transpose": + self._emit_transpose(node) + elif method == "permute": + self._emit_permute(node) + elif method == "matmul": + self._add_simple_node(node, "MatMul", self._binop_inputs(node)) + elif method == "squeeze": + self._emit_squeeze(node) + elif method == "unsqueeze": + self._emit_unsqueeze(node) + else: + raise TypeError(f"Unsupported call_method for FX ONNX export: {node.target!r}") + + def _collect_outputs(self, node): + ret = node.args[0] + rets = list(ret) if isinstance(ret, (tuple, list)) else [ret] + for r in rets: + if not isinstance(r, fx.Node): + raise TypeError("FX ONNX export: unsupported (non-tensor) model output") + val = self.node_to_name[r] + # MHA nodes store a tuple (out, avg_attn); expose the attention output. + self.output_names.append(val[0] if isinstance(val, tuple) else val) + + def _add_simple_node(self, node, op_type, inputs, **attrs): + out = f"{node.name}_{op_type.lower()}" + self.nodes.append(oh.make_node(op_type, inputs=inputs, outputs=[out], **attrs)) + self.node_to_name[node] = out + + def _emit_getitem(self, node): + container = self.node_to_name[node.args[0]] + if isinstance(container, tuple): + # Unpack a tuple output (e.g. from PQMultiheadAttention). + self.node_to_name[node] = container[node.args[1]] + else: + # Tensor slicing: x[:, 0], x[..., :4], ... → Slice (+ Squeeze) + rank = self._node_rank(node.args[0]) + self.node_to_name[node] = emit_getitem(node.name, container, node.args[1], rank, self.nodes, self.initializers) + + def _emit_transpose(self, node): + # torch.transpose(t, d0, d1) swaps two dims; ONNX needs a full perm. + perm = _swap_perm(self._node_rank(node.args[0]), int(node.args[1]), int(node.args[2])) + self._add_simple_node(node, "Transpose", [self._name_of(node.args[0])], perm=perm) + + def _emit_permute(self, node): + perm = _resolve_perm_dims(node.args[1:], self._node_rank(node.args[0])) + self._add_simple_node(node, "Transpose", [self._name_of(node.args[0])], perm=perm) + + def _emit_flatten(self, node, default_start_dim): + start_dim = node.args[1] if len(node.args) > 1 else node.kwargs.get("start_dim", default_start_dim) + self._add_simple_node(node, "Flatten", [self._name_of(node.args[0])], axis=int(start_dim)) + + def _emit_reshape(self, node): + shape_vals = [] + for a in node.args[1:]: + if not isinstance(a, int): + raise TypeError("Dynamic reshape (non-constant shape) is not supported in FX ONNX export") + shape_vals.append(a) + shape_name = add_int64_array(self.initializers, f"{node.name}_shape", shape_vals) + self._add_simple_node(node, "Reshape", [self._name_of(node.args[0]), shape_name]) + + def _emit_squeeze(self, node): + axes = self._squeeze_axes(node) + self.node_to_name[node] = emit_squeeze(node.name, self._name_of(node.args[0]), axes, self.nodes, self.initializers) + + def _squeeze_axes(self, node) -> list[int]: + """Resolve the axes a torch squeeze()/.squeeze() call removes.""" + in_shape = self._node_shape(node.args[0]) + if len(node.args) > 1 or "dim" in node.kwargs: + dim = int(node.args[1]) if len(node.args) > 1 else int(node.kwargs["dim"]) + dim %= len(in_shape) + return [dim] if in_shape[dim] == 1 else [] + return [i for i, s in enumerate(in_shape) if s == 1 and i != 0] + + def _emit_unsqueeze(self, node): + dim = int(node.args[1]) if len(node.args) > 1 else int(node.kwargs["dim"]) + axes = [dim % (self._node_rank(node.args[0]) + 1)] + self.node_to_name[node] = emit_unsqueeze(node.name, self._name_of(node.args[0]), axes, self.nodes, self.initializers) + + +def _prune_untranslatable_nodes(gm): + """Remove trace artifacts with no ONNX equivalent: assertions, dead comparisons, + and placeholders specialized away by concrete_args.""" + for n in reversed(list(gm.graph.find_nodes(op="call_function", target=torch._assert))): + gm.graph.erase_node(n) + for n in reversed(list(gm.graph.find_nodes(op="call_function", target=operator.eq))): + if len(n.users) == 0: + gm.graph.erase_node(n) + for n in reversed(list(gm.graph.find_nodes(op="placeholder"))): + if len(n.users) == 0 and len(n.args) > 0: # specialized: has a baked default, now unused + gm.graph.erase_node(n) + gm.recompile() + + +def _graph_input_names(gm, n_expected: int) -> dict: + """Map each tensor placeholder to its ONNX graph-input name.""" + placeholders = list(gm.graph.find_nodes(op="placeholder")) + if len(placeholders) != n_expected: + raise ValueError( + f"FX export: model.forward has {len(placeholders)} tensor input(s) but " + f"input_shape describes {n_expected}. Specialize non-tensor " + f"arguments via concrete_args={{...}}." + ) + # A single input keeps the graph-input name "input"; with multiple inputs + # each graph input is named after its forward parameter. + names = ["input"] if n_expected == 1 else [str(p.target) for p in placeholders] + return dict(zip(placeholders, names)) + + +def _route_input_passthrough_outputs(output_names, input_names, nodes): + """ONNX forbids a graph input from also being a graph output; insert Identity nodes.""" + graph_input_names = set(input_names) + for idx, name in enumerate(output_names): + if name in graph_input_names: + identity_out = f"{name}_identity_out{idx}" + nodes.append(oh.make_node("Identity", inputs=[name], outputs=[identity_out])) + output_names[idx] = identity_out + + +def convert_to_onnx( + model: nn.Module, + input_shape: tuple, + output_path: str = "model.onnx", + opset: int = 13, + use_qonnx: bool = False, + store_integer_weights: bool = False, + integer_ops: bool = False, + include_clip: bool = True, + concrete_args: dict | None = None, + input_dtypes=None, + batch_size: int | None = None, +) -> onnx.ModelProto: + """ + Convert a PQuant nn.Module to ONNX or QONNX using torch.fx symbolic tracing. + + Works with arbitrary model topologies including residual/skip connections, + branches, and concatenations. The model must be symbolically traceable + (no data-dependent control flow). + + Multiple inputs are supported: pass a sequence of per-input shapes as + ``input_shape`` (e.g. ``[(3, 32, 32), (16,)]``) and the model's ``forward`` + must take one tensor argument per shape, in the same order. A single input + keeps the graph-input name ``"input"``; with multiple inputs each graph input + is named after its ``forward`` parameter. + + Non-tensor inputs (bool flags, int sizes, ``None`` masks, ...) are not ONNX + graph inputs. Specialize them to constants at trace time by passing + ``concrete_args={"flag": False, ...}``; only the remaining tensor arguments + become graph inputs (see ``concrete_args`` below). + + Args: + model: Trained nn.Module. Call apply_final_compression() + on all PQ modules before passing here. + input_shape: Shape of a single sample (excluding batch), e.g. (3, 32, 32), + or a sequence of per-input shapes for multi-input models. + output_path: Where to save the .onnx file. + opset: ONNX opset version (≥13 required for per-channel DequantizeLinear). + use_qonnx: If True, emit QONNX Quant custom nodes (requires qonnx runtime). + If False (default), emit Clip+QuantizeLinear+DequantizeLinear + nodes runnable with plain onnxruntime. + store_integer_weights: If True (and use_qonnx=False), store weight/bias initializers + as int8/uint8 followed by DequantizeLinear instead of float32. + Ignored when use_qonnx=True or integer_ops=True. + integer_ops: If True (and use_qonnx=False), use MatMulInteger for Dense layers + so the inner product runs in int32 arithmetic. Weights are stored + as int8 (pre-transposed) and a single DequantizeLinear converts the + int32 accumulator back to float using the combined scale s_x * s_w. + Implies integer weight storage; store_integer_weights is ignored. + include_clip: Prepend a Clip node before each QuantizeLinear when True (default). + Set to False to emit bare QuantizeLinear+DequantizeLinear pairs — + safe when values are guaranteed in-range at inference time since + QuantizeLinear saturates naturally. Ignored when use_qonnx=True. + concrete_args: Forwarded to ``torch.fx.Tracer.trace`` to bake non-tensor + ``forward`` arguments in as constants. Keys are + ``forward`` parameter names. Specialized arguments are + dropped from the ONNX graph inputs. + input_dtypes: Optional dtype per input (single value or a list parallel + to ``input_shape``). Each is a torch.dtype or a string + (``"float32"``, ``"bool"``, ``"int64"``, ``"int32"``). + Defaults to float32. Use ``"bool"`` for a runtime + attention ``key_padding_mask`` input, for example. + batch_size: If not None, fix the batch dimension of every graph input + and output to this value. If None (default), the batch + dimension is left dynamic. + + Returns: + The constructed onnx.ModelProto. + """ + model.eval() + quant_fn = quant_node if use_qonnx else functools.partial(qdq_node, include_clip=include_clip) + + input_shapes = _normalize_input_shapes(input_shape) + input_torch_dtypes, input_tp_dtypes = _normalize_input_dtypes(input_dtypes, len(input_shapes)) + + gm = fx.GraphModule(model, PQTracer().trace(model, concrete_args=concrete_args)) + _prune_untranslatable_nodes(gm) + ph_to_name = _graph_input_names(gm, len(input_shapes)) + input_names = list(ph_to_name.values()) + + device = next((p.device for p in model.parameters()), None) + probes = [torch.zeros(1, *shp, device=device, dtype=dt) for shp, dt in zip(input_shapes, input_torch_dtypes)] + with torch.no_grad(): + ShapeProp(gm).propagate(*probes) + + emitter = FxGraphEmitter(gm, ph_to_name, quant_fn, use_qonnx, store_integer_weights, integer_ops) + output_names = emitter._run() + _route_input_passthrough_outputs(output_names, input_names, emitter.nodes) + + with torch.no_grad(): + dummy_out = model(*probes, **(concrete_args or {})) + dummy_outs = list(dummy_out) if isinstance(dummy_out, (tuple, list)) else [dummy_out] + + batch_dim = batch_size # None → dynamic, int → fixed + input_vis = [ + oh.make_tensor_value_info(name, tp, [batch_dim, *shp]) + for name, shp, tp in zip(input_names, input_shapes, input_tp_dtypes) + ] + output_vis = [ + oh.make_tensor_value_info(name, TensorProto.FLOAT, [batch_dim] + list(t.shape[1:])) + for name, t in zip(output_names, dummy_outs) + ] + + graph = oh.make_graph( + nodes=emitter.nodes, + name="pquant_onnx_fx", + inputs=input_vis, + outputs=output_vis, + initializer=emitter.initializers, + ) + model_proto = save_model(graph, output_path, opset, use_qonnx=use_qonnx) + logging.info("Saved %s model (FX) → %s", "QONNX" if use_qonnx else "ONNX (QDQ)", output_path) + return model_proto diff --git a/src/pquant/core/torch/onnx/layer_builders.py b/src/pquant/core/torch/onnx/layer_builders.py new file mode 100644 index 0000000..fb33b7a --- /dev/null +++ b/src/pquant/core/torch/onnx/layer_builders.py @@ -0,0 +1,464 @@ +"""Per-layer ONNX graph builders (Dense/Conv/BN/LN/Pool/Activation/MHA) for the PQuant torch converter.""" + +import numpy as np +import onnx.helper as oh + +from pquant.core.onnx_common import ( + add_float_scalar, + add_initializer, + add_transpose, + conv_padding_attrs, + emit_mha_core, + emit_param, + fixed_point_clip_range, + maybe_quant_input, + maybe_quant_output, + qdq_node, + symmetric_pads, + to_list, + to_np, +) + + +def _quantize_input_to_int(quantizer, prefix, current, nodes, initializers): + """Clip + QuantizeLinear the input to int8/uint8, stopping before DequantizeLinear. + + Returns (int_tensor_name, zero_point_name, input_scale). + """ + k, i, f = quantizer.get_quantization_bits() + signed = int(to_np(k).ravel()[0]) == 1 + i_val = float(to_np(i).ravel()[0]) + f_val = float(to_np(f).ravel()[0]) + scale = float(2.0 ** (-f_val)) + clip_min, clip_max = fixed_point_clip_range(signed, i_val, f_val, "SAT") + + clip_min_name = add_float_scalar(initializers, f"{prefix}_in_clip_min", clip_min) + clip_max_name = add_float_scalar(initializers, f"{prefix}_in_clip_max", clip_max) + scale_name = add_float_scalar(initializers, f"{prefix}_in_scale", scale) + zp_name = add_initializer(initializers, f"{prefix}_in_zp", np.array(np.int8(0) if signed else np.uint8(0))) + + clipped_name = f"{prefix}_in_clipped" + int_name = f"{prefix}_in_int" + nodes.append(oh.make_node("Clip", inputs=[current, clip_min_name, clip_max_name], outputs=[clipped_name])) + nodes.append(oh.make_node("QuantizeLinear", inputs=[clipped_name, scale_name, zp_name], outputs=[int_name])) + return int_name, zp_name, scale + + +def _integer_weights_transposed(module, prefix, initializers): + """Quantize the dense weight to int8/uint8, pre-transposed to [in, out] so + MatMulInteger needs no runtime Transpose node. + + Returns (weight_name, weight_zp_name, scale_1d, per_channel) where scale_1d + has shape [1] (per-tensor) or [out] (per-channel). + """ + weight_np = to_np(module._weight) # PyTorch layout: [out, in] + k, _, f = module.weight_quantizer.get_quantization_bits() + dtype = np.int8 if int(to_np(k).ravel()[0]) == 1 else np.uint8 + out_channels = weight_np.shape[0] + + f_np = to_np(f) + if f_np.size == 1: + f_1d = np.array([float(f_np.ravel()[0])]) + per_channel = False + else: + f_1d = f_np.reshape(out_channels, -1).min(axis=1) # min f → max scale → covers all values + per_channel = True + + scale_1d = (2.0 ** (-f_1d)).astype(np.float32) + broadcast = scale_1d.reshape((out_channels,) + (1,) * (weight_np.ndim - 1)) if per_channel else float(scale_1d[0]) + int_weights = np.round(weight_np / broadcast).astype(dtype).T # [in, out] + + weight_name = add_initializer(initializers, f"{prefix}_weight_int", int_weights) + zp_name = add_initializer(initializers, f"{prefix}_weight_zp", np.array(dtype(0))) # zero for symmetric quantization + return weight_name, zp_name, scale_1d, per_channel + + +def _dequantize_accumulator(prefix, current, combined_scale_1d, per_channel, nodes, initializers): + """DequantizeLinear the int32 accumulator back to float32 with the combined scale s_x * s_w. + + Per-channel: axis=-1 because out features are the last axis of the accumulator + ([batch, out] for rank-2 inputs, [batch, ..., out] otherwise). + """ + if per_channel: + scale_np = combined_scale_1d.astype(np.float32) + zp_np = np.zeros(len(combined_scale_1d), dtype=np.int32) + dql_kwargs = {"axis": -1} + else: + scale_np = np.array(float(combined_scale_1d[0]), dtype=np.float32) + zp_np = np.array(np.int32(0)) + dql_kwargs = {} + + scale_name = add_initializer(initializers, f"{prefix}_combined_scale", scale_np) + zp_name = add_initializer(initializers, f"{prefix}_combined_zp", zp_np) + out = f"{prefix}_dequantized" + nodes.append(oh.make_node("DequantizeLinear", inputs=[current, scale_name, zp_name], outputs=[out], **dql_kwargs)) + return out + + +def _add_dense_integer(module, prefix, current, nodes, initializers): + """Dense layer whose inner product runs in int32 via MatMulInteger.""" + if getattr(module, "input_quantizer", None) is None or not module.quantize_input: + raise ValueError(f"{prefix}: integer_ops requires quantize_input=True on the layer") + + x_int, x_zp, input_scale = _quantize_input_to_int(module.input_quantizer, prefix, current, nodes, initializers) + w_int, w_zp, weight_scale_1d, per_channel = _integer_weights_transposed(module, prefix, initializers) + combined_scale_1d = input_scale * weight_scale_1d + + current = f"{prefix}_matmul_int" # MatMulInteger([batch, in], [in, out]) → int32 [batch, out] + nodes.append(oh.make_node("MatMulInteger", inputs=[x_int, w_int, x_zp, w_zp], outputs=[current])) + + if module._bias is not None: + bias_scale = combined_scale_1d if per_channel else float(combined_scale_1d[0]) + bias_int32 = np.round(to_np(module._bias) / bias_scale).astype(np.int32) + bias_name = add_initializer(initializers, f"{prefix}_bias_int", bias_int32) + biased_name = f"{prefix}_matmul_biased" + nodes.append(oh.make_node("Add", inputs=[current, bias_name], outputs=[biased_name])) + current = biased_name + + current = _dequantize_accumulator(prefix, current, combined_scale_1d, per_channel, nodes, initializers) + + # Optional output quantization (e.g. last layer with quantize_output=True) + return maybe_quant_output(module, prefix, current, nodes, initializers, qdq_node) + + +def _add_dense_nd(module, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights): + """Dense layer as MatMul + Add, for inputs of rank > 2 (Gemm only takes rank-2).""" + current = maybe_quant_input(module, prefix, current, nodes, initializers, quant_fn) + + weight_np = to_np(module._weight) # [out, in] + if use_qonnx or store_integer_weights: + # Quantized/int-stored weight is emitted in native [out, in] layout, then transposed. + q_weight_native = emit_param( + prefix, "weight", weight_np, module.weight_quantizer, nodes, initializers, use_qonnx, store_integer_weights + ) + q_weight = f"{prefix}_weight_T" + nodes.append(oh.make_node("Transpose", inputs=[q_weight_native], outputs=[q_weight], perm=[1, 0])) + else: + q_weight = f"{prefix}_weight_T" + add_initializer(initializers, q_weight, weight_np.T) # pre-transposed [in, out] + + matmul_out = f"{prefix}_matmul" + nodes.append(oh.make_node("MatMul", inputs=[current, q_weight], outputs=[matmul_out])) + current = matmul_out + + if module._bias is not None: + q_bias = emit_param( + prefix, "bias", to_np(module._bias), module.bias_quantizer, nodes, initializers, use_qonnx, store_integer_weights + ) + biased_out = f"{prefix}_biased" + nodes.append(oh.make_node("Add", inputs=[matmul_out, q_bias], outputs=[biased_out])) + current = biased_out + + return maybe_quant_output(module, prefix, current, nodes, initializers, quant_fn) + + +def _add_dense(module, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights, integer_ops=False): + if integer_ops and not use_qonnx: + return _add_dense_integer(module, prefix, current, nodes, initializers) + current = maybe_quant_input(module, prefix, current, nodes, initializers, quant_fn) + + q_weight = emit_param( + prefix, + "weight", + to_np(module._weight), + module.weight_quantizer, + nodes, + initializers, + use_qonnx, + store_integer_weights, + ) + gemm_inputs = [current, q_weight] + + if module._bias is not None: + gemm_inputs.append( + emit_param( + prefix, + "bias", + to_np(module._bias), + module.bias_quantizer, + nodes, + initializers, + use_qonnx, + store_integer_weights, + ) + ) + + gemm_out = f"{prefix}_gemm" + nodes.append(oh.make_node("Gemm", inputs=gemm_inputs, outputs=[gemm_out], transB=1)) + + return maybe_quant_output(module, prefix, gemm_out, nodes, initializers, quant_fn) + + +def _add_conv(module, prefix, current, nodes, initializers, ndim, quant_fn, use_qonnx, store_integer_weights): + current = maybe_quant_input(module, prefix, current, nodes, initializers, quant_fn) + + q_weight = emit_param( + prefix, + "weight", + to_np(module._weight), + module.weight_quantizer, + nodes, + initializers, + use_qonnx, + store_integer_weights, + ) + conv_inputs = [current, q_weight] + + if module._bias is not None: + conv_inputs.append( + emit_param( + prefix, + "bias", + to_np(module._bias), + module.bias_quantizer, + nodes, + initializers, + use_qonnx, + store_integer_weights, + ) + ) + + auto_pad, pads = conv_padding_attrs(module.padding, ndim) + conv_attrs = dict( + kernel_shape=to_list(module.kernel_size, ndim), + strides=to_list(module.stride, ndim), + dilations=to_list(module.dilation, ndim), + group=module.groups, + auto_pad=auto_pad, + ) + if pads is not None: + conv_attrs["pads"] = pads + + conv_out = f"{prefix}_conv" + nodes.append(oh.make_node("Conv", inputs=conv_inputs, outputs=[conv_out], **conv_attrs)) + + return maybe_quant_output(module, prefix, conv_out, nodes, initializers, quant_fn) + + +def _add_batchnorm(module, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights): + current = maybe_quant_input(module, prefix, current, nodes, initializers, quant_fn) + + q_gamma = emit_param( + prefix, + "gamma", + to_np(module._weight), + module.weight_quantizer, + nodes, + initializers, + use_qonnx, + store_integer_weights, + ) + q_beta = emit_param( + prefix, "beta", to_np(module._bias), module.bias_quantizer, nodes, initializers, use_qonnx, store_integer_weights + ) + mean_name = add_initializer(initializers, f"{prefix}_running_mean", to_np(module.running_mean)) + var_name = add_initializer(initializers, f"{prefix}_running_var", to_np(module.running_var)) + + bn_out = f"{prefix}_bn" + nodes.append( + oh.make_node( + "BatchNormalization", + inputs=[current, q_gamma, q_beta, mean_name, var_name], + outputs=[bn_out], + epsilon=float(module.eps), + ) + ) + return bn_out + + +def _add_layernorm(module, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights): + current = maybe_quant_input(module, prefix, current, nodes, initializers, quant_fn) + + normalized_shape = tuple(to_list(module.normalized_shape, 1)) + axis = -len(normalized_shape) + + has_weight = module._weight is not None + has_bias = module._bias is not None + # elementwise_affine=False has no quantizers; emit plain float parameters. + use_qonnx = use_qonnx and has_weight + store_integer_weights = store_integer_weights and has_weight + + gamma_np = to_np(module._weight) if has_weight else np.ones(normalized_shape, dtype=np.float32) + quantizer = module.weight_quantizer if has_weight else None + q_gamma = emit_param(prefix, "gamma", gamma_np, quantizer, nodes, initializers, use_qonnx, store_integer_weights) + + ln_inputs = [current, q_gamma] + if has_bias: + bias_quantizer = module.bias_quantizer if has_weight else None + q_beta = emit_param( + prefix, "beta", to_np(module._bias), bias_quantizer, nodes, initializers, use_qonnx, store_integer_weights + ) + ln_inputs.append(q_beta) + + ln_out = f"{prefix}_ln" + nodes.append( + oh.make_node("LayerNormalization", inputs=ln_inputs, outputs=[ln_out], axis=axis, epsilon=float(module.eps)) + ) + return maybe_quant_output(module, prefix, ln_out, nodes, initializers, quant_fn) + + +def _add_avgpool(module, prefix, current, nodes, initializers, ndim, quant_fn): + current = maybe_quant_input(module, prefix, current, nodes, initializers, quant_fn) + + pool_out = f"{prefix}_pool" + nodes.append( + oh.make_node( + "AveragePool", + inputs=[current], + outputs=[pool_out], + kernel_shape=to_list(module.kernel_size, ndim), + strides=to_list(module.stride, ndim), + pads=symmetric_pads(module.padding, ndim), + ceil_mode=int(module.ceil_mode), + count_include_pad=int(module.count_include_pad), + ) + ) + return maybe_quant_output(module, prefix, pool_out, nodes, initializers, quant_fn) + + +def _add_maxpool(module, prefix, current, nodes): + out = f"{prefix}_maxpool" + nodes.append( + oh.make_node( + "MaxPool", + inputs=[current], + outputs=[out], + kernel_shape=to_list(module.kernel_size, 2), + strides=to_list(module.stride, 2), + pads=symmetric_pads(module.padding, 2), + ) + ) + return out + + +def _add_upsample(module, prefix, current, nodes, initializers): + """Emit a Resize node with nearest/linear mode and constant scale factors.""" + roi_name = add_initializer(initializers, f"{prefix}_upsample_roi", np.array([], dtype=np.float32)) + scale_factor = module.scale_factor + if isinstance(scale_factor, (int, float)): + scale_factor = (scale_factor, scale_factor) + scales = np.array([1.0, 1.0, float(scale_factor[0]), float(scale_factor[1])], dtype=np.float32) + scales_name = add_initializer(initializers, f"{prefix}_upsample_scales", scales) + mode = "nearest" if module.mode == "nearest" else "linear" + + out = f"{prefix}_upsample" + nodes.append( + oh.make_node( + "Resize", + inputs=[current, roi_name, scales_name], + outputs=[out], + mode=mode, + coordinate_transformation_mode="asymmetric", + ) + ) + return out + + +def _add_activation(module, prefix, current, nodes, initializers, quant_fn): + """PQActivation: optional input quantization, the activation itself, optional output quantization.""" + current = maybe_quant_input(module, prefix, current, nodes, initializers, quant_fn) + + activation = module.activation_name + act_out = f"{prefix}_act" + if activation == "relu": + nodes.append(oh.make_node("Relu", inputs=[current], outputs=[act_out])) + elif activation == "tanh": + nodes.append(oh.make_node("Tanh", inputs=[current], outputs=[act_out])) + elif activation == "hard_tanh": + cmin_name = add_float_scalar(initializers, f"{prefix}_htanh_min", -1.0) + cmax_name = add_float_scalar(initializers, f"{prefix}_htanh_max", 1.0) + nodes.append(oh.make_node("Clip", inputs=[current, cmin_name, cmax_name], outputs=[act_out])) + elif activation == "leaky_relu": + alpha = module.activation_function.negative_slope + nodes.append(oh.make_node("LeakyRelu", inputs=[current], outputs=[act_out], alpha=alpha)) + elif activation == "gelu": + _add_gelu(module, prefix, current, act_out, nodes, initializers) + else: + raise TypeError(f"PQActivation: unsupported activation {activation!r} for ONNX export") + + return maybe_quant_output(module, prefix, act_out, nodes, initializers, quant_fn) + + +def _add_gelu(module, prefix, current, act_out, nodes, initializers): + """Decompose gelu so the default opset (13) works; ONNX added a Gelu op only in opset 20.""" + approximate = getattr(module.activation_function, "approximate", "none") + half_name = add_float_scalar(initializers, f"{prefix}_gelu_half", 0.5) + one_name = add_float_scalar(initializers, f"{prefix}_gelu_one", 1.0) + plus_one = f"{prefix}_gelu_plus1" + + if approximate == "tanh": + # 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))) + sqrt_2_over_pi_name = add_float_scalar(initializers, f"{prefix}_gelu_sqrt2_over_pi", np.sqrt(2.0 / np.pi)) + cubic_coeff_name = add_float_scalar(initializers, f"{prefix}_gelu_c1", 0.044715) + three_name = add_float_scalar(initializers, f"{prefix}_gelu_three", 3.0) + x3 = f"{prefix}_gelu_x3" + cx3 = f"{prefix}_gelu_cx3" + inner = f"{prefix}_gelu_inner" + scaled = f"{prefix}_gelu_scaled" + tanh_out = f"{prefix}_gelu_tanh" + nodes += [ + oh.make_node("Pow", inputs=[current, three_name], outputs=[x3]), + oh.make_node("Mul", inputs=[x3, cubic_coeff_name], outputs=[cx3]), + oh.make_node("Add", inputs=[current, cx3], outputs=[inner]), + oh.make_node("Mul", inputs=[inner, sqrt_2_over_pi_name], outputs=[scaled]), + oh.make_node("Tanh", inputs=[scaled], outputs=[tanh_out]), + oh.make_node("Add", inputs=[tanh_out, one_name], outputs=[plus_one]), + ] + else: + # Exact: 0.5 * x * (1 + erf(x / sqrt(2))) + inv_sqrt2_name = add_float_scalar(initializers, f"{prefix}_gelu_inv_sqrt2", 1.0 / np.sqrt(2.0)) + scaled = f"{prefix}_gelu_scaled" + erf_out = f"{prefix}_gelu_erf" + nodes += [ + oh.make_node("Mul", inputs=[current, inv_sqrt2_name], outputs=[scaled]), + oh.make_node("Erf", inputs=[scaled], outputs=[erf_out]), + oh.make_node("Add", inputs=[erf_out, one_name], outputs=[plus_one]), + ] + + x_times = f"{prefix}_gelu_xprod" + nodes += [ + oh.make_node("Mul", inputs=[current, plus_one], outputs=[x_times]), + oh.make_node("Mul", inputs=[x_times, half_name], outputs=[act_out]), + ] + + +def _add_mha( + module, + prefix, + q_input, + k_input, + v_input, + nodes, + initializers, + quant_fn, + use_qonnx, + store_integer_weights, + key_padding_mask=None, + attn_mask=None, +): + if not module.batch_first: + q_input = add_transpose(f"{prefix}_q_in", q_input, [1, 0, 2], nodes) + k_input = add_transpose(f"{prefix}_k_in", k_input, [1, 0, 2], nodes) + v_input = add_transpose(f"{prefix}_v_in", v_input, [1, 0, 2], nodes) + + # Q / K / V projections: (B, L, E) → (B, L, E) via MatMul (input is rank-3) + q_proj_out = _add_dense_nd( + module.q_proj, f"{prefix}_q_proj", q_input, nodes, initializers, quant_fn, use_qonnx, store_integer_weights + ) + k_proj_out = _add_dense_nd( + module.k_proj, f"{prefix}_k_proj", k_input, nodes, initializers, quant_fn, use_qonnx, store_integer_weights + ) + v_proj_out = _add_dense_nd( + module.v_proj, f"{prefix}_v_proj", v_input, nodes, initializers, quant_fn, use_qonnx, store_integer_weights + ) + + context, avg_attn = emit_mha_core( + module, prefix, q_proj_out, k_proj_out, v_proj_out, nodes, initializers, quant_fn, key_padding_mask, attn_mask + ) + out = _add_dense_nd( + module.out_proj, f"{prefix}_out_proj", context, nodes, initializers, quant_fn, use_qonnx, store_integer_weights + ) + + if not module.batch_first: + out = add_transpose(f"{prefix}_out_seq_first", out, [1, 0, 2], nodes) + return out, avg_attn diff --git a/tests/test_keras_onnx_converter.py b/tests/test_keras_onnx_converter.py index 91dc46f..01a562e 100644 --- a/tests/test_keras_onnx_converter.py +++ b/tests/test_keras_onnx_converter.py @@ -10,52 +10,48 @@ import keras import numpy as np +import onnxruntime as ort import pytest import pquant -from pquant.core.keras.convert_to_onnx import convert_to_onnx from pquant.core.keras.layers import ( + PQActivation, PQBatchNormalization, PQConv1d, PQConv2d, PQDense, PQDepthwiseConv2d, + PQMultiheadAttention, apply_final_compression, ) - -ort = pytest.importorskip("onnxruntime", reason="onnxruntime not installed") +from pquant.core.keras.onnx import convert_to_onnx ATOL = 1e-4 +QUANT_ATOL = 5e-3 -# --------------------------------------------------------------------------- -# fixtures -# --------------------------------------------------------------------------- +def atol(cfg): + return QUANT_ATOL if cfg.quantization_parameters.enable_quantization else ATOL -@pytest.fixture -def cfg(): +@pytest.fixture(params=[False, True], ids=["float", "quant"]) +def cfg(request): c = pquant.cs_config() - c.quantization_parameters.enable_quantization = False + c.quantization_parameters.enable_quantization = request.param return c -# --------------------------------------------------------------------------- -# helpers -# --------------------------------------------------------------------------- - - -def _channels_first(): +def channels_first(): return keras.backend.image_data_format() == "channels_first" -def _keras_out(model, x: np.ndarray) -> np.ndarray: +def keras_out(model, x: np.ndarray) -> np.ndarray: from keras import ops return ops.convert_to_numpy(model(x, training=False)) -def _onnx_run(model, x: np.ndarray, input_shape: tuple, tmp_path) -> np.ndarray: +def onnx_run(model, x: np.ndarray, input_shape: tuple, tmp_path) -> np.ndarray: path = str(tmp_path / "model.onnx") convert_to_onnx(model, input_shape=input_shape, output_path=path) sess = ort.InferenceSession(path) @@ -63,135 +59,274 @@ def _onnx_run(model, x: np.ndarray, input_shape: tuple, tmp_path) -> np.ndarray: return sess.run(None, {in_name: x})[0] -# --------------------------------------------------------------------------- -# PQDense -# --------------------------------------------------------------------------- - +# (layer factory, channels, spatial dims, batch size, warm-up call kwargs) +SINGLE_LAYER_CASES = [ + pytest.param(lambda cfg: PQDense(cfg, units=8, use_bias=True), 16, (), 4, {}, id="dense-bias"), + pytest.param(lambda cfg: PQDense(cfg, units=8, use_bias=False), 16, (), 4, {}, id="dense-nobias"), + pytest.param( + lambda cfg: PQConv2d(cfg, 8, kernel_size=3, padding="same", use_bias=True), 3, (8, 8), 2, {}, id="conv2d-bias" + ), + pytest.param( + lambda cfg: PQConv2d(cfg, 8, kernel_size=3, padding="same", use_bias=False), 3, (8, 8), 2, {}, id="conv2d-nobias" + ), + pytest.param( + lambda cfg: PQConv1d(cfg, 8, kernel_size=3, padding="same", use_bias=True), 4, (16,), 2, {}, id="conv1d-bias" + ), + pytest.param( + lambda cfg: PQConv1d(cfg, 8, kernel_size=3, padding="same", use_bias=False), 4, (16,), 2, {}, id="conv1d-nobias" + ), + pytest.param( + lambda cfg: PQBatchNormalization(cfg, axis=1 if channels_first() else -1), + 8, + (4, 4), + 4, + {"training": True}, # warm up running stats + id="batchnorm", + ), + pytest.param(lambda cfg: PQDepthwiseConv2d(cfg, kernel_size=3, padding="same"), 4, (8, 8), 2, {}, id="depthwise_conv2d"), +] + + +@pytest.mark.parametrize("make_layer,channels,spatial,batch,warmup_kwargs", SINGLE_LAYER_CASES) +def test_single_layer_onnx(cfg, make_layer, channels, spatial, batch, warmup_kwargs, tmp_path): + input_shape = (channels, *spatial) if channels_first() else (*spatial, channels) + x_np = np.random.randn(batch, *input_shape).astype(np.float32) -@pytest.mark.parametrize("bias", [True, False]) -def test_dense_onnx(cfg, bias, tmp_path): - IN, OUT = 16, 8 - inputs = keras.Input(shape=(IN,)) - x = PQDense(cfg, units=OUT, use_bias=bias)(inputs) + inputs = keras.Input(shape=input_shape) + x = make_layer(cfg)(inputs) model = keras.Model(inputs, x) - dummy = np.zeros((1, IN), dtype=np.float32) - model(dummy) + model(np.zeros((1, *input_shape), dtype=np.float32), **warmup_kwargs) apply_final_compression(model) - x_np = np.random.randn(4, IN).astype(np.float32) - keras_out = _keras_out(model, x_np) - onnx_out = _onnx_run(model, x_np, input_shape=(IN,), tmp_path=tmp_path) - np.testing.assert_allclose(keras_out, onnx_out, atol=ATOL, err_msg=f"PQDense bias={bias}: keras vs ONNX mismatch") + keras_output = keras_out(model, x_np) + onnx_out = onnx_run(model, x_np, input_shape=input_shape, tmp_path=tmp_path) + np.testing.assert_allclose(keras_output, onnx_out, atol=atol(cfg), err_msg="keras vs ONNX mismatch") -# --------------------------------------------------------------------------- -# PQConv2d -# --------------------------------------------------------------------------- +@pytest.mark.parametrize("activation", ["relu", "tanh", "hard_tanh"]) +def test_pqactivation_onnx(cfg, activation, tmp_path): + DIM = 16 + inputs = keras.Input(shape=(DIM,)) + x = PQActivation(cfg, activation)(inputs) + model = keras.Model(inputs, x) + model(np.zeros((1, DIM), dtype=np.float32)) + apply_final_compression(model) -@pytest.mark.parametrize("bias", [True, False]) -def test_conv2d_onnx(cfg, bias, tmp_path): - IN_C, OUT_C, H, W = 3, 8, 8, 8 - if _channels_first(): - input_shape = (IN_C, H, W) - x_np = np.random.randn(2, IN_C, H, W).astype(np.float32) - else: - input_shape = (H, W, IN_C) - x_np = np.random.randn(2, H, W, IN_C).astype(np.float32) + x_np = np.random.randn(4, DIM).astype(np.float32) + keras_output = keras_out(model, x_np) + onnx_out = onnx_run(model, x_np, input_shape=(DIM,), tmp_path=tmp_path) + np.testing.assert_allclose( + keras_output, onnx_out, atol=atol(cfg), err_msg=f"PQActivation {activation}: keras vs ONNX mismatch" + ) + + +def test_residual_concat_onnx(cfg, tmp_path): + DIM = 16 + inputs = keras.Input(shape=(DIM,)) + h = PQDense(cfg, units=DIM)(inputs) + h2 = PQDense(cfg, units=DIM)(h) + add = keras.layers.Add()([h, h2]) # residual / skip add + cat = keras.layers.Concatenate(axis=-1)([add, inputs]) # branch merge + out = PQDense(cfg, units=8)(cat) + model = keras.Model(inputs, out) + + model(np.zeros((1, DIM), dtype=np.float32)) + apply_final_compression(model) - inputs = keras.Input(shape=input_shape) - x = PQConv2d(cfg, OUT_C, kernel_size=3, padding="same", use_bias=bias)(inputs) - model = keras.Model(inputs, x) + x_np = np.random.randn(4, DIM).astype(np.float32) + keras_output = keras_out(model, x_np) + onnx_out = onnx_run(model, x_np, input_shape=(DIM,), tmp_path=tmp_path) + np.testing.assert_allclose(keras_output, onnx_out, atol=atol(cfg), err_msg="residual+concat: keras vs ONNX mismatch") - dummy = np.zeros((1, *input_shape), dtype=np.float32) - model(dummy) + +def test_two_input_onnx(cfg, tmp_path): + IN_A, IN_B, OUT = 16, 4, 8 + a = keras.Input(shape=(IN_A,)) + b = keras.Input(shape=(IN_B,)) + ha = PQDense(cfg, units=OUT)(a) + hb = PQDense(cfg, units=OUT)(b) + out = keras.layers.Add()([ha, hb]) + model = keras.Model([a, b], out) + + model([np.zeros((1, IN_A), np.float32), np.zeros((1, IN_B), np.float32)]) apply_final_compression(model) - keras_out = _keras_out(model, x_np) - onnx_out = _onnx_run(model, x_np, input_shape=input_shape, tmp_path=tmp_path) - np.testing.assert_allclose(keras_out, onnx_out, atol=ATOL, err_msg=f"PQConv2d bias={bias}: keras vs ONNX mismatch") + xa = np.random.randn(3, IN_A).astype(np.float32) + xb = np.random.randn(3, IN_B).astype(np.float32) + keras_output = keras_out(model, [xa, xb]) + path = str(tmp_path / "two_input.onnx") + proto = convert_to_onnx(model, input_shape=[(IN_A,), (IN_B,)], output_path=path) + in_shapes = {i.name: [d.dim_value for d in i.type.tensor_type.shape.dim] for i in proto.graph.input} + assert in_shapes["input_0"] == [0, IN_A] # dim_value 0 == dynamic batch + assert in_shapes["input_1"] == [0, IN_B] -# --------------------------------------------------------------------------- -# PQConv1d -# --------------------------------------------------------------------------- + sess = ort.InferenceSession(path) + names = [i.name for i in sess.get_inputs()] + onnx_out = sess.run(None, {names[0]: xa, names[1]: xb})[0] + np.testing.assert_allclose(keras_output, onnx_out, atol=atol(cfg), err_msg="two-input: keras vs ONNX mismatch") @pytest.mark.parametrize("bias", [True, False]) -def test_conv1d_onnx(cfg, bias, tmp_path): - IN_C, OUT_C, L = 4, 8, 16 - if _channels_first(): - input_shape = (IN_C, L) - x_np = np.random.randn(2, IN_C, L).astype(np.float32) - else: - input_shape = (L, IN_C) - x_np = np.random.randn(2, L, IN_C).astype(np.float32) +def test_mha_onnx(cfg, bias, tmp_path): + E, H, T = 16, 4, 8 + inputs = keras.Input(shape=(T, E)) + mha = PQMultiheadAttention(cfg, embed_dim=E, num_heads=H, bias=bias) + out, _ = mha([inputs, inputs, inputs]) + model = keras.Model(inputs, out) + + model(np.zeros((1, T, E), dtype=np.float32)) + apply_final_compression(model) - inputs = keras.Input(shape=input_shape) - x = PQConv1d(cfg, OUT_C, kernel_size=3, padding="same", use_bias=bias)(inputs) - model = keras.Model(inputs, x) + x_np = np.random.randn(2, T, E).astype(np.float32) + keras_output = keras_out(model, x_np) + onnx_out = onnx_run(model, x_np, input_shape=(T, E), tmp_path=tmp_path) + np.testing.assert_allclose( + keras_output, onnx_out, atol=atol(cfg), err_msg=f"PQMultiheadAttention bias={bias}: keras vs ONNX mismatch" + ) - dummy = np.zeros((1, *input_shape), dtype=np.float32) - model(dummy) - apply_final_compression(model) - keras_out = _keras_out(model, x_np) - onnx_out = _onnx_run(model, x_np, input_shape=input_shape, tmp_path=tmp_path) - np.testing.assert_allclose(keras_out, onnx_out, atol=ATOL, err_msg=f"PQConv1d bias={bias}: keras vs ONNX mismatch") +def test_mha_causal_attn_mask_onnx(cfg, tmp_path): + E, H, T = 16, 4, 8 + inputs = keras.Input(shape=(T, E)) + mha = PQMultiheadAttention(cfg, embed_dim=E, num_heads=H) + # (T, S) additive causal mask: 0 on/below the diagonal, large-negative above it. + attn_mask = np.triu(np.full((T, T), -1e4, dtype=np.float32), k=1) + out, _ = mha([inputs, inputs, inputs], attn_mask=attn_mask) + model = keras.Model(inputs, out) + model(np.zeros((1, T, E), dtype=np.float32)) + apply_final_compression(model) -# --------------------------------------------------------------------------- -# PQBatchNormalization -# --------------------------------------------------------------------------- + x_np = np.random.randn(2, T, E).astype(np.float32) + keras_output = keras_out(model, x_np) + onnx_out = onnx_run(model, x_np, input_shape=(T, E), tmp_path=tmp_path) + np.testing.assert_allclose( + keras_output, onnx_out, atol=atol(cfg), err_msg="MHA causal attn_mask: keras vs ONNX mismatch" + ) -def test_batchnorm_onnx(cfg, tmp_path): - IN_C, H, W = 8, 4, 4 - if _channels_first(): - input_shape = (IN_C, H, W) - x_np = np.random.randn(4, IN_C, H, W).astype(np.float32) - bn_axis = 1 - else: - input_shape = (H, W, IN_C) - x_np = np.random.randn(4, H, W, IN_C).astype(np.float32) - bn_axis = -1 +def test_mha_key_padding_mask_onnx(cfg, tmp_path): + import onnx - inputs = keras.Input(shape=input_shape) - x = PQBatchNormalization(cfg, axis=bn_axis)(inputs) - model = keras.Model(inputs, x) + E, H, T = 16, 4, 8 + inputs = keras.Input(shape=(T, E)) + kpm = keras.Input(shape=(T,), dtype="bool") # runtime bool padding mask, True == padding + mha = PQMultiheadAttention(cfg, embed_dim=E, num_heads=H) + out, _ = mha([inputs, inputs, inputs], key_padding_mask=kpm) + model = keras.Model([inputs, kpm], out) - dummy = np.zeros((1, *input_shape), dtype=np.float32) - model(dummy, training=True) # warm up running stats + model([np.zeros((1, T, E), np.float32), np.zeros((1, T), bool)]) apply_final_compression(model) - keras_out = _keras_out(model, x_np) - onnx_out = _onnx_run(model, x_np, input_shape=input_shape, tmp_path=tmp_path) - np.testing.assert_allclose(keras_out, onnx_out, atol=ATOL, err_msg="PQBatchNormalization: keras vs ONNX mismatch") + x_np = np.random.randn(2, T, E).astype(np.float32) + mask_np = np.zeros((2, T), dtype=bool) + mask_np[:, -2:] = True # last two key positions are padding + keras_output = keras_out(model, [x_np, mask_np]) + path = str(tmp_path / "mha_kpm.onnx") + proto = convert_to_onnx(model, input_shape=[(T, E), (T,)], output_path=path) + # The padding mask must be a genuine bool graph input, not baked away. + kpm_vi = next(i for i in proto.graph.input if i.name == "input_1") + assert kpm_vi.type.tensor_type.elem_type == onnx.TensorProto.BOOL -# --------------------------------------------------------------------------- -# PQDepthwiseConv2d -# --------------------------------------------------------------------------- - + sess = ort.InferenceSession(path) + names = [i.name for i in sess.get_inputs()] + onnx_out = sess.run(None, {names[0]: x_np, names[1]: mask_np})[0] + np.testing.assert_allclose( + keras_output, onnx_out, atol=atol(cfg), err_msg="MHA key_padding_mask: keras vs ONNX mismatch" + ) + + +@pytest.mark.parametrize( + "slicer", + [ + lambda y: y[:, 2:6], + lambda y: y[:, 0], + lambda y: y[..., 1:8:2], + lambda y: y[:, -1], + ], + ids=["range", "int_squeeze", "ellipsis_step", "neg_int"], +) +def test_tensor_slicing_onnx(cfg, slicer, tmp_path): + """KerasTensor slicing (GetItem ops) must export as ONNX Slice (+ Squeeze).""" + IN, OUT = 16, 8 + inputs = keras.Input(shape=(IN,)) + y = PQDense(cfg, units=OUT)(inputs) + model = keras.Model(inputs, slicer(y)) -def test_depthwise_conv2d_onnx(cfg, tmp_path): - IN_C, H, W = 4, 8, 8 - if _channels_first(): - input_shape = (IN_C, H, W) - x_np = np.random.randn(2, IN_C, H, W).astype(np.float32) - else: - input_shape = (H, W, IN_C) - x_np = np.random.randn(2, H, W, IN_C).astype(np.float32) + dummy = np.zeros((1, IN), dtype=np.float32) + model(dummy) + apply_final_compression(model) - inputs = keras.Input(shape=input_shape) - x = PQDepthwiseConv2d(cfg, kernel_size=3, padding="same")(inputs) - model = keras.Model(inputs, x) + x_np = np.random.randn(4, IN).astype(np.float32) + keras_output = keras_out(model, x_np) + onnx_out = onnx_run(model, x_np, input_shape=(IN,), tmp_path=tmp_path) + np.testing.assert_allclose(keras_output, onnx_out, atol=atol(cfg), err_msg="tensor slicing: keras vs ONNX mismatch") + + +@pytest.mark.parametrize( + "reshaper", + [ + lambda y: keras.ops.expand_dims(y, 1), + lambda y: keras.ops.expand_dims(y, -1), + lambda y: keras.ops.squeeze(keras.ops.expand_dims(y, 2), 2), + lambda y: keras.ops.squeeze(keras.ops.expand_dims(y, 1)), + ], + ids=["expand_dims", "expand_dims_neg", "roundtrip", "squeeze_all"], +) +def test_squeeze_unsqueeze_onnx(cfg, reshaper, tmp_path): + """keras.ops.squeeze / expand_dims must export as ONNX Squeeze/Unsqueeze.""" + IN, OUT = 16, 8 + inputs = keras.Input(shape=(IN,)) + y = PQDense(cfg, units=OUT)(inputs) + model = keras.Model(inputs, reshaper(y)) - dummy = np.zeros((1, *input_shape), dtype=np.float32) + dummy = np.zeros((1, IN), dtype=np.float32) model(dummy) apply_final_compression(model) - keras_out = _keras_out(model, x_np) - onnx_out = _onnx_run(model, x_np, input_shape=input_shape, tmp_path=tmp_path) - np.testing.assert_allclose(keras_out, onnx_out, atol=ATOL, err_msg="PQDepthwiseConv2d: keras vs ONNX mismatch") + x_np = np.random.randn(4, IN).astype(np.float32) + keras_output = keras_out(model, x_np) + onnx_out = onnx_run(model, x_np, input_shape=(IN,), tmp_path=tmp_path) + assert keras_output.shape == onnx_out.shape + np.testing.assert_allclose(keras_output, onnx_out, atol=atol(cfg), err_msg="squeeze/expand_dims: keras vs ONNX mismatch") + + +def rank3_dense_model(cfg): + """PQDense applied to a [batch, seq, dim] input (Dense maps over the last axis).""" + SEQ, DIM, OUT = 5, 16, 8 + inputs = keras.Input(shape=(SEQ, DIM)) + out = PQDense(cfg, units=OUT)(inputs) + model = keras.Model(inputs, out) + model(np.zeros((1, SEQ, DIM), dtype=np.float32)) + apply_final_compression(model) + return model, (SEQ, DIM) + + +def test_dense_rank3_input_onnx(cfg, tmp_path): + """A standalone PQDense on a rank-3 input must export as MatMul + Add (Gemm is rank-2 only).""" + model, input_shape = rank3_dense_model(cfg) + + x_np = np.random.randn(4, *input_shape).astype(np.float32) + keras_output = keras_out(model, x_np) + onnx_out = onnx_run(model, x_np, input_shape=input_shape, tmp_path=tmp_path) + np.testing.assert_allclose(keras_output, onnx_out, atol=atol(cfg), err_msg="rank-3 dense: keras vs ONNX mismatch") + + +def test_dense_rank3_integer_weights_onnx(tmp_path): + """store_integer_weights on a rank-3 dense exercises the quantized-weight Transpose branch.""" + cfg = pquant.cs_config() + cfg.quantization_parameters.enable_quantization = True + model, input_shape = rank3_dense_model(cfg) + + x_np = np.random.randn(4, *input_shape).astype(np.float32) + keras_output = keras_out(model, x_np) + + path = str(tmp_path / "rank3_int.onnx") + convert_to_onnx(model, input_shape=input_shape, output_path=path, store_integer_weights=True) + sess = ort.InferenceSession(path) + onnx_out = sess.run(None, {sess.get_inputs()[0].name: x_np})[0] + np.testing.assert_allclose(keras_output, onnx_out, atol=QUANT_ATOL, err_msg="rank-3 dense int weights: mismatch") diff --git a/tests/test_torch_onnx_converter.py b/tests/test_torch_onnx_converter.py index 2064d30..4927915 100644 --- a/tests/test_torch_onnx_converter.py +++ b/tests/test_torch_onnx_converter.py @@ -1,4 +1,4 @@ -"""Tests for convert_to_onnx / convert_to_onnx_fx. +"""Tests for convert_to_onnx. Each test builds a small model (one PQ layer + ReLU where applicable), runs a forward pass to initialise any running statistics, calls apply_final_compression @@ -11,6 +11,7 @@ import os import numpy as np +import onnxruntime as ort import pytest import torch import torch.nn as nn @@ -18,12 +19,10 @@ os.environ["KERAS_BACKEND"] = "torch" import pquant # noqa: E402 -from pquant.core.torch.convert_to_onnx import ( # noqa: E402 - convert_to_onnx, - convert_to_onnx_fx, - export_qdq_layernorm, -) +from pquant.core.torch.layers import Quantizer # noqa: E402 +from pquant.core.torch.onnx import convert_to_onnx # noqa: E402 from pquant.layers import ( # noqa: E402 + PQActivation, PQAvgPool1d, PQAvgPool2d, PQBatchNorm1d, @@ -31,38 +30,39 @@ PQConv1d, PQConv2d, PQDense, + PQLayerNorm, PQMultiheadAttention, ) -ort = pytest.importorskip("onnxruntime", reason="onnxruntime not installed") - -ATOL = 1e-4 # float32 Gemm/Conv can differ by ~1 ULP; keep some slack +ATOL = 1e-4 +QUANT_ATOL = 5e-3 -# --------------------------------------------------------------------------- -# fixtures -# --------------------------------------------------------------------------- +def atol(cfg): + return QUANT_ATOL if cfg.quantization_parameters.enable_quantization else ATOL -@pytest.fixture -def cfg(): +@pytest.fixture(params=[False, True], ids=["float", "quant"]) +def cfg(request): c = pquant.cs_config() - c.quantization_parameters.enable_quantization = False + c.quantization_parameters.enable_quantization = request.param return c -# --------------------------------------------------------------------------- -# helpers -# --------------------------------------------------------------------------- +@pytest.fixture +def cfg_quant(): + c = pquant.cs_config() + c.quantization_parameters.enable_quantization = True + return c -def _apply_compression(model: nn.Module): +def apply_compression(model: nn.Module): for m in model.modules(): if hasattr(m, "apply_final_compression"): m.apply_final_compression() -def _onnx_run(model: nn.Module, x: torch.Tensor, input_shape: tuple, tmp_path) -> np.ndarray: +def onnx_run(model: nn.Module, x: torch.Tensor, input_shape: tuple, tmp_path) -> np.ndarray: """Export model → ONNX file in tmp_path, run with onnxruntime, return output.""" path = str(tmp_path / "model.onnx") convert_to_onnx(model, input_shape=input_shape, output_path=path) @@ -71,314 +71,689 @@ def _onnx_run(model: nn.Module, x: torch.Tensor, input_shape: tuple, tmp_path) - return sess.run(None, {in_name: x.cpu().numpy()})[0] -def _onnx_run_fx(model: nn.Module, x: torch.Tensor, input_shape: tuple, tmp_path) -> np.ndarray: +def onnx_run_fx(model: nn.Module, x: torch.Tensor, input_shape: tuple, tmp_path) -> np.ndarray: """FX-based export → ONNX, run with onnxruntime.""" path = str(tmp_path / "model_fx.onnx") - convert_to_onnx_fx(model, input_shape=input_shape, output_path=path) + convert_to_onnx(model, input_shape=input_shape, output_path=path) sess = ort.InferenceSession(path) in_name = sess.get_inputs()[0].name return sess.run(None, {in_name: x.cpu().numpy()})[0] -def _torch_out(model: nn.Module, x: torch.Tensor) -> np.ndarray: +def torch_out(model: nn.Module, x: torch.Tensor) -> np.ndarray: model.eval() with torch.no_grad(): return model(x).cpu().numpy() -# --------------------------------------------------------------------------- -# PQDense -# --------------------------------------------------------------------------- +# (model factory, input shape without batch dim, batch size) +SINGLE_LAYER_CASES = [ + pytest.param( + lambda cfg: nn.Sequential(PQDense(cfg, in_features=16, out_features=8, bias=True), nn.ReLU()), + (16,), + 4, + id="dense-bias", + ), + pytest.param( + lambda cfg: nn.Sequential(PQDense(cfg, in_features=16, out_features=8, bias=False), nn.ReLU()), + (16,), + 4, + id="dense-nobias", + ), + pytest.param( + lambda cfg: nn.Sequential( + PQConv2d(cfg, in_channels=3, out_channels=8, kernel_size=3, padding=1, bias=True), nn.ReLU() + ), + (3, 8, 8), + 2, + id="conv2d-bias", + ), + pytest.param( + lambda cfg: nn.Sequential( + PQConv2d(cfg, in_channels=3, out_channels=8, kernel_size=3, padding=1, bias=False), nn.ReLU() + ), + (3, 8, 8), + 2, + id="conv2d-nobias", + ), + pytest.param( + lambda cfg: nn.Sequential( + PQConv1d(cfg, in_channels=4, out_channels=8, kernel_size=3, padding=1, bias=True), nn.ReLU() + ), + (4, 16), + 2, + id="conv1d-bias", + ), + pytest.param( + lambda cfg: nn.Sequential( + PQConv1d(cfg, in_channels=4, out_channels=8, kernel_size=3, padding=1, bias=False), nn.ReLU() + ), + (4, 16), + 2, + id="conv1d-nobias", + ), + pytest.param(lambda cfg: nn.Sequential(PQBatchNorm2d(cfg, num_features=8), nn.ReLU()), (8, 4, 4), 4, id="batchnorm2d"), + pytest.param(lambda cfg: nn.Sequential(PQBatchNorm1d(cfg, num_features=8), nn.ReLU()), (8, 16), 4, id="batchnorm1d"), + pytest.param(lambda cfg: nn.Sequential(PQAvgPool2d(cfg, kernel_size=2, stride=2)), (8, 8, 8), 2, id="avgpool2d"), + pytest.param(lambda cfg: nn.Sequential(PQAvgPool1d(cfg, kernel_size=2, stride=2)), (8, 16), 2, id="avgpool1d"), +] + + +@pytest.mark.parametrize("make_model,input_shape,batch", SINGLE_LAYER_CASES) +def test_single_layer_onnx(cfg, make_model, input_shape, batch, tmp_path): + model = make_model(cfg) + x = torch.randn(batch, *input_shape) + with torch.no_grad(): + model(x) # warm-up in train mode (initialises any running stats) + apply_compression(model) + torch_output = torch_out(model, x) # eval mode: BN uses running stats + onnx_out = onnx_run(model, x, input_shape=input_shape, tmp_path=tmp_path) + np.testing.assert_allclose(torch_output, onnx_out, atol=ATOL, err_msg="torch vs ONNX mismatch") -@pytest.mark.parametrize("bias", [True, False]) -def test_dense_onnx(cfg, bias, tmp_path): - IN, OUT = 16, 8 - model = nn.Sequential( - PQDense(cfg, in_features=IN, out_features=OUT, bias=bias), - nn.ReLU(), - ) - x = torch.randn(4, IN) - with torch.no_grad(): - model(x) # warm-up (needed for any running stats) - _apply_compression(model) - torch_out = _torch_out(model, x) - onnx_out = _onnx_run(model, x, input_shape=(IN,), tmp_path=tmp_path) - np.testing.assert_allclose(torch_out, onnx_out, atol=ATOL, err_msg=f"PQDense bias={bias}: torch vs ONNX mismatch") +class SelfAttnModel(nn.Module): + """Thin wrapper so FX tracing sees a single-input model.""" + def __init__(self, mha: PQMultiheadAttention): + super().__init__() + self.mha = mha -# --------------------------------------------------------------------------- -# PQConv2d -# --------------------------------------------------------------------------- + def forward(self, x): + out, _ = self.mha(x, x, x) + return out @pytest.mark.parametrize("bias", [True, False]) -def test_conv2d_onnx(cfg, bias, tmp_path): - IN_C, OUT_C, H, W = 3, 8, 8, 8 - model = nn.Sequential( - PQConv2d(cfg, in_channels=IN_C, out_channels=OUT_C, kernel_size=3, padding=1, bias=bias), - nn.ReLU(), - ) - x = torch.randn(2, IN_C, H, W) +def test_mha_onnx(cfg, bias, tmp_path): + E, H, T = 16, 4, 8 + mha = PQMultiheadAttention(cfg, embed_dim=E, num_heads=H, bias=bias, batch_first=True) + model = SelfAttnModel(mha) + + x = torch.randn(2, T, E) with torch.no_grad(): model(x) - _apply_compression(model) + apply_compression(model) + + # The quantized softmax (exp/inv LUTs) re-quantizes intermediates, so allow ~1 LSB + # of rounding-boundary slack when quantization is enabled (see atol). + torch_output = torch_out(model, x) + onnx_out = onnx_run_fx(model, x, input_shape=(T, E), tmp_path=tmp_path) + np.testing.assert_allclose( + torch_output, onnx_out, atol=atol(cfg), err_msg=f"PQMultiheadAttention bias={bias}: torch vs ONNX mismatch" + ) - torch_out = _torch_out(model, x) - onnx_out = _onnx_run(model, x, input_shape=(IN_C, H, W), tmp_path=tmp_path) - np.testing.assert_allclose(torch_out, onnx_out, atol=ATOL, err_msg=f"PQConv2d bias={bias}: torch vs ONNX mismatch") +class CausalSelfAttnModel(nn.Module): + """Self-attention with a constant additive causal mask (the decoder-inference case).""" -# --------------------------------------------------------------------------- -# PQConv1d -# --------------------------------------------------------------------------- + def __init__(self, mha: PQMultiheadAttention, seq_len: int): + super().__init__() + self.mha = mha + # (T, S) additive mask: 0 on/below the diagonal, large-negative above it. + self.register_buffer("attn_mask", torch.triu(torch.full((seq_len, seq_len), -1e4), diagonal=1)) + + def forward(self, x): + out, _ = self.mha(x, x, x, attn_mask=self.attn_mask) + return out @pytest.mark.parametrize("bias", [True, False]) -def test_conv1d_onnx(cfg, bias, tmp_path): - IN_C, OUT_C, L = 4, 8, 16 - model = nn.Sequential( - PQConv1d(cfg, in_channels=IN_C, out_channels=OUT_C, kernel_size=3, padding=1, bias=bias), - nn.ReLU(), - ) - x = torch.randn(2, IN_C, L) +def test_mha_causal_attn_mask_onnx(cfg, bias, tmp_path): + E, H, T = 16, 4, 8 + mha = PQMultiheadAttention(cfg, embed_dim=E, num_heads=H, bias=bias, batch_first=True) + model = CausalSelfAttnModel(mha, T) + + x = torch.randn(2, T, E) with torch.no_grad(): model(x) - _apply_compression(model) + apply_compression(model) - torch_out = _torch_out(model, x) - onnx_out = _onnx_run(model, x, input_shape=(IN_C, L), tmp_path=tmp_path) - np.testing.assert_allclose(torch_out, onnx_out, atol=ATOL, err_msg=f"PQConv1d bias={bias}: torch vs ONNX mismatch") + torch_output = torch_out(model, x) + onnx_out = onnx_run_fx(model, x, input_shape=(T, E), tmp_path=tmp_path) + np.testing.assert_allclose( + torch_output, onnx_out, atol=atol(cfg), err_msg=f"MHA causal attn_mask bias={bias}: torch vs ONNX mismatch" + ) -# --------------------------------------------------------------------------- -# PQBatchNorm2d -# --------------------------------------------------------------------------- +class PaddedSelfAttnModel(nn.Module): + """Self-attention with a runtime bool key_padding_mask input (True == padding).""" + def __init__(self, mha: PQMultiheadAttention): + super().__init__() + self.mha = mha -def test_batchnorm2d_onnx(cfg, tmp_path): - C, H, W = 8, 4, 4 - model = nn.Sequential( - PQBatchNorm2d(cfg, num_features=C), - nn.ReLU(), - ) - x = torch.randn(4, C, H, W) + def forward(self, x, key_padding_mask): + out, _ = self.mha(x, x, x, key_padding_mask=key_padding_mask) + return out + + +@pytest.mark.parametrize("bias", [True, False]) +def test_mha_key_padding_mask_onnx(cfg, bias, tmp_path): + import onnx + + E, H, T = 16, 4, 8 + mha = PQMultiheadAttention(cfg, embed_dim=E, num_heads=H, bias=bias, batch_first=True) + model = PaddedSelfAttnModel(mha) + + x = torch.randn(2, T, E) + key_padding_mask = torch.zeros(2, T, dtype=torch.bool) + key_padding_mask[:, -2:] = True # last two key positions are padding with torch.no_grad(): - model(x) - _apply_compression(model) - model.eval() # switch BN to use running stats + model(x, key_padding_mask) + apply_compression(model) - torch_out = _torch_out(model, x) - onnx_out = _onnx_run(model, x, input_shape=(C, H, W), tmp_path=tmp_path) - np.testing.assert_allclose(torch_out, onnx_out, atol=ATOL, err_msg="PQBatchNorm2d: torch vs ONNX mismatch") + model.eval() + with torch.no_grad(): + torch_output = model(x, key_padding_mask).cpu().numpy() + path = str(tmp_path / "mha_kpm.onnx") + proto = convert_to_onnx(model, input_shape=[(T, E), (T,)], output_path=path, input_dtypes=["float32", "bool"]) -# --------------------------------------------------------------------------- -# PQBatchNorm1d -# --------------------------------------------------------------------------- + # The padding mask must be a genuine bool graph input, not baked away. + kpm_vi = next(i for i in proto.graph.input if i.name == "key_padding_mask") + assert kpm_vi.type.tensor_type.elem_type == onnx.TensorProto.BOOL + sess = ort.InferenceSession(path) + onnx_out = sess.run(None, {"x": x.cpu().numpy(), "key_padding_mask": key_padding_mask.cpu().numpy()})[0] -def test_batchnorm1d_onnx(cfg, tmp_path): - C, L = 8, 16 - model = nn.Sequential( - PQBatchNorm1d(cfg, num_features=C), - nn.ReLU(), + np.testing.assert_allclose( + torch_output, onnx_out, atol=atol(cfg), err_msg=f"MHA key_padding_mask bias={bias}: torch vs ONNX mismatch" ) - x = torch.randn(4, C, L) + + +class TwoInputModel(nn.Module): + """Two tensor inputs merged by addition after independent Dense layers.""" + + def __init__(self, cfg, in_a: int, in_b: int, out: int, bias: bool): + super().__init__() + self.dense_a = PQDense(cfg, in_features=in_a, out_features=out, bias=bias) + self.dense_b = PQDense(cfg, in_features=in_b, out_features=out, bias=bias) + + def forward(self, a, b): + return torch.relu(self.dense_a(a) + self.dense_b(b)) + + +@pytest.mark.parametrize("bias", [True, False]) +def test_two_input_onnx(cfg, bias, tmp_path): + IN_A, IN_B, OUT = 16, 4, 8 + model = TwoInputModel(cfg, IN_A, IN_B, OUT, bias) + + a = torch.randn(3, IN_A) + b = torch.randn(3, IN_B) with torch.no_grad(): - model(x) - _apply_compression(model) + model(a, b) # warm-up + apply_compression(model) + model.eval() + with torch.no_grad(): + torch_output = model(a, b).cpu().numpy() + + path = str(tmp_path / "two_input.onnx") + model_proto = convert_to_onnx(model, input_shape=[(IN_A,), (IN_B,)], output_path=path) + + # Graph must declare exactly two inputs, named after the forward parameters, + # each with a dynamic batch dim and its own feature shape. + assert [i.name for i in model_proto.graph.input] == ["a", "b"] + in_shapes = {i.name: [d.dim_value for d in i.type.tensor_type.shape.dim] for i in model_proto.graph.input} + assert in_shapes["a"] == [0, IN_A] # dim_value 0 == dynamic (no batch fixed) + assert in_shapes["b"] == [0, IN_B] + + sess = ort.InferenceSession(path) + names = [i.name for i in sess.get_inputs()] + assert set(names) == {"a", "b"} + onnx_out = sess.run(None, {"a": a.cpu().numpy(), "b": b.cpu().numpy()})[0] - torch_out = _torch_out(model, x) - onnx_out = _onnx_run(model, x, input_shape=(C, L), tmp_path=tmp_path) - np.testing.assert_allclose(torch_out, onnx_out, atol=ATOL, err_msg="PQBatchNorm1d: torch vs ONNX mismatch") + np.testing.assert_allclose(torch_output, onnx_out, atol=ATOL, err_msg=f"two-input bias={bias}: torch vs ONNX mismatch") -# --------------------------------------------------------------------------- -# PQAvgPool2d -# --------------------------------------------------------------------------- +def test_two_input_shape_count_mismatch(cfg, tmp_path): + """Wrong number of shapes for the model's tensor inputs is a clear error.""" + model = TwoInputModel(cfg, 16, 4, 8, bias=True) + with torch.no_grad(): + model(torch.randn(2, 16), torch.randn(2, 4)) + apply_compression(model) + + path = str(tmp_path / "bad_count.onnx") + with pytest.raises(ValueError, match="tensor input"): + convert_to_onnx(model, input_shape=(16,), output_path=path) # only one shape + + +class FlaggedModel(nn.Module): + """One tensor input plus a bool flag selecting an optional scaling branch.""" + + def __init__(self, cfg, in_features: int, out: int): + super().__init__() + self.dense = PQDense(cfg, in_features=in_features, out_features=out, bias=True) + def forward(self, x, scale_up: bool = False): + out = self.dense(x) + if scale_up: + out = out * 2.0 + return torch.relu(out) -def test_avgpool2d_onnx(cfg, tmp_path): - C, H, W = 8, 8, 8 - model = nn.Sequential( - PQAvgPool2d(cfg, kernel_size=2, stride=2), + +@pytest.mark.parametrize("scale_up", [False, True]) +def test_concrete_args_specialization(cfg, scale_up, tmp_path): + IN, OUT = 16, 8 + model = FlaggedModel(cfg, IN, OUT) + + x = torch.randn(3, IN) + with torch.no_grad(): + model(x, scale_up) + apply_compression(model) + + model.eval() + with torch.no_grad(): + torch_output = model(x, scale_up).cpu().numpy() + + path = str(tmp_path / f"flag_{scale_up}.onnx") + model_proto = convert_to_onnx(model, input_shape=(IN,), output_path=path, concrete_args={"scale_up": scale_up}) + + # The bool flag is baked in as a constant, so it must NOT appear as a graph + # input — only the single tensor input "input" remains. + assert [i.name for i in model_proto.graph.input] == ["input"] + + sess = ort.InferenceSession(path) + assert [i.name for i in sess.get_inputs()] == ["input"] + onnx_out = sess.run(None, {"input": x.cpu().numpy()})[0] + + np.testing.assert_allclose( + torch_output, onnx_out, atol=ATOL, err_msg=f"concrete_args scale_up={scale_up}: torch vs ONNX mismatch" ) - x = torch.randn(2, C, H, W) + + +class ResidualConcatModel(nn.Module): + """Exercises the FX converter's branch handling: a skip-add and a concat.""" + + def __init__(self, cfg, dim: int, out: int): + super().__init__() + self.d1 = PQDense(cfg, in_features=dim, out_features=dim) + self.d2 = PQDense(cfg, in_features=dim, out_features=dim) + self.d3 = PQDense(cfg, in_features=2 * dim, out_features=out) + + def forward(self, x): + h = self.d1(x) + h = h + self.d2(h) # residual / skip add + h = torch.cat([h, x], dim=1) # branch merge by concatenation + return self.d3(h) + + +def test_residual_concat_onnx(cfg_quant, tmp_path): + DIM, OUT = 16, 8 + model = ResidualConcatModel(cfg_quant, DIM, OUT) + + x = torch.randn(4, DIM) with torch.no_grad(): model(x) - _apply_compression(model) + apply_compression(model) - torch_out = _torch_out(model, x) - onnx_out = _onnx_run(model, x, input_shape=(C, H, W), tmp_path=tmp_path) - np.testing.assert_allclose(torch_out, onnx_out, atol=ATOL, err_msg="PQAvgPool2d: torch vs ONNX mismatch") + model.eval() + with torch.no_grad(): + torch_output = model(x).cpu().numpy() + path = str(tmp_path / "residual_concat.onnx") + model_proto = convert_to_onnx(model, input_shape=(DIM,), output_path=path) + op_types = [n.op_type for n in model_proto.graph.node] + assert "Add" in op_types # the skip connection + assert "Concat" in op_types # the branch merge -# --------------------------------------------------------------------------- -# PQAvgPool1d -# --------------------------------------------------------------------------- + sess = ort.InferenceSession(path) + onnx_out = sess.run(None, {sess.get_inputs()[0].name: x.cpu().numpy()})[0] + np.testing.assert_allclose(torch_output, onnx_out, atol=QUANT_ATOL, err_msg="residual+concat: torch vs ONNX mismatch") -def test_avgpool1d_onnx(cfg, tmp_path): - C, L = 8, 16 - model = nn.Sequential( - PQAvgPool1d(cfg, kernel_size=2, stride=2), - ) - x = torch.randn(2, C, L) +@pytest.mark.parametrize("activation", ["relu", "tanh", "hard_tanh", "leaky_relu", "gelu"]) +def test_pqactivation_onnx(cfg_quant, activation, tmp_path): + DIM = 16 + model = nn.Sequential(PQActivation(cfg_quant, activation=activation)) + + x = torch.randn(4, DIM) with torch.no_grad(): model(x) - _apply_compression(model) + apply_compression(model) - torch_out = _torch_out(model, x) - onnx_out = _onnx_run(model, x, input_shape=(C, L), tmp_path=tmp_path) - np.testing.assert_allclose(torch_out, onnx_out, atol=ATOL, err_msg="PQAvgPool1d: torch vs ONNX mismatch") + torch_output = torch_out(model, x) + onnx_out = onnx_run(model, x, input_shape=(DIM,), tmp_path=tmp_path) + np.testing.assert_allclose( + torch_output, onnx_out, atol=ATOL, err_msg=f"PQActivation {activation}: torch vs ONNX mismatch" + ) + + +def test_standalone_quantizer_onnx(cfg_quant, tmp_path): + qp = cfg_quant.quantization_parameters + quant = Quantizer( + k=qp.default_data_keep_negatives, + i=qp.default_data_integer_bits, + f=qp.default_data_fractional_bits, + overflow=qp.overflow_mode_data, + round_mode=qp.round_mode, + is_heterogeneous=False, + is_data=True, + granularity="per_tensor", + hgq_gamma=qp.hgq_gamma, + ) + model = nn.Sequential(quant) + x = torch.randn(4, 16) + with torch.no_grad(): + model(x) + apply_compression(model) -# --------------------------------------------------------------------------- -# PQMultiheadAttention (uses FX converter; self-attention, batch_first=True) -# --------------------------------------------------------------------------- + # A standalone quantizer must emit a Quantize/Dequantize pair. + path = str(tmp_path / "quantizer.onnx") + model_proto = convert_to_onnx(model, input_shape=(16,), output_path=path) + op_types = [n.op_type for n in model_proto.graph.node] + assert "QuantizeLinear" in op_types and "DequantizeLinear" in op_types + torch_output = torch_out(model, x) + sess = ort.InferenceSession(path) + onnx_out = sess.run(None, {sess.get_inputs()[0].name: x.cpu().numpy()})[0] + np.testing.assert_allclose(torch_output, onnx_out, atol=ATOL, err_msg="standalone Quantizer: torch vs ONNX mismatch") -class _SelfAttnModel(nn.Module): - """Thin wrapper so FX tracing sees a single-input model.""" - def __init__(self, mha: PQMultiheadAttention): +class CNNFlattenModel(nn.Module): + def __init__(self, cfg, in_c: int, hw: int, out: int, use_reshape: bool): super().__init__() - self.mha = mha + self.conv = PQConv2d(cfg, in_channels=in_c, out_channels=4, kernel_size=3, padding=1) + self.dense = PQDense(cfg, in_features=4 * hw * hw, out_features=out) + self.use_reshape = use_reshape + self._flat = 4 * hw * hw def forward(self, x): - out, _ = self.mha(x, x, x) - return out + h = torch.relu(self.conv(x)) + # Both are common CNN→Dense transitions; reshape uses a constant (-1, N) + # shape because the FX exporter requires static reshape targets. + h = h.reshape(-1, self._flat) if self.use_reshape else torch.flatten(h, 1) + return self.dense(h) -@pytest.mark.parametrize("bias", [True, False]) -def test_mha_onnx(cfg, bias, tmp_path): - E, H, T = 16, 4, 8 - mha = PQMultiheadAttention(cfg, embed_dim=E, num_heads=H, bias=bias, batch_first=True) - model = _SelfAttnModel(mha) +@pytest.mark.parametrize("use_reshape", [False, True]) +def test_cnn_flatten_to_dense_onnx(cfg_quant, use_reshape, tmp_path): + IN_C, HW, OUT = 3, 8, 8 + model = CNNFlattenModel(cfg_quant, IN_C, HW, OUT, use_reshape) - x = torch.randn(2, T, E) + x = torch.randn(2, IN_C, HW, HW) with torch.no_grad(): model(x) - _apply_compression(model) + apply_compression(model) - torch_out = _torch_out(model, x) - onnx_out = _onnx_run_fx(model, x, input_shape=(T, E), tmp_path=tmp_path) + model.eval() + with torch.no_grad(): + torch_output = model(x).cpu().numpy() + + path = str(tmp_path / f"cnn_flatten_{use_reshape}.onnx") + convert_to_onnx(model, input_shape=(IN_C, HW, HW), output_path=path) + sess = ort.InferenceSession(path) + onnx_out = sess.run(None, {sess.get_inputs()[0].name: x.cpu().numpy()})[0] np.testing.assert_allclose( - torch_out, onnx_out, atol=ATOL, err_msg=f"PQMultiheadAttention bias={bias}: torch vs ONNX mismatch" + torch_output, onnx_out, atol=QUANT_ATOL, err_msg=f"CNN→Dense reshape={use_reshape}: torch vs ONNX mismatch" ) -# --------------------------------------------------------------------------- -# Static-QDQ LayerNormalization graph -# --------------------------------------------------------------------------- +class ScalarOpsModel(nn.Module): + def __init__(self, cfg, dim: int): + super().__init__() + self.d = PQDense(cfg, in_features=dim, out_features=dim) + def forward(self, x): + h = self.d(x) + h = h * 2.0 # Mul with scalar literal + h = h - 1.0 # Sub with scalar literal + h = h / 3.0 # Div with scalar literal + return torch.sigmoid(h) -@pytest.mark.parametrize("input_shape", [(4, 64), (1, 4, 64)]) -def test_qdq_layernorm_export(input_shape, tmp_path): - import onnx - D = input_shape[-1] - rng = np.random.default_rng(0) - # Q7 representable: gamma = k / 128, k integer, |k| < 32768 - gamma_q = rng.integers(low=64, high=192, size=(D,), dtype=np.int32) # ~0.5 .. 1.5 - gamma = (gamma_q.astype(np.float32)) / (1 << 7) - # Q15 representable: beta = k / 32768, k integer, |k| < 32768 (so |beta| < 1) - beta_q = rng.integers(low=-1024, high=1024, size=(D,), dtype=np.int32) - beta = (beta_q.astype(np.float32)) / (1 << 15) - - input_scale_log2 = -7 # input_scale = 2**-7 - output_scale_log2 = -6 # output_scale = 2**-6 - eps_q0 = 1 - - path = str(tmp_path / "qdq_layernorm.onnx") - model_proto = export_qdq_layernorm( - output_path=path, - input_shape=input_shape, - gamma=gamma, - beta=beta, - input_scale_log2=input_scale_log2, - output_scale_log2=output_scale_log2, - eps_q0=eps_q0, - ) +def test_scalar_ops_onnx(cfg_quant, tmp_path): + DIM = 16 + model = ScalarOpsModel(cfg_quant, DIM) - # ----- structural checks ----- + x = torch.randn(4, DIM) + with torch.no_grad(): + model(x) + apply_compression(model) + + model.eval() + with torch.no_grad(): + torch_output = model(x).cpu().numpy() + + path = str(tmp_path / "scalar_ops.onnx") + model_proto = convert_to_onnx(model, input_shape=(DIM,), output_path=path) op_types = [n.op_type for n in model_proto.graph.node] - assert op_types == ["DequantizeLinear", "LayerNormalization", "QuantizeLinear", "DequantizeLinear"] - - ln_node = model_proto.graph.node[1] - axis = next(a.i for a in ln_node.attribute if a.name == "axis") - eps_attr = next(a.f for a in ln_node.attribute if a.name == "epsilon") - assert axis == -1 - expected_eps = eps_q0 * (2.0**input_scale_log2) ** 2 - assert abs(eps_attr - expected_eps) < 1e-12 - - # input must be int8, output float - assert len(model_proto.graph.input) == 1 - assert model_proto.graph.input[0].type.tensor_type.elem_type == onnx.TensorProto.INT8 - assert model_proto.graph.output[0].type.tensor_type.elem_type == onnx.TensorProto.FLOAT - in_dims = [d.dim_value for d in model_proto.graph.input[0].type.tensor_type.shape.dim] - assert tuple(in_dims) == input_shape - - # zero-points must be int8 zero - inits = {t.name: t for t in model_proto.graph.initializer} - for zp_name in ("input_zero_point", "output_zero_point"): - zp = onnx.numpy_helper.to_array(inits[zp_name]) - assert zp.dtype == np.int8 - assert int(zp) == 0 - - # scales must be exact powers of two - in_scale = float(onnx.numpy_helper.to_array(inits["input_scale"])) - out_scale = float(onnx.numpy_helper.to_array(inits["output_scale"])) - assert in_scale == 2.0**input_scale_log2 - assert out_scale == 2.0**output_scale_log2 - - # ----- numerical check via onnxruntime ----- + for expected in ("Mul", "Sub", "Div", "Sigmoid"): + assert expected in op_types, f"missing {expected} node" + sess = ort.InferenceSession(path) - in_name = sess.get_inputs()[0].name - x_q = rng.integers(low=-64, high=64, size=input_shape, dtype=np.int8) - onnx_out = sess.run(None, {in_name: x_q})[0] - - # Reference: dequantize -> layernorm(axis=-1) -> quantize -> dequantize - x_f = x_q.astype(np.float32) * in_scale - mean = x_f.mean(axis=-1, keepdims=True) - var = x_f.var(axis=-1, keepdims=True) - x_norm = (x_f - mean) / np.sqrt(var + expected_eps) - y_f = x_norm * gamma + beta - y_q = np.clip(np.round(y_f / out_scale), -128, 127).astype(np.int8) - y_ref = y_q.astype(np.float32) * out_scale - - np.testing.assert_allclose(onnx_out, y_ref, atol=out_scale * 0.5) - - -def test_qdq_layernorm_validation(tmp_path): - path = str(tmp_path / "bad.onnx") - D = 64 - gamma = np.ones(D, dtype=np.float32) - beta = np.zeros(D, dtype=np.float32) - - # rank-1 input: rejected - with pytest.raises(ValueError, match="rank"): - export_qdq_layernorm(path, (D,), gamma, beta, -7, -6) - - # last dim not multiple of 32 - with pytest.raises(ValueError, match="multiple of 32"): - export_qdq_layernorm(path, (4, 16), np.ones(16, np.float32), np.zeros(16, np.float32), -7, -6) - - # last dim not power of two (96 = 32*3) - with pytest.raises(ValueError, match="power of two"): - export_qdq_layernorm(path, (4, 96), np.ones(96, np.float32), np.zeros(96, np.float32), -7, -6) - - # gamma not Q7-representable (1/3 is not k/128 exactly) - with pytest.raises(ValueError, match="gamma"): - export_qdq_layernorm(path, (4, D), np.full(D, 1.0 / 3.0, np.float32), beta, -7, -6) - - # beta not Q15-representable (1/3 is not k/32768 exactly) - with pytest.raises(ValueError, match="beta"): - export_qdq_layernorm(path, (4, D), gamma, np.full(D, 1.0 / 3.0, np.float32), -7, -6) - - # eps_q0 < 1 - with pytest.raises(ValueError, match="eps_q0"): - export_qdq_layernorm(path, (4, D), gamma, beta, -7, -6, eps_q0=0) + onnx_out = sess.run(None, {sess.get_inputs()[0].name: x.cpu().numpy()})[0] + np.testing.assert_allclose(torch_output, onnx_out, atol=ATOL, err_msg="scalar ops+sigmoid: torch vs ONNX mismatch") + + +class MultiOutputModel(nn.Module): + def __init__(self, cfg, dim: int): + super().__init__() + self.a = PQDense(cfg, in_features=dim, out_features=8) + self.b = PQDense(cfg, in_features=dim, out_features=4) + + def forward(self, x): + return self.a(x), self.b(x) + + +def test_multi_output_onnx(cfg_quant, tmp_path): + DIM = 16 + model = MultiOutputModel(cfg_quant, DIM) + + x = torch.randn(3, DIM) + with torch.no_grad(): + model(x) + apply_compression(model) + + model.eval() + with torch.no_grad(): + t0, t1 = (t.cpu().numpy() for t in model(x)) + + path = str(tmp_path / "multi_output.onnx") + model_proto = convert_to_onnx(model, input_shape=(DIM,), output_path=path) + assert len(model_proto.graph.output) == 2 + + sess = ort.InferenceSession(path) + out_names = [o.name for o in sess.get_outputs()] + outs = sess.run(None, {sess.get_inputs()[0].name: x.cpu().numpy()}) + by_name = dict(zip(out_names, outs)) + # Match outputs by shape (order is preserved, but be explicit about which is which). + o0 = next(v for v in by_name.values() if v.shape[1] == 8) + o1 = next(v for v in by_name.values() if v.shape[1] == 4) + np.testing.assert_allclose(t0, o0, atol=ATOL, err_msg="multi-output[0]: torch vs ONNX mismatch") + np.testing.assert_allclose(t1, o1, atol=ATOL, err_msg="multi-output[1]: torch vs ONNX mismatch") + + +def test_pqlayernorm_onnx(cfg_quant, tmp_path): + DIM = 16 + model = nn.Sequential(PQLayerNorm(cfg_quant, normalized_shape=DIM)) + + x = torch.randn(4, DIM) + with torch.no_grad(): + model(x) + apply_compression(model) + + model.eval() + with torch.no_grad(): + torch_output = model(x).cpu().numpy() + + path = str(tmp_path / "pqlayernorm.onnx") + # LayerNormalization is an opset-17 op; the converter default (13) cannot host it. + convert_to_onnx(model, input_shape=(DIM,), output_path=path, opset=17) + sess = ort.InferenceSession(path) + onnx_out = sess.run(None, {sess.get_inputs()[0].name: x.cpu().numpy()})[0] + np.testing.assert_allclose(torch_output, onnx_out, atol=ATOL, err_msg="PQLayerNorm: torch vs ONNX mismatch") + + +@pytest.mark.parametrize( + "make_model,input_shape,batch", + [ + (lambda: nn.Sequential(nn.LeakyReLU(0.1)), (8, 16), 4), + (lambda: nn.Sequential(nn.MaxPool2d(2, 2)), (3, 8, 8), 2), + (lambda: nn.Sequential(nn.Upsample(scale_factor=2, mode="nearest")), (3, 4, 4), 2), + (lambda: nn.Sequential(nn.Dropout(0.5)), (16,), 4), + ], + ids=["leaky_relu", "maxpool2d", "upsample", "dropout"], +) +def test_plain_passthrough_layers_onnx(make_model, input_shape, batch, tmp_path): + model = make_model() + model.eval() # Dropout/BatchNorm must be in eval mode for a deterministic compare + x = torch.randn(batch, *input_shape) + + torch_output = torch_out(model, x) + onnx_out = onnx_run(model, x, input_shape=input_shape, tmp_path=tmp_path) + np.testing.assert_allclose(torch_output, onnx_out, atol=ATOL, err_msg="plain passthrough layer: torch vs ONNX mismatch") + + +def quantized_dense_model(cfg_quant): + model = nn.Sequential(PQDense(cfg_quant, in_features=16, out_features=8), nn.ReLU()) + x = torch.randn(4, 16) + with torch.no_grad(): + model(x) + apply_compression(model) + return model, x + + +@pytest.mark.parametrize("integer_ops", [False, True]) +def test_integer_weight_storage_onnx(cfg_quant, integer_ops, tmp_path): + """store_integer_weights and integer_ops (MatMulInteger) must stay numerically exact.""" + model, x = quantized_dense_model(cfg_quant) + torch_output = torch_out(model, x) + + path = str(tmp_path / f"int_{integer_ops}.onnx") + kwargs = {"integer_ops": True} if integer_ops else {"store_integer_weights": True} + convert_to_onnx(model, input_shape=(16,), output_path=path, **kwargs) + sess = ort.InferenceSession(path) + onnx_out = sess.run(None, {sess.get_inputs()[0].name: x.cpu().numpy()})[0] + np.testing.assert_allclose(torch_output, onnx_out, atol=ATOL, err_msg=f"integer_ops={integer_ops}: mismatch") + + +def rank3_dense_model(cfg): + """PQDense applied to a [batch, seq, dim] input (F.linear maps over the last axis).""" + SEQ, DIM, OUT = 5, 16, 8 + model = nn.Sequential(PQDense(cfg, in_features=DIM, out_features=OUT), nn.ReLU()) + x = torch.randn(4, SEQ, DIM) + with torch.no_grad(): + model(x) + apply_compression(model) + return model, x, (SEQ, DIM) + + +def test_dense_rank3_input_onnx(cfg, tmp_path): + """A standalone PQDense on a rank-3 input must export as MatMul + Add (Gemm is rank-2 only).""" + model, x, input_shape = rank3_dense_model(cfg) + torch_output = torch_out(model, x) + onnx_out = onnx_run(model, x, input_shape=input_shape, tmp_path=tmp_path) + np.testing.assert_allclose(torch_output, onnx_out, atol=atol(cfg), err_msg="rank-3 dense: torch vs ONNX mismatch") + + +@pytest.mark.parametrize("integer_ops", [False, True]) +def test_dense_rank3_integer_onnx(cfg_quant, integer_ops, tmp_path): + """Integer weight storage and MatMulInteger must also handle rank-3 dense inputs.""" + model, x, input_shape = rank3_dense_model(cfg_quant) + torch_output = torch_out(model, x) + + path = str(tmp_path / f"rank3_int_{integer_ops}.onnx") + kwargs = {"integer_ops": True} if integer_ops else {"store_integer_weights": True} + convert_to_onnx(model, input_shape=input_shape, output_path=path, **kwargs) + sess = ort.InferenceSession(path) + onnx_out = sess.run(None, {sess.get_inputs()[0].name: x.cpu().numpy()})[0] + np.testing.assert_allclose(torch_output, onnx_out, atol=ATOL, err_msg=f"rank-3 integer_ops={integer_ops}: mismatch") + + +def test_include_clip_toggle_structure(cfg_quant, tmp_path): + """include_clip controls whether a Clip node precedes each input QuantizeLinear.""" + model, _ = quantized_dense_model(cfg_quant) + + proto_clip = convert_to_onnx(model, input_shape=(16,), output_path=str(tmp_path / "clip.onnx"), include_clip=True) + proto_noclip = convert_to_onnx(model, input_shape=(16,), output_path=str(tmp_path / "noclip.onnx"), include_clip=False) + assert "Clip" in [n.op_type for n in proto_clip.graph.node] + assert "Clip" not in [n.op_type for n in proto_noclip.graph.node] + + +def test_batch_size_fixes_input_dim(cfg_quant, tmp_path): + """batch_size pins the graph's batch dimension instead of leaving it dynamic.""" + model, _ = quantized_dense_model(cfg_quant) + + proto = convert_to_onnx(model, input_shape=(16,), output_path=str(tmp_path / "bs.onnx"), batch_size=4) + in_dims = [d.dim_value for d in proto.graph.input[0].type.tensor_type.shape.dim] + assert in_dims[0] == 4 # batch fixed + assert in_dims[1] == 16 + + +def test_qonnx_export_builds(cfg_quant, tmp_path): + """use_qonnx emits QONNX Quant nodes and produces a structurally valid model.""" + import onnx + + model, _ = quantized_dense_model(cfg_quant) + path = str(tmp_path / "qonnx.onnx") + proto = convert_to_onnx(model, input_shape=(16,), output_path=path, use_qonnx=True) + onnx.checker.check_model(onnx.load(path)) + assert any(n.op_type == "Quant" for n in proto.graph.node) + + +@pytest.mark.parametrize( + "slicer", + [ + lambda y: y[:, 2:6], + lambda y: y[:, 0], + lambda y: y[..., 1:8:2], + lambda y: y[:, -1], + ], + ids=["range", "int_squeeze", "ellipsis_step", "neg_int"], +) +def test_tensor_slicing_onnx(cfg, slicer, tmp_path): + """Constant tensor slicing in forward must export as ONNX Slice (+ Squeeze).""" + + class SliceModel(nn.Module): + def __init__(self): + super().__init__() + self.dense = PQDense(cfg, in_features=16, out_features=8) + + def forward(self, x): + return slicer(self.dense(x)) + + model = SliceModel() + x = torch.randn(4, 16) + with torch.no_grad(): + model(x) + apply_compression(model) + + torch_output = torch_out(model, x) + onnx_out = onnx_run(model, x, input_shape=(16,), tmp_path=tmp_path) + np.testing.assert_allclose(torch_output, onnx_out, atol=atol(cfg), err_msg="tensor slicing: torch vs ONNX mismatch") + + +@pytest.mark.parametrize( + "reshaper", + [ + lambda y: y.unsqueeze(1), + lambda y: torch.unsqueeze(y, -1), + lambda y: y.unsqueeze(2).squeeze(2), + lambda y: torch.squeeze(y.unsqueeze(1)), + ], + ids=["method_unsqueeze", "fn_unsqueeze_neg", "roundtrip", "squeeze_all"], +) +def test_squeeze_unsqueeze_onnx(cfg, reshaper, tmp_path): + """squeeze/unsqueeze in forward must export as ONNX Squeeze/Unsqueeze.""" + + class ReshapeModel(nn.Module): + def __init__(self): + super().__init__() + self.dense = PQDense(cfg, in_features=16, out_features=8) + + def forward(self, x): + return reshaper(self.dense(x)) + + model = ReshapeModel() + x = torch.randn(4, 16) + with torch.no_grad(): + model(x) + apply_compression(model) + + torch_output = torch_out(model, x) + onnx_out = onnx_run(model, x, input_shape=(16,), tmp_path=tmp_path) + assert torch_output.shape == onnx_out.shape + np.testing.assert_allclose(torch_output, onnx_out, atol=atol(cfg), err_msg="squeeze/unsqueeze: torch vs ONNX mismatch")