From df71727b6168600a9311d240482675ab50d96778 Mon Sep 17 00:00:00 2001 From: Roope Niemi Date: Thu, 9 Jul 2026 19:13:50 +0200 Subject: [PATCH 1/8] onnx converter improvements: multi-input models, more tests --- src/pquant/core/keras/convert_to_onnx.py | 692 ++++++++----------- src/pquant/core/torch/convert_to_onnx.py | 828 +++++++++-------------- tests/test_keras_onnx_converter.py | 192 ++++-- tests/test_torch_onnx_converter.py | 552 +++++++++++++-- 4 files changed, 1276 insertions(+), 988 deletions(-) diff --git a/src/pquant/core/keras/convert_to_onnx.py b/src/pquant/core/keras/convert_to_onnx.py index 35263c2..57aa39b 100644 --- a/src/pquant/core/keras/convert_to_onnx.py +++ b/src/pquant/core/keras/convert_to_onnx.py @@ -238,6 +238,18 @@ def _int_weight_node(name_prefix, weight_np, k, i, f, initializers): # noqa: AR # --------------------------------------------------------------------------- +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 _np(tensor): """Convert a Keras tensor / variable / scalar to a float32 numpy array.""" return np.array(tensor, dtype=np.float32) @@ -275,12 +287,47 @@ def _bn_transpose_info(layer): return True, perm_fwd, perm_bwd +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 _emit_param(prefix, name, arr, quantizer, nodes, initializers, use_qonnx, store_integer_weights, out_channels=None): + """Emit the ONNX value for a learnable parameter (kernel/bias/gamma/beta) and return its name""" + if use_qonnx: + fp_name = f"{prefix}_{name}_fp" + initializers.append(onh.from_array(arr, name=fp_name)) + k, i, f = quantizer.get_quantization_bits() + q_nodes, out = _quant_node( + f"{prefix}_{name}", + fp_name, + quantizer.round_mode, + _np(k), + _np(i), + _np(f), + initializers, + overflow_mode=quantizer.overflow, + ) + nodes.extend(q_nodes) + return out + if store_integer_weights: + k, i, f = quantizer.get_quantization_bits() + if out_channels is not None: + k_a = _weight_f_for_onnx(_np(k), out_channels) + i_a = _weight_f_for_onnx(_np(i), out_channels) + f_a = _weight_f_for_onnx(_np(f), out_channels) + else: + k_a, i_a, f_a = _np(k), _np(i), _np(f) + q_nodes, out = _int_weight_node(f"{prefix}_{name}", arr, k_a, i_a, f_a, initializers) + nodes.extend(q_nodes) + return out + out = f"{prefix}_{name}" + initializers.append(onh.from_array(arr, name=out)) + return out + + 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) - ): + if getattr(layer, "input_quantizer", None) is not None and layer.quantize_input and layer.enable_quantization: q = layer.input_quantizer k, i, f = q.get_quantization_bits() new_nodes, current = quant_fn( @@ -291,18 +338,14 @@ def _maybe_quant_input(layer, prefix, current, nodes, initializers, quant_fn): _np(i), _np(f), initializers, - overflow_mode=getattr(q, "overflow", "SAT"), + overflow_mode=q.overflow, ) 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) - ): + if getattr(layer, "output_quantizer", None) is not None and layer.quantize_output and layer.enable_quantization: q = layer.output_quantizer k, i, f = q.get_quantization_bits() new_nodes, current = quant_fn( @@ -313,7 +356,7 @@ def _maybe_quant_output(layer, prefix, current, nodes, initializers, quant_fn): _np(i), _np(f), initializers, - overflow_mode=getattr(q, "overflow", "SAT"), + overflow_mode=q.overflow, ) nodes.extend(new_nodes) return current @@ -347,70 +390,22 @@ def _weight_f_for_onnx(f_np, out_channels): 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)) + q_weight = _emit_param( + prefix, "weight", kernel_np, layer.weight_quantizer, nodes, initializers, use_qonnx, store_integer_weights, out_units + ) 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)) + q_bias = _emit_param( + prefix, "bias", bias_np, layer.bias_quantizer, nodes, initializers, use_qonnx, store_integer_weights + ) gemm_inputs.append(q_bias) gemm_out = f"{prefix}_gemm" @@ -422,7 +417,6 @@ def _add_dense(layer, prefix, current, nodes, initializers, quant_fn, use_qonnx, 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: @@ -441,63 +435,27 @@ def _add_conv(layer, prefix, current, nodes, initializers, ndim, quant_fn, use_q 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)) + q_weight = _emit_param( + prefix, + "weight", + kernel_onnx, + layer.weight_quantizer, + nodes, + initializers, + use_qonnx, + store_integer_weights, + out_channels, + ) 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)) + q_bias = _emit_param( + prefix, "bias", bias_np, layer.bias_quantizer, nodes, initializers, use_qonnx, store_integer_weights + ) conv_inputs.append(q_bias) - # Padding padding = layer.padding if isinstance(padding, str): auto_pad = "SAME_UPPER" if padding == "same" else "VALID" @@ -507,11 +465,10 @@ def _add_conv(layer, prefix, current, nodes, initializers, ndim, quant_fn, use_q 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), + 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, ) @@ -544,65 +501,29 @@ def _add_depthwise_conv(layer, prefix, current, nodes, initializers, quant_fn, u 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)) + q_weight = _emit_param( + prefix, + "weight", + kernel_onnx, + layer.weight_quantizer, + nodes, + initializers, + use_qonnx, + store_integer_weights, + out_channels, + ) 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)) + q_bias = _emit_param( + prefix, "bias", bias_np, layer.bias_quantizer, nodes, initializers, use_qonnx, store_integer_weights + ) conv_inputs.append(q_bias) padding = layer.padding @@ -614,11 +535,10 @@ def _add_depthwise_conv(layer, prefix, current, nodes, initializers, quant_fn, u 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), + 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, ) @@ -659,48 +579,14 @@ def _add_batchnorm(layer, prefix, current, nodes, initializers, quant_fn, use_qo 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)) + qonnx_p = use_qonnx and is_pq + intstore_p = store_integer_weights and is_pq + q_gamma = _emit_param( + prefix, "gamma", gamma_np, layer.weight_quantizer if is_pq else None, nodes, initializers, qonnx_p, intstore_p + ) + q_beta = _emit_param( + prefix, "beta", beta_np, layer.bias_quantizer if is_pq else None, nodes, initializers, qonnx_p, intstore_p + ) mean_name = f"{prefix}_running_mean" var_name = f"{prefix}_running_var" @@ -724,43 +610,14 @@ def _add_batchnorm(layer, prefix, current, nodes, initializers, quant_fn, use_qo 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)) + q_weight = _emit_param( + prefix, "weight", kernel_np, layer.weight_quantizer, nodes, initializers, use_qonnx, store_integer_weights, out_units + ) # Transpose [out, in] → [in, out] so MatMul(input[..., in], kernel_t[in, out]) works kernel_t_name = f"{prefix}_weight_t" @@ -772,28 +629,9 @@ def _add_dense_nd(layer, prefix, current, nodes, initializers, quant_fn, use_qon 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)) + q_bias = _emit_param( + prefix, "bias", bias_np, 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 @@ -802,21 +640,101 @@ def _add_dense_nd(layer, prefix, current, nodes, initializers, quant_fn, use_qon 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). +def _add_quantized_softmax(sm, prefix, current, nodes, initializers, quant_fn, kpm_mask=None): + enable = sm.enable_quantization + scaler = float(sm.input_scaler) + stable = bool(sm.stable) + eps = float(sm.epsilon) - Decomposes multi-head attention into primitive ONNX ops: + def qdq(q, pfx, x): + k, i, f = q.get_quantization_bits() + q_nodes, out = quant_fn(pfx, x, q.round_mode, _np(k), _np(i), _np(f), initializers, overflow_mode=q.overflow) + nodes.extend(q_nodes) + return out - 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 + # 1) Softmax input quantizer. + if sm.quantize_input and enable: + current = qdq(sm.input_quantizer, f"{prefix}_sm_in_q", current) - Returns (out_name, avg_attn_weights_name). - """ + # 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 (only when quantize_input==stable), + # Exp of (-scaler * x) for the stable branch (+scaler otherwise), output QDQ. + exp_t = sm.exp_table + if exp_t.quantize_input and enable: + exp_in = qdq(exp_t.input_quantizer, f"{prefix}_sm_exp_in_q", exp_in) + coeff = -scaler if stable else scaler + exp_arg = exp_in + if coeff != 1.0: + coeff_name = f"{prefix}_sm_exp_coeff" + initializers.append(onh.from_array(np.array(coeff, dtype=np.float32), name=coeff_name)) + exp_arg = f"{prefix}_sm_exp_arg" + nodes.append(oh.make_node("Mul", inputs=[exp_in, coeff_name], outputs=[exp_arg])) + exp_inp = f"{prefix}_sm_exp" + nodes.append(oh.make_node("Exp", inputs=[exp_arg], outputs=[exp_inp])) + if exp_t.quantize_output and enable: + exp_inp = qdq(exp_t.output_quantizer, f"{prefix}_sm_exp_out_q", exp_inp) + + # 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=TensorProto.FLOAT)) + masked = f"{prefix}_sm_masked" + nodes.append(oh.make_node("Mul", inputs=[kpm_f, exp_inp], outputs=[masked])) + exp_inp = masked + + # 4) Sum over the last axis (ReduceSum takes axes as an input from opset 13). + sum_axes = f"{prefix}_sm_sum_axes" + initializers.append(onh.from_array(np.array([-1], dtype=np.int64), name=sum_axes)) + sums = f"{prefix}_sm_sum" + nodes.append(oh.make_node("ReduceSum", inputs=[exp_inp, sum_axes], outputs=[sums], keepdims=1)) + + # 5) Quantized reciprocal table: input QDQ, 1/(x+eps), output QDQ. + inv_t = sm.inv_table + inv_in = sums + if inv_t.quantize_input and enable: + inv_in = qdq(inv_t.input_quantizer, f"{prefix}_sm_inv_in_q", inv_in) + eps_name = f"{prefix}_sm_eps" + initializers.append(onh.from_array(np.array(eps, dtype=np.float32), name=eps_name)) + 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_t.quantize_output and enable: + divisor = qdq(inv_t.output_quantizer, f"{prefix}_sm_inv_out_q", divisor) + + # 6) Multiply numerator by reciprocal. + out = f"{prefix}_sm_out" + nodes.append(oh.make_node("Mul", inputs=[exp_inp, divisor], outputs=[out])) + current = out + + # 7) Softmax output quantizer. + if sm.quantize_output and enable: + current = qdq(sm.output_quantizer, f"{prefix}_sm_out_q", current) + 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, +): H = layer.num_heads head_dim = layer.head_dim E = layer.embed_dim @@ -884,41 +802,24 @@ def _split_heads(x_name, pfx): 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) + 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 = f"{prefix}_kpm_axes" + initializers.append(onh.from_array(np.array([1, 2], dtype=np.int64), name=kpm_axes)) + 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])) + + current = _add_quantized_softmax( + layer.softmax, f"{prefix}_attn", current, nodes, initializers, quant_fn, kpm_mask=kpm_mult + ) + attn_w_name = current # softmax output = attention weights (also averaged over heads below) ctx_raw = f"{prefix}_ctx_raw" nodes.append(oh.make_node("MatMul", inputs=[current, v_h], outputs=[ctx_raw])) @@ -974,16 +875,14 @@ def _add_avgpool(layer, prefix, current, nodes, initializers, ndim, quant_fn): 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), + kernel_shape=_to_list(layer.pool_size, ndim), + strides=_to_list(layer.strides, ndim), pads=[0] * (ndim * 2), count_include_pad=0, ) @@ -1009,10 +908,6 @@ def _add_global_avgpool(layer, prefix, current, nodes, ndim): 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 @@ -1021,21 +916,9 @@ def _add_global_avgpool(layer, prefix, current, nodes, ndim): 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") - ): + if layer.use_multiplier and 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" @@ -1044,15 +927,13 @@ def _add_pq_activation(layer, prefix, current, nodes, initializers, quant_fn): nodes.append(oh.make_node("Mul", inputs=[current, scale_name], outputs=[scaled_out])) current = scaled_out - # --- activation --- - act = getattr(layer, "activation_name", "relu") + act = layer.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": - # hard_tanh(x) = clip(x, -1, 1) cmin_name = f"{prefix}_htanh_min" cmax_name = f"{prefix}_htanh_max" initializers += [ @@ -1074,21 +955,64 @@ def _add_pq_activation(layer, prefix, current, nodes, initializers, quant_fn): # --------------------------------------------------------------------------- +def _resolve_mask_arg(mask, prefix, kind, tensor_to_onnx, initializers): + """Resolve an MHA mask call-argument to an ONNX value name (or None). + + A KerasTensor mask (e.g. a runtime keras.Input) maps through tensor_to_onnx; a + constant array mask (e.g. a fixed causal mask) becomes an initializer. + """ + 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)] + arr = np.asarray(_np(mask)) + name = f"{prefix}_{kind}_const" + initializers.append(onh.from_array(arr, name=name)) + return name + + def _emit_layer( - layer, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights, input_onnx_names=None + 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.""" # --- PQuant layers --- if isinstance(layer, PQMultiheadAttention): - # input_onnx_names = [query, key, value] or [single_input] for self-attention + # input_onnx_names = [query, key, value] (+ any mask tensors, ignored here) + # or [single_input] for self-attention. q/k/v are always the first three. 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) + # Masks are passed as call kwargs on the layer's inbound node. + kwargs = layer._inbound_nodes[0].arguments.kwargs if layer._inbound_nodes 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_integer_weights, + key_padding_mask=kpm, + attn_mask=attn_mask, + ) if isinstance(layer, PQActivation): return _add_pq_activation(layer, prefix, current, nodes, initializers, quant_fn) @@ -1242,11 +1166,10 @@ def _add_conv_plain(layer, prefix, current, nodes, initializers): 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), + 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, ) @@ -1265,11 +1188,6 @@ def _add_conv_plain(layer, prefix, current, nodes, initializers): 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}" @@ -1298,11 +1216,6 @@ def _inbound_input_names(layer, tensor_to_onnx): 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] @@ -1393,6 +1306,7 @@ def convert_to_onnx( 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) @@ -1400,21 +1314,26 @@ def convert_to_onnx( # 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) + n_in = len(model.inputs) + if n_in == 1: + input_names = ["input"] + input_shapes = [tuple(input_shape)] + else: + input_names = [f"input_{i}" for i in range(n_in)] + 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 n_in == 1 else dummies, 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)) - ] + 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, output_shape) graph = oh.make_graph( @@ -1431,10 +1350,6 @@ def convert_to_onnx( 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[:] @@ -1445,34 +1360,3 @@ def convert_to_onnx( 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/torch/convert_to_onnx.py b/src/pquant/core/torch/convert_to_onnx.py index a61aeaa..69a0cc1 100644 --- a/src/pquant/core/torch/convert_to_onnx.py +++ b/src/pquant/core/torch/convert_to_onnx.py @@ -71,14 +71,8 @@ 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: + if f.numel() > 1: i = i.reshape(-1).max() f = f.reshape(-1).min() i_val = float(i.item()) @@ -117,12 +111,6 @@ def _quant_node(name_prefix, input_name, rounding_mode, k, i, f, initializers, o 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()) @@ -190,12 +178,12 @@ def _int_weight_node(name_prefix, weight_np, k, i, f, initializers): # noqa: AR Returns ([node], output_name). """ - k_val = int(k.item()) if hasattr(k, "item") else int(k) + k_val = int(k.item()) 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) + f_t = f.detach().cpu() if f_t.numel() == 1: # per-tensor @@ -238,56 +226,59 @@ def _torch_padding_to_onnx(padding, ndim): return list(padding) + list(padding) +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 _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) - ): + # input_quantizer is created conditionally, so guard it; the bool flags are always present. + if getattr(module, "input_quantizer", None) is not None and module.quantize_input and module.enable_quantization: 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") - ) + new_nodes, current = quant_fn(f"{prefix}_in", current, q.round_mode, k, i, f, initializers, overflow_mode=q.overflow) 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) - ): + if getattr(module, "output_quantizer", None) is not None and module.quantize_output and module.enable_quantization: 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") + f"{prefix}_out", current, q.round_mode, k, i, f, initializers, overflow_mode=q.overflow ) nodes.extend(new_nodes) return current +def _emit_param(prefix, name, arr, quantizer, nodes, initializers, use_qonnx, store_integer_weights): + if use_qonnx: + fp_name = f"{prefix}_{name}_fp" + initializers.append(onh.from_array(arr, name=fp_name)) + 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, i, f = quantizer.get_quantization_bits() + q_nodes, out = _int_weight_node(f"{prefix}_{name}", arr, k, i, f, initializers) + nodes.extend(q_nodes) + return out + out = f"{prefix}_{name}" + initializers.append(onh.from_array(arr, name=out)) + return out + + # --------------------------------------------------------------------------- # 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)): + 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") # --- Input: Clip + QuantizeLinear → int8 (stop before DequantizeLinear) --- @@ -323,11 +314,11 @@ def _add_dense_integer(module, prefix, current, nodes, initializers): # 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) + k_w_val = int(k_w.item()) # get_quantization_bits() always returns tensors 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) + f_w_t = f_w.detach().cpu() if f_w_t.numel() == 1: f_w_1d = np.array([float(f_w_t.item())]) per_channel_w = False @@ -406,41 +397,16 @@ def _add_dense_integer(module, prefix, current, nodes, initializers): 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"), + 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 ) - 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 + 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" initializers.append(onh.from_array(weight_np.T, name=q_weight)) # pre-transposed [in, out] @@ -451,28 +417,9 @@ def _add_dense_nd(module, prefix, current, nodes, initializers, quant_fn, use_qo 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)) + q_bias = _emit_param( + prefix, "bias", bias_np, 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 @@ -487,57 +434,17 @@ def _add_dense(module, prefix, current, nodes, initializers, quant_fn, use_qonnx 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)) + q_weight = _emit_param( + prefix, "weight", weight_np, module.weight_quantizer, nodes, initializers, use_qonnx, store_integer_weights + ) - # 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)) + q_bias = _emit_param( + prefix, "bias", bias_np, module.bias_quantizer, nodes, initializers, use_qonnx, store_integer_weights + ) gemm_inputs.append(q_bias) gemm_out = f"{prefix}_gemm" @@ -552,55 +459,17 @@ def _add_conv(module, prefix, current, nodes, initializers, ndim, quant_fn, use_ 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)) + q_weight = _emit_param( + prefix, "weight", weight_np, module.weight_quantizer, nodes, initializers, use_qonnx, store_integer_weights + ) 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)) + q_bias = _emit_param( + prefix, "bias", bias_np, module.bias_quantizer, nodes, initializers, use_qonnx, store_integer_weights + ) conv_inputs.append(q_bias) padding = module.padding @@ -611,11 +480,10 @@ def _add_conv(module, prefix, current, nodes, initializers, ndim, quant_fn, use_ 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), + 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, ) @@ -636,48 +504,12 @@ def _add_batchnorm(module, prefix, current, nodes, initializers, quant_fn, use_q 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)) + q_gamma = _emit_param( + prefix, "gamma", gamma_np, module.weight_quantizer, nodes, initializers, use_qonnx, store_integer_weights + ) + q_beta = _emit_param( + prefix, "beta", beta_np, module.bias_quantizer, nodes, initializers, use_qonnx, store_integer_weights + ) mean_name = f"{prefix}_running_mean" var_name = f"{prefix}_running_var" @@ -697,7 +529,6 @@ def _add_batchnorm(module, prefix, current, nodes, initializers, quant_fn, use_q 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 = ( @@ -713,50 +544,29 @@ def _add_layernorm(module, prefix, current, nodes, initializers, quant_fn, use_q 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, + qonnx_p = use_qonnx and has_weight + intstore_p = store_integer_weights and has_weight + q_gamma = _emit_param( + prefix, + "gamma", + gamma_np, + module.weight_quantizer if has_weight else None, + nodes, + initializers, + qonnx_p, + intstore_p, + ) + if has_bias: + q_beta = _emit_param( + prefix, + "beta", + beta_np, + module.bias_quantizer if has_weight else None, + nodes, initializers, - overflow_mode=getattr(module.weight_quantizer, "overflow", "SAT"), + qonnx_p, + intstore_p, ) - 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: @@ -779,15 +589,14 @@ def _add_layernorm(module, prefix, current, nodes, initializers, quant_fn, use_q 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), + 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), @@ -804,30 +613,92 @@ def _add_avgpool(module, prefix, current, nodes, initializers, ndim, quant_fn): # --------------------------------------------------------------------------- -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. +def _add_quantized_softmax(sm, prefix, current, nodes, initializers, quant_fn, kpm_mask=None): + enable = sm.enable_quantization + scaler = float(sm.input_scaler) + stable = bool(sm.stable) + eps = float(sm.epsilon) - Decomposes multi-head attention into primitive ONNX ops: + def qdq(q, pfx, x): + k, i, f = q.get_quantization_bits() + q_nodes, out = quant_fn(pfx, x, q.round_mode, k, i, f, initializers, overflow_mode=q.overflow) + nodes.extend(q_nodes) + return out - [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] + if sm.quantize_input and enable: + current = qdq(sm.input_quantizer, f"{prefix}_sm_in_q", current) - 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. + 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 + + exp_t = sm.exp_table + if exp_t.quantize_input and enable: + exp_in = qdq(exp_t.input_quantizer, f"{prefix}_sm_exp_in_q", exp_in) + coeff = -scaler if stable else scaler + exp_arg = exp_in + if coeff != 1.0: + coeff_name = f"{prefix}_sm_exp_coeff" + initializers.append(onh.from_array(np.array(coeff, dtype=np.float32), name=coeff_name)) + exp_arg = f"{prefix}_sm_exp_arg" + nodes.append(oh.make_node("Mul", inputs=[exp_in, coeff_name], outputs=[exp_arg])) + exp_inp = f"{prefix}_sm_exp" + nodes.append(oh.make_node("Exp", inputs=[exp_arg], outputs=[exp_inp])) + if exp_t.quantize_output and enable: + exp_inp = qdq(exp_t.output_quantizer, f"{prefix}_sm_exp_out_q", exp_inp) + + 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=TensorProto.FLOAT)) + masked = f"{prefix}_sm_masked" + nodes.append(oh.make_node("Mul", inputs=[kpm_f, exp_inp], outputs=[masked])) + exp_inp = masked + + sum_axes = f"{prefix}_sm_sum_axes" + initializers.append(onh.from_array(np.array([-1], dtype=np.int64), name=sum_axes)) + sums = f"{prefix}_sm_sum" + nodes.append(oh.make_node("ReduceSum", inputs=[exp_inp, sum_axes], outputs=[sums], keepdims=1)) + + inv_t = sm.inv_table + inv_in = sums + if inv_t.quantize_input and enable: + inv_in = qdq(inv_t.input_quantizer, f"{prefix}_sm_inv_in_q", inv_in) + eps_name = f"{prefix}_sm_eps" + initializers.append(onh.from_array(np.array(eps, dtype=np.float32), name=eps_name)) + 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_t.quantize_output and enable: + divisor = qdq(inv_t.output_quantizer, f"{prefix}_sm_inv_out_q", divisor) + + out = f"{prefix}_sm_out" + nodes.append(oh.make_node("Mul", inputs=[exp_inp, divisor], outputs=[out])) + current = out + + if sm.quantize_output and enable: + current = qdq(sm.output_quantizer, f"{prefix}_sm_out_q", current) + return current - 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. - """ + +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, +): H = module.num_heads head_dim = module.head_dim E = module.embed_dim @@ -907,50 +778,29 @@ def _split_heads(x_name, pfx): 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) + 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 = f"{prefix}_kpm_axes" + initializers.append(onh.from_array(np.array([1, 2], dtype=np.int64), name=kpm_axes)) + 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])) + + current = _add_quantized_softmax( + module.softmax, f"{prefix}_attn", current, nodes, initializers, quant_fn, kpm_mask=kpm_mult + ) + attn_w_name = current # softmax output = attention weights (also averaged over heads below) - # --- 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" @@ -1205,9 +1055,7 @@ def _emit_module( # 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") - ) + new_nodes, out = quant_fn(prefix, current, module.round_mode, k, i, f, initializers, overflow_mode=module.overflow) nodes.extend(new_nodes) return out raise TypeError(f"Unsupported module type for ONNX export: {type(module).__name__}") @@ -1259,48 +1107,24 @@ def convert_to_onnx( 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, + Note: + This is a thin wrapper over convert_to_onnx_fx(). An ``nn.Sequential`` is a + plain linear chain, so torch.fx always traces it successfully; routing through + the FX converter keeps a single code path for both linear and branched models. + """ + return convert_to_onnx_fx( + model, + input_shape, + output_path=output_path, + opset=opset, + use_qonnx=use_qonnx, + store_integer_weights=store_integer_weights, + integer_ops=integer_ops, + include_clip=include_clip, + batch_size=batch_size, ) - 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 @@ -1321,38 +1145,6 @@ def export_qdq_layernorm( 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): @@ -1479,8 +1271,6 @@ def _check_q_int16(arr: np.ndarray, frac_bits: int, name: str) -> None: class _PQTracer(_fx.Tracer): - """Tracer that treats all PQuant layer types (and standard torch.nn leaves) as atomic.""" - _LEAF_TYPES = ( PQDense, PQConv2d, @@ -1499,6 +1289,48 @@ 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 convert_to_onnx_fx( model: nn.Module, input_shape: tuple, @@ -1508,6 +1340,9 @@ def convert_to_onnx_fx( 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 any PQuant nn.Module to ONNX using torch.fx symbolic tracing. @@ -1516,23 +1351,68 @@ def convert_to_onnx_fx( 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. + 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: + 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. + + Remaining args match the per-layer quantization behaviour described in the module docstring. """ model.eval() quant_fn = _quant_node if use_qonnx else functools.partial(_qdq_node, include_clip=include_clip) - graph = _PQTracer().trace(model) + input_shapes = _normalize_input_shapes(input_shape) + input_torch_dtypes, input_tp_dtypes = _normalize_input_dtypes(input_dtypes, len(input_shapes)) + + graph = _PQTracer().trace(model, concrete_args=concrete_args) 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. + 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() + + tensor_phs = list(gm.graph.find_nodes(op="placeholder")) + if len(tensor_phs) != len(input_shapes): + raise ValueError( + f"FX export: model.forward has {len(tensor_phs)} tensor input(s) but " + f"input_shape describes {len(input_shapes)}. Specialize non-tensor " + f"arguments via concrete_args={{...}}." + ) + input_names = ["input"] if len(tensor_phs) == 1 else [str(p.target) for p in tensor_phs] + ph_to_name = {p: n for p, n in zip(tensor_phs, input_names)} + 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) + 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(torch.zeros(1, *input_shape, device=device)) + ShapeProp(gm).propagate(*probes) onnx_nodes: list[onnx.NodeProto] = [] initializers: list[onnx.TensorProto] = [] @@ -1545,8 +1425,6 @@ def _res(arg) -> str: 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): @@ -1581,11 +1459,9 @@ def _resolve_perm_dims(args, rank: int) -> list[int]: for node in gm.graph.nodes: if node.op == "placeholder": - node_to_name[node] = "input" + node_to_name[node] = ph_to_name[node] 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) @@ -1598,10 +1474,21 @@ def _resolve_perm_dims(args, rank: int) -> list[int]: 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, ...]) + # forward(query, key, value, key_padding_mask=None, attn_mask=None, ...) 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 + + def _mask_name(pos, kw, node=node): + 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 node_to_name[arg] + + kpm_name = _mask_name(3, "key_padding_mask") + attn_mask_name = _mask_name(4, "attn_mask") out_name, avg_attn_name = _add_mha( mod, mod_prefix, @@ -1613,6 +1500,8 @@ def _resolve_perm_dims(args, rank: int) -> list[int]: quant_fn, use_qonnx, store_integer_weights, + key_padding_mask=kpm_name, + attn_mask=attn_mask_name, ) # Store tuple so operator.getitem(node, 0/1) resolves correctly. node_to_name[node] = (out_name, avg_attn_name) @@ -1633,6 +1522,9 @@ def _resolve_perm_dims(args, rank: int) -> list[int]: elif node.op == "call_function": fn = node.target + if fn is torch._assert or getattr(fn, "__name__", "") == "_assert" or fn is _operator.eq: + continue + if fn is _operator.getitem: # Unpack a tuple output (e.g. from PQMultiheadAttention). container = node_to_name[node.args[0]] @@ -1767,20 +1659,31 @@ def _resolve_perm_dims(args, rank: int) -> list[int]: # MHA nodes store a tuple (out, avg_attn); expose the attention output. output_names.append(val[0] if isinstance(val, tuple) else val) + graph_input_names = set(input_names) + for idx, nm in enumerate(output_names): + if nm in graph_input_names: + ident = f"{nm}_identity_out{idx}" + onnx_nodes.append(oh.make_node("Identity", inputs=[nm], outputs=[ident])) + output_names[idx] = ident + with torch.no_grad(): - dummy_out = model(torch.zeros(1, *input_shape, device=device)) + dummy_out = model(*probes, **(concrete_args or {})) 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]) + 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, [None] + list(t.shape[1:])) + oh.make_tensor_value_info(name, TensorProto.FLOAT, [batch_dim] + 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], + inputs=input_vis, outputs=output_vis, initializer=initializers, ) @@ -1796,64 +1699,3 @@ def _resolve_perm_dims(args, rank: int) -> list[int]: 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/tests/test_keras_onnx_converter.py b/tests/test_keras_onnx_converter.py index 91dc46f..b7e840e 100644 --- a/tests/test_keras_onnx_converter.py +++ b/tests/test_keras_onnx_converter.py @@ -15,36 +15,39 @@ 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") ATOL = 1e-4 +# When quantization is enabled, torch/keras fake-quant and ONNX QuantizeLinear can round +# a few values to opposite sides of a 0.5 boundary (op/accumulation-order ULP differences), +# so allow ~1 quantization level of slack for graphs that re-quantize intermediates. +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): + # Run every cfg-based test twice: float path and quantization-enabled, so the + # emitted Quantize/DequantizeLinear nodes are actually exercised against onnxruntime + # (they are skipped entirely when enable_quantization is False). c = pquant.cs_config() - c.quantization_parameters.enable_quantization = False + c.quantization_parameters.enable_quantization = request.param return c -# --------------------------------------------------------------------------- -# helpers -# --------------------------------------------------------------------------- - - def _channels_first(): return keras.backend.image_data_format() == "channels_first" @@ -63,11 +66,6 @@ def _onnx_run(model, x: np.ndarray, input_shape: tuple, tmp_path) -> np.ndarray: return sess.run(None, {in_name: x})[0] -# --------------------------------------------------------------------------- -# PQDense -# --------------------------------------------------------------------------- - - @pytest.mark.parametrize("bias", [True, False]) def test_dense_onnx(cfg, bias, tmp_path): IN, OUT = 16, 8 @@ -82,12 +80,7 @@ def test_dense_onnx(cfg, bias, tmp_path): 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") - - -# --------------------------------------------------------------------------- -# PQConv2d -# --------------------------------------------------------------------------- + np.testing.assert_allclose(keras_out, onnx_out, atol=_atol(cfg), err_msg=f"PQDense bias={bias}: keras vs ONNX mismatch") @pytest.mark.parametrize("bias", [True, False]) @@ -110,12 +103,7 @@ def test_conv2d_onnx(cfg, bias, tmp_path): 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") - - -# --------------------------------------------------------------------------- -# PQConv1d -# --------------------------------------------------------------------------- + np.testing.assert_allclose(keras_out, onnx_out, atol=_atol(cfg), err_msg=f"PQConv2d bias={bias}: keras vs ONNX mismatch") @pytest.mark.parametrize("bias", [True, False]) @@ -138,12 +126,7 @@ def test_conv1d_onnx(cfg, bias, tmp_path): 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") - - -# --------------------------------------------------------------------------- -# PQBatchNormalization -# --------------------------------------------------------------------------- + np.testing.assert_allclose(keras_out, onnx_out, atol=_atol(cfg), err_msg=f"PQConv1d bias={bias}: keras vs ONNX mismatch") def test_batchnorm_onnx(cfg, tmp_path): @@ -167,12 +150,7 @@ def test_batchnorm_onnx(cfg, tmp_path): 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") - - -# --------------------------------------------------------------------------- -# PQDepthwiseConv2d -# --------------------------------------------------------------------------- + np.testing.assert_allclose(keras_out, onnx_out, atol=_atol(cfg), err_msg="PQBatchNormalization: keras vs ONNX mismatch") def test_depthwise_conv2d_onnx(cfg, tmp_path): @@ -194,4 +172,136 @@ def test_depthwise_conv2d_onnx(cfg, tmp_path): 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") + np.testing.assert_allclose(keras_out, onnx_out, atol=_atol(cfg), err_msg="PQDepthwiseConv2d: keras vs ONNX mismatch") + + +@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) + + x_np = np.random.randn(4, DIM).astype(np.float32) + keras_out = _keras_out(model, x_np) + onnx_out = _onnx_run(model, x_np, input_shape=(DIM,), tmp_path=tmp_path) + np.testing.assert_allclose( + keras_out, 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) + + x_np = np.random.randn(4, DIM).astype(np.float32) + keras_out = _keras_out(model, x_np) + onnx_out = _onnx_run(model, x_np, input_shape=(DIM,), tmp_path=tmp_path) + np.testing.assert_allclose(keras_out, onnx_out, atol=_atol(cfg), err_msg="residual+concat: keras vs ONNX mismatch") + + +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) + + xa = np.random.randn(3, IN_A).astype(np.float32) + xb = np.random.randn(3, IN_B).astype(np.float32) + keras_out = _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] + + 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_out, onnx_out, atol=_atol(cfg), err_msg="two-input: keras vs ONNX mismatch") + + +@pytest.mark.parametrize("bias", [True, False]) +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) + + x_np = np.random.randn(2, T, E).astype(np.float32) + keras_out = _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_out, onnx_out, atol=_atol(cfg), err_msg=f"PQMultiheadAttention 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) + + x_np = np.random.randn(2, T, E).astype(np.float32) + keras_out = _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_out, onnx_out, atol=_atol(cfg), err_msg="MHA causal attn_mask: keras vs ONNX mismatch") + + +def test_mha_key_padding_mask_onnx(cfg, tmp_path): + import onnx + + 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) + + model([np.zeros((1, T, E), np.float32), np.zeros((1, T), bool)]) + apply_final_compression(model) + + 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_out = _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 + + 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_out, onnx_out, atol=_atol(cfg), err_msg="MHA key_padding_mask: keras vs ONNX mismatch") diff --git a/tests/test_torch_onnx_converter.py b/tests/test_torch_onnx_converter.py index 2064d30..5edeadf 100644 --- a/tests/test_torch_onnx_converter.py +++ b/tests/test_torch_onnx_converter.py @@ -23,7 +23,9 @@ convert_to_onnx_fx, export_qdq_layernorm, ) +from pquant.core.torch.layers import Quantizer # noqa: E402 from pquant.layers import ( # noqa: E402 + PQActivation, PQAvgPool1d, PQAvgPool2d, PQBatchNorm1d, @@ -31,12 +33,22 @@ 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 +# When quantization is enabled, torch fake-quant and ONNX QuantizeLinear can round a +# few values to opposite sides of a 0.5 boundary (the rounding inputs differ by float +# ULPs from differing op/accumulation order), so allow ~1 quantization level of slack +# for graphs that re-quantize intermediate activations. +QUANT_ATOL = 5e-3 + + +def _atol(cfg): + return QUANT_ATOL if cfg.quantization_parameters.enable_quantization else ATOL # --------------------------------------------------------------------------- @@ -44,16 +56,23 @@ # --------------------------------------------------------------------------- -@pytest.fixture -def cfg(): +@pytest.fixture(params=[False, True], ids=["float", "quant"]) +def cfg(request): + # Run every cfg-based test twice: once with the plain float path and once with + # quantization enabled so the emitted Quantize/DequantizeLinear nodes are + # actually exercised against onnxruntime (they are skipped entirely when + # enable_quantization is False). 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(): + # Quantization-enabled config for tests that specifically target the QDQ path. + c = pquant.cs_config() + c.quantization_parameters.enable_quantization = True + return c def _apply_compression(model: nn.Module): @@ -86,11 +105,6 @@ def _torch_out(model: nn.Module, x: torch.Tensor) -> np.ndarray: return model(x).cpu().numpy() -# --------------------------------------------------------------------------- -# PQDense -# --------------------------------------------------------------------------- - - @pytest.mark.parametrize("bias", [True, False]) def test_dense_onnx(cfg, bias, tmp_path): IN, OUT = 16, 8 @@ -108,11 +122,6 @@ def test_dense_onnx(cfg, bias, tmp_path): np.testing.assert_allclose(torch_out, onnx_out, atol=ATOL, err_msg=f"PQDense bias={bias}: torch vs ONNX mismatch") -# --------------------------------------------------------------------------- -# PQConv2d -# --------------------------------------------------------------------------- - - @pytest.mark.parametrize("bias", [True, False]) def test_conv2d_onnx(cfg, bias, tmp_path): IN_C, OUT_C, H, W = 3, 8, 8, 8 @@ -130,11 +139,6 @@ def test_conv2d_onnx(cfg, bias, tmp_path): np.testing.assert_allclose(torch_out, onnx_out, atol=ATOL, err_msg=f"PQConv2d bias={bias}: torch vs ONNX mismatch") -# --------------------------------------------------------------------------- -# PQConv1d -# --------------------------------------------------------------------------- - - @pytest.mark.parametrize("bias", [True, False]) def test_conv1d_onnx(cfg, bias, tmp_path): IN_C, OUT_C, L = 4, 8, 16 @@ -152,11 +156,6 @@ def test_conv1d_onnx(cfg, bias, tmp_path): np.testing.assert_allclose(torch_out, onnx_out, atol=ATOL, err_msg=f"PQConv1d bias={bias}: torch vs ONNX mismatch") -# --------------------------------------------------------------------------- -# PQBatchNorm2d -# --------------------------------------------------------------------------- - - def test_batchnorm2d_onnx(cfg, tmp_path): C, H, W = 8, 4, 4 model = nn.Sequential( @@ -174,11 +173,6 @@ def test_batchnorm2d_onnx(cfg, tmp_path): np.testing.assert_allclose(torch_out, onnx_out, atol=ATOL, err_msg="PQBatchNorm2d: torch vs ONNX mismatch") -# --------------------------------------------------------------------------- -# PQBatchNorm1d -# --------------------------------------------------------------------------- - - def test_batchnorm1d_onnx(cfg, tmp_path): C, L = 8, 16 model = nn.Sequential( @@ -196,11 +190,6 @@ def test_batchnorm1d_onnx(cfg, tmp_path): np.testing.assert_allclose(torch_out, onnx_out, atol=ATOL, err_msg="PQBatchNorm1d: torch vs ONNX mismatch") -# --------------------------------------------------------------------------- -# PQAvgPool2d -# --------------------------------------------------------------------------- - - def test_avgpool2d_onnx(cfg, tmp_path): C, H, W = 8, 8, 8 model = nn.Sequential( @@ -216,11 +205,6 @@ def test_avgpool2d_onnx(cfg, tmp_path): np.testing.assert_allclose(torch_out, onnx_out, atol=ATOL, err_msg="PQAvgPool2d: torch vs ONNX mismatch") -# --------------------------------------------------------------------------- -# PQAvgPool1d -# --------------------------------------------------------------------------- - - def test_avgpool1d_onnx(cfg, tmp_path): C, L = 8, 16 model = nn.Sequential( @@ -236,11 +220,6 @@ def test_avgpool1d_onnx(cfg, tmp_path): np.testing.assert_allclose(torch_out, onnx_out, atol=ATOL, err_msg="PQAvgPool1d: torch vs ONNX mismatch") -# --------------------------------------------------------------------------- -# PQMultiheadAttention (uses FX converter; self-attention, batch_first=True) -# --------------------------------------------------------------------------- - - class _SelfAttnModel(nn.Module): """Thin wrapper so FX tracing sees a single-input model.""" @@ -264,16 +243,91 @@ def test_mha_onnx(cfg, bias, tmp_path): model(x) _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_out = _torch_out(model, x) onnx_out = _onnx_run_fx(model, x, input_shape=(T, E), tmp_path=tmp_path) np.testing.assert_allclose( - torch_out, onnx_out, atol=ATOL, err_msg=f"PQMultiheadAttention bias={bias}: torch vs ONNX mismatch" + torch_out, onnx_out, atol=_atol(cfg), err_msg=f"PQMultiheadAttention bias={bias}: torch vs ONNX mismatch" ) -# --------------------------------------------------------------------------- -# Static-QDQ LayerNormalization graph -# --------------------------------------------------------------------------- +class _CausalSelfAttnModel(nn.Module): + """Self-attention with a constant additive causal mask (the decoder-inference case).""" + + 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_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) + + torch_out = _torch_out(model, x) + onnx_out = _onnx_run_fx(model, x, input_shape=(T, E), tmp_path=tmp_path) + np.testing.assert_allclose( + torch_out, onnx_out, atol=_atol(cfg), err_msg=f"MHA causal attn_mask bias={bias}: torch vs ONNX mismatch" + ) + + +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 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, key_padding_mask) + _apply_compression(model) + + model.eval() + with torch.no_grad(): + torch_out = model(x, key_padding_mask).cpu().numpy() + + path = str(tmp_path / "mha_kpm.onnx") + proto = convert_to_onnx_fx(model, input_shape=[(T, E), (T,)], output_path=path, input_dtypes=["float32", "bool"]) + + # 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] + + np.testing.assert_allclose( + torch_out, onnx_out, atol=_atol(cfg), err_msg=f"MHA key_padding_mask bias={bias}: torch vs ONNX mismatch" + ) @pytest.mark.parametrize("input_shape", [(4, 64), (1, 4, 64)]) @@ -353,6 +407,404 @@ def test_qdq_layernorm_export(input_shape, tmp_path): np.testing.assert_allclose(onnx_out, y_ref, atol=out_scale * 0.5) +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(a, b) # warm-up + _apply_compression(model) + + model.eval() + with torch.no_grad(): + torch_out = model(a, b).cpu().numpy() + + path = str(tmp_path / "two_input.onnx") + model_proto = convert_to_onnx_fx(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] + + np.testing.assert_allclose(torch_out, onnx_out, atol=ATOL, err_msg=f"two-input bias={bias}: torch vs ONNX mismatch") + + +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_fx(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) + + +@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_out = model(x, scale_up).cpu().numpy() + + path = str(tmp_path / f"flag_{scale_up}.onnx") + model_proto = convert_to_onnx_fx(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_out, onnx_out, atol=ATOL, err_msg=f"concrete_args scale_up={scale_up}: torch vs ONNX mismatch" + ) + + +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) + + model.eval() + with torch.no_grad(): + torch_out = model(x).cpu().numpy() + + path = str(tmp_path / "residual_concat.onnx") + model_proto = convert_to_onnx_fx(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 + + 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=QUANT_ATOL, err_msg="residual+concat: torch vs ONNX mismatch") + + +@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) + + torch_out = _torch_out(model, x) + onnx_out = _onnx_run(model, x, input_shape=(DIM,), tmp_path=tmp_path) + np.testing.assert_allclose(torch_out, 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) + + # 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_out = _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_out, onnx_out, atol=ATOL, err_msg="standalone Quantizer: torch vs ONNX mismatch") + + +class _CNNFlattenModel(nn.Module): + def __init__(self, cfg, in_c: int, hw: int, out: int, use_reshape: bool): + super().__init__() + 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): + 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("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, IN_C, HW, HW) + with torch.no_grad(): + model(x) + _apply_compression(model) + + model.eval() + with torch.no_grad(): + torch_out = model(x).cpu().numpy() + + path = str(tmp_path / f"cnn_flatten_{use_reshape}.onnx") + convert_to_onnx_fx(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=QUANT_ATOL, err_msg=f"CNN→Dense reshape={use_reshape}: torch vs ONNX mismatch" + ) + + +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) + + +def test_scalar_ops_onnx(cfg_quant, tmp_path): + DIM = 16 + model = _ScalarOpsModel(cfg_quant, DIM) + + x = torch.randn(4, DIM) + with torch.no_grad(): + model(x) + _apply_compression(model) + + model.eval() + with torch.no_grad(): + torch_out = model(x).cpu().numpy() + + path = str(tmp_path / "scalar_ops.onnx") + model_proto = convert_to_onnx_fx(model, input_shape=(DIM,), output_path=path) + op_types = [n.op_type for n in model_proto.graph.node] + for expected in ("Mul", "Sub", "Div", "Sigmoid"): + assert expected in op_types, f"missing {expected} node" + + 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="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_fx(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_out = 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_out, 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), + (lambda: nn.Sequential(nn.BatchNorm2d(3)), (3, 8, 8), 2), + ], + ids=["leaky_relu", "maxpool2d", "upsample", "dropout", "batchnorm2d"], +) +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_out = _torch_out(model, x) + onnx_out = _onnx_run(model, x, input_shape=input_shape, tmp_path=tmp_path) + np.testing.assert_allclose(torch_out, 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_out = _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_out, onnx_out, atol=ATOL, err_msg=f"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) + + def test_qdq_layernorm_validation(tmp_path): path = str(tmp_path / "bad.onnx") D = 64 From 59f39c7f47ab5596eb7d47d2b772b0a7186cae69 Mon Sep 17 00:00:00 2001 From: Roope Niemi Date: Thu, 9 Jul 2026 19:22:58 +0200 Subject: [PATCH 2/8] explicit build/compute_output_shape for keras PQMultiheadAttention --- src/pquant/core/keras/layers.py | 36 +++++++++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/src/pquant/core/keras/layers.py b/src/pquant/core/keras/layers.py index 4a058c0..50d80ad 100644 --- a/src/pquant/core/keras/layers.py +++ b/src/pquant/core/keras/layers.py @@ -1955,6 +1955,42 @@ 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): + # Resolve (query, key, value) shapes from either a single shape (self-attention) + # or a list/tuple of per-input shapes; missing key/value fall back to query/key. + 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): + # Provide static output shapes so Keras does not run call() symbolically for + # shape inference (which would build the quantized/pruned sublayers inside a + # scratch graph and fail). Mirrors the (output, avg_attn_weights) tuple that + # call() returns: output is (B, Tq, embed_dim), attn weights are (B, Tq, Tk). + 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): + # Build the projection/softmax sublayers explicitly. Without this Keras would + # try to auto-build the layer by tracing call() in a scratch FuncGraph, which + # fails for the quantized/pruned PQDense sublayers (their build() creates tensors + # that escape the scratch graph). Mirrors how PQDense itself defines build(). + 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, From d999f73b06a809e5f5120509ac22a74def55838f Mon Sep 17 00:00:00 2001 From: Roope Niemi Date: Fri, 10 Jul 2026 15:35:00 +0200 Subject: [PATCH 3/8] split onnx conversion logic to multiple files, clean up code --- src/pquant/core/keras/convert_to_onnx.py | 1362 ------------- src/pquant/core/keras/layers.py | 10 - src/pquant/core/keras/onnx/__init__.py | 5 + src/pquant/core/keras/onnx/convert_to_onnx.py | 417 ++++ src/pquant/core/keras/onnx/helpers.py | 405 ++++ src/pquant/core/keras/onnx/layers.py | 579 ++++++ src/pquant/core/torch/convert_to_onnx.py | 1701 ----------------- src/pquant/core/torch/onnx/__init__.py | 5 + src/pquant/core/torch/onnx/convert_to_onnx.py | 854 +++++++++ src/pquant/core/torch/onnx/helpers.py | 304 +++ src/pquant/core/torch/onnx/layers.py | 580 ++++++ tests/test_keras_onnx_converter.py | 200 +- tests/test_torch_onnx_converter.py | 285 ++- 13 files changed, 3382 insertions(+), 3325 deletions(-) delete mode 100644 src/pquant/core/keras/convert_to_onnx.py create mode 100644 src/pquant/core/keras/onnx/__init__.py create mode 100644 src/pquant/core/keras/onnx/convert_to_onnx.py create mode 100644 src/pquant/core/keras/onnx/helpers.py create mode 100644 src/pquant/core/keras/onnx/layers.py delete mode 100644 src/pquant/core/torch/convert_to_onnx.py create mode 100644 src/pquant/core/torch/onnx/__init__.py create mode 100644 src/pquant/core/torch/onnx/convert_to_onnx.py create mode 100644 src/pquant/core/torch/onnx/helpers.py create mode 100644 src/pquant/core/torch/onnx/layers.py 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 57aa39b..0000000 --- a/src/pquant/core/keras/convert_to_onnx.py +++ /dev/null @@ -1,1362 +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 _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 _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 _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 _emit_param(prefix, name, arr, quantizer, nodes, initializers, use_qonnx, store_integer_weights, out_channels=None): - """Emit the ONNX value for a learnable parameter (kernel/bias/gamma/beta) and return its name""" - if use_qonnx: - fp_name = f"{prefix}_{name}_fp" - initializers.append(onh.from_array(arr, name=fp_name)) - k, i, f = quantizer.get_quantization_bits() - q_nodes, out = _quant_node( - f"{prefix}_{name}", - fp_name, - quantizer.round_mode, - _np(k), - _np(i), - _np(f), - initializers, - overflow_mode=quantizer.overflow, - ) - nodes.extend(q_nodes) - return out - if store_integer_weights: - k, i, f = quantizer.get_quantization_bits() - if out_channels is not None: - k_a = _weight_f_for_onnx(_np(k), out_channels) - i_a = _weight_f_for_onnx(_np(i), out_channels) - f_a = _weight_f_for_onnx(_np(f), out_channels) - else: - k_a, i_a, f_a = _np(k), _np(i), _np(f) - q_nodes, out = _int_weight_node(f"{prefix}_{name}", arr, k_a, i_a, f_a, initializers) - nodes.extend(q_nodes) - return out - out = f"{prefix}_{name}" - initializers.append(onh.from_array(arr, name=out)) - return out - - -def _maybe_quant_input(layer, prefix, current, nodes, initializers, quant_fn): - if getattr(layer, "input_quantizer", None) is not None and layer.quantize_input and layer.enable_quantization: - 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=q.overflow, - ) - 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 layer.quantize_output and layer.enable_quantization: - 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=q.overflow, - ) - 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): - 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] - - q_weight = _emit_param( - prefix, "weight", kernel_np, layer.weight_quantizer, nodes, initializers, use_qonnx, store_integer_weights, out_units - ) - - gemm_inputs = [current, q_weight] - - if layer._bias is not None: - bias_np = _np(layer._bias) - q_bias = _emit_param( - prefix, "bias", bias_np, layer.bias_quantizer, nodes, initializers, use_qonnx, store_integer_weights - ) - 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): - 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] - - q_weight = _emit_param( - prefix, - "weight", - kernel_onnx, - layer.weight_quantizer, - nodes, - initializers, - use_qonnx, - store_integer_weights, - out_channels, - ) - - conv_inputs = [current, q_weight] - - if layer._bias is not None: - bias_np = _np(layer._bias) - q_bias = _emit_param( - prefix, "bias", bias_np, layer.bias_quantizer, nodes, initializers, use_qonnx, store_integer_weights - ) - 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] * ndim - pads = p + p # ONNX format: [begin_0, begin_1, ..., end_0, end_1, ...] - auto_pad = "NOTSET" - - 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] - 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] - - q_weight = _emit_param( - prefix, - "weight", - kernel_onnx, - layer.weight_quantizer, - nodes, - initializers, - use_qonnx, - store_integer_weights, - out_channels, - ) - - conv_inputs = [current, q_weight] - - if layer._bias is not None: - bias_np = _np(layer._bias) - q_bias = _emit_param( - prefix, "bias", bias_np, layer.bias_quantizer, nodes, initializers, use_qonnx, store_integer_weights - ) - 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" - - 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) - - qonnx_p = use_qonnx and is_pq - intstore_p = store_integer_weights and is_pq - q_gamma = _emit_param( - prefix, "gamma", gamma_np, layer.weight_quantizer if is_pq else None, nodes, initializers, qonnx_p, intstore_p - ) - q_beta = _emit_param( - prefix, "beta", beta_np, layer.bias_quantizer if is_pq else None, nodes, initializers, qonnx_p, intstore_p - ) - - 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): - 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] - - q_weight = _emit_param( - prefix, "weight", kernel_np, layer.weight_quantizer, nodes, initializers, use_qonnx, store_integer_weights, out_units - ) - - # 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) - q_bias = _emit_param( - prefix, "bias", bias_np, 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 - - current = _maybe_quant_output(layer, prefix, current, nodes, initializers, quant_fn) - return current - - -def _add_quantized_softmax(sm, prefix, current, nodes, initializers, quant_fn, kpm_mask=None): - enable = sm.enable_quantization - scaler = float(sm.input_scaler) - stable = bool(sm.stable) - eps = float(sm.epsilon) - - def qdq(q, pfx, x): - k, i, f = q.get_quantization_bits() - q_nodes, out = quant_fn(pfx, x, q.round_mode, _np(k), _np(i), _np(f), initializers, overflow_mode=q.overflow) - nodes.extend(q_nodes) - return out - - # 1) Softmax input quantizer. - if sm.quantize_input and enable: - current = qdq(sm.input_quantizer, f"{prefix}_sm_in_q", current) - - # 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 (only when quantize_input==stable), - # Exp of (-scaler * x) for the stable branch (+scaler otherwise), output QDQ. - exp_t = sm.exp_table - if exp_t.quantize_input and enable: - exp_in = qdq(exp_t.input_quantizer, f"{prefix}_sm_exp_in_q", exp_in) - coeff = -scaler if stable else scaler - exp_arg = exp_in - if coeff != 1.0: - coeff_name = f"{prefix}_sm_exp_coeff" - initializers.append(onh.from_array(np.array(coeff, dtype=np.float32), name=coeff_name)) - exp_arg = f"{prefix}_sm_exp_arg" - nodes.append(oh.make_node("Mul", inputs=[exp_in, coeff_name], outputs=[exp_arg])) - exp_inp = f"{prefix}_sm_exp" - nodes.append(oh.make_node("Exp", inputs=[exp_arg], outputs=[exp_inp])) - if exp_t.quantize_output and enable: - exp_inp = qdq(exp_t.output_quantizer, f"{prefix}_sm_exp_out_q", exp_inp) - - # 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=TensorProto.FLOAT)) - masked = f"{prefix}_sm_masked" - nodes.append(oh.make_node("Mul", inputs=[kpm_f, exp_inp], outputs=[masked])) - exp_inp = masked - - # 4) Sum over the last axis (ReduceSum takes axes as an input from opset 13). - sum_axes = f"{prefix}_sm_sum_axes" - initializers.append(onh.from_array(np.array([-1], dtype=np.int64), name=sum_axes)) - sums = f"{prefix}_sm_sum" - nodes.append(oh.make_node("ReduceSum", inputs=[exp_inp, sum_axes], outputs=[sums], keepdims=1)) - - # 5) Quantized reciprocal table: input QDQ, 1/(x+eps), output QDQ. - inv_t = sm.inv_table - inv_in = sums - if inv_t.quantize_input and enable: - inv_in = qdq(inv_t.input_quantizer, f"{prefix}_sm_inv_in_q", inv_in) - eps_name = f"{prefix}_sm_eps" - initializers.append(onh.from_array(np.array(eps, dtype=np.float32), name=eps_name)) - 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_t.quantize_output and enable: - divisor = qdq(inv_t.output_quantizer, f"{prefix}_sm_inv_out_q", divisor) - - # 6) Multiply numerator by reciprocal. - out = f"{prefix}_sm_out" - nodes.append(oh.make_node("Mul", inputs=[exp_inp, divisor], outputs=[out])) - current = out - - # 7) Softmax output quantizer. - if sm.quantize_output and enable: - current = qdq(sm.output_quantizer, f"{prefix}_sm_out_q", current) - 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, -): - 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 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 = f"{prefix}_kpm_axes" - initializers.append(onh.from_array(np.array([1, 2], dtype=np.int64), name=kpm_axes)) - 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])) - - current = _add_quantized_softmax( - layer.softmax, f"{prefix}_attn", current, nodes, initializers, quant_fn, kpm_mask=kpm_mult - ) - attn_w_name = current # softmax output = attention weights (also averaged over heads below) - - 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) - - 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: - 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"): - 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 - - act = layer.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])) - 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 _resolve_mask_arg(mask, prefix, kind, tensor_to_onnx, initializers): - """Resolve an MHA mask call-argument to an ONNX value name (or None). - - A KerasTensor mask (e.g. a runtime keras.Input) maps through tensor_to_onnx; a - constant array mask (e.g. a fixed causal mask) becomes an initializer. - """ - 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)] - arr = np.asarray(_np(mask)) - name = f"{prefix}_{kind}_const" - initializers.append(onh.from_array(arr, name=name)) - return name - - -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.""" - - # --- PQuant layers --- - if isinstance(layer, PQMultiheadAttention): - # input_onnx_names = [query, key, value] (+ any mask tensors, ignored here) - # or [single_input] for self-attention. q/k/v are always the first three. - 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] - # Masks are passed as call kwargs on the layer's inbound node. - kwargs = layer._inbound_nodes[0].arguments.kwargs if layer._inbound_nodes 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_integer_weights, - key_padding_mask=kpm, - attn_mask=attn_mask, - ) - - 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" - 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): - 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): - 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, - tensor_to_onnx=tensor_to_onnx, - ) - - _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 - - n_in = len(model.inputs) - if n_in == 1: - input_names = ["input"] - input_shapes = [tuple(input_shape)] - else: - input_names = [f"input_{i}" for i in range(n_in)] - 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 n_in == 1 else dummies, 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 - 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, 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 - - _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 diff --git a/src/pquant/core/keras/layers.py b/src/pquant/core/keras/layers.py index 50d80ad..ed61973 100644 --- a/src/pquant/core/keras/layers.py +++ b/src/pquant/core/keras/layers.py @@ -1957,8 +1957,6 @@ def hgq_loss(self): @staticmethod def _split_qkv_shapes(input_shape): - # Resolve (query, key, value) shapes from either a single shape (self-attention) - # or a list/tuple of per-input shapes; missing key/value fall back to query/key. 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 @@ -1968,20 +1966,12 @@ def _split_qkv_shapes(input_shape): return q_shape, k_shape, v_shape def compute_output_shape(self, input_shape): - # Provide static output shapes so Keras does not run call() symbolically for - # shape inference (which would build the quantized/pruned sublayers inside a - # scratch graph and fail). Mirrors the (output, avg_attn_weights) tuple that - # call() returns: output is (B, Tq, embed_dim), attn weights are (B, Tq, Tk). 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): - # Build the projection/softmax sublayers explicitly. Without this Keras would - # try to auto-build the layer by tracing call() in a scratch FuncGraph, which - # fails for the quantized/pruned PQDense sublayers (their build() creates tensors - # that escape the scratch graph). Mirrors how PQDense itself defines build(). q_shape, k_shape, v_shape = self._split_qkv_shapes(input_shape) self.q_proj.build(q_shape) self.k_proj.build(k_shape) 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..0813d46 --- /dev/null +++ b/src/pquant/core/keras/onnx/convert_to_onnx.py @@ -0,0 +1,417 @@ +""" +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 +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, +) +from pquant.core.keras.onnx.helpers import ( + emit_getitem, + emit_squeeze, + emit_unsqueeze, + keras_dtype_to_tp, + qdq_node, + quant_node, + to_np, +) +from pquant.core.keras.onnx.layers import ( + add_avgpool, + add_batchnorm, + add_conv, + add_dense, + add_depthwise_conv, + add_global_avgpool, + add_mha, + add_pq_activation, +) + + +def resolve_mask_arg(mask, prefix, kind, tensor_to_onnx, 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)] + arr = np.asarray(to_np(mask)) + name = f"{prefix}_{kind}_const" + initializers.append(onh.from_array(arr, name=name)) + return name + + +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): + 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] + kwargs = layer._inbound_nodes[0].arguments.kwargs if layer._inbound_nodes 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_integer_weights, + key_padding_mask=kpm, + attn_mask=attn_mask, + ) + + 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 (weightless/structural only; weighted layers + # must be PQ variants — plain Conv/Dense/BatchNorm are not supported) --- + if type(layer).__name__ == "GetItem": + # keras.ops GetItem operation recorded by ``x[...]`` KerasTensor syntax. + node = layer._inbound_nodes[0] + args = node.arguments.args + spec = args[1] if len(args) > 1 else node.arguments.kwargs["key"] + rank = len(args[0].shape) + return emit_getitem(prefix, current, spec, rank, nodes, initializers) + + if type(layer).__name__ == "ExpandDims": + # keras.ops.expand_dims operation; the axis is stored on the op. + rank = len(layer._inbound_nodes[0].arguments.args[0].shape) + return emit_unsqueeze(prefix, current, [int(layer.axis) % (rank + 1)], nodes, initializers) + + if type(layer).__name__ == "Squeeze": + # 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 = layer._inbound_nodes[0].arguments.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) + + 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}") + + +# --------------------------------------------------------------------------- +# Keras functional model graph traversal +# --------------------------------------------------------------------------- + + +def build_tensor_onnx_map(model): + 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): + 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 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 + + current = input_onnx_names[0] + 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, + 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 + + n_in = len(model.inputs) + if n_in == 1: + input_names = ["input"] + input_shapes = [tuple(input_shape)] + else: + input_names = [f"input_{i}" for i in range(n_in)] + 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 n_in == 1 else dummies, 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:]) + + 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, 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 + + _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 diff --git a/src/pquant/core/keras/onnx/helpers.py b/src/pquant/core/keras/onnx/helpers.py new file mode 100644 index 0000000..abaf8ef --- /dev/null +++ b/src/pquant/core/keras/onnx/helpers.py @@ -0,0 +1,405 @@ +""" +Low-level ONNX node emitters and small utilities shared by the PQuant +Keras → ONNX converter. + +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). +""" + +import keras +import numpy as np +import onnx.helper as oh +import onnx.numpy_helper as onh +from onnx import TensorProto + +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 + + +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). + """ + 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 + + +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 + + +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 to_np(tensor): + 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 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 emit_param(prefix, name, arr, quantizer, nodes, initializers, use_qonnx, store_integer_weights, out_channels=None): + """Emit the ONNX value for a learnable parameter (kernel/bias/gamma/beta) and return its name""" + if use_qonnx: + fp_name = f"{prefix}_{name}_fp" + initializers.append(onh.from_array(arr, name=fp_name)) + k, i, f = quantizer.get_quantization_bits() + q_nodes, out = quant_node( + f"{prefix}_{name}", + fp_name, + quantizer.round_mode, + to_np(k), + to_np(i), + to_np(f), + initializers, + overflow_mode=quantizer.overflow, + ) + nodes.extend(q_nodes) + return out + if store_integer_weights: + k, i, f = quantizer.get_quantization_bits() + if out_channels is not None: + k_a = weight_f_for_onnx(to_np(k), out_channels) + i_a = weight_f_for_onnx(to_np(i), out_channels) + f_a = weight_f_for_onnx(to_np(f), out_channels) + else: + k_a, i_a, f_a = to_np(k), to_np(i), to_np(f) + q_nodes, out = int_weight_node(f"{prefix}_{name}", arr, k_a, i_a, f_a, initializers) + nodes.extend(q_nodes) + return out + out = f"{prefix}_{name}" + initializers.append(onh.from_array(arr, name=out)) + return out + + +def maybe_quant_input(layer, prefix, current, nodes, initializers, quant_fn): + if getattr(layer, "input_quantizer", None) is not None and layer.quantize_input and layer.enable_quantization: + q = layer.input_quantizer + k, i, f = q.get_quantization_bits() + new_nodes, current = quant_fn( + f"{prefix}_in", + current, + q.round_mode, + to_np(k), + to_np(i), + to_np(f), + initializers, + overflow_mode=q.overflow, + ) + 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 layer.quantize_output and layer.enable_quantization: + q = layer.output_quantizer + k, i, f = q.get_quantization_bits() + new_nodes, current = quant_fn( + f"{prefix}_out", + current, + q.round_mode, + to_np(k), + to_np(i), + to_np(f), + initializers, + overflow_mode=q.overflow, + ) + 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", keras.config.image_data_format()) == "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) + + +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)): + name = f"{prefix}_slice_{part}" + initializers.append(onh.from_array(np.array(vals, dtype=np.int64), name=name)) + slice_inputs.append(name) + current = f"{prefix}_slice" + nodes.append(oh.make_node("Slice", inputs=slice_inputs, outputs=[current])) + if squeeze_axes: + # Squeeze takes axes as an input tensor from opset 13 on (the converter minimum). + ax_name = f"{prefix}_squeeze_axes" + initializers.append(onh.from_array(np.array(squeeze_axes, dtype=np.int64), name=ax_name)) + out = f"{prefix}_squeeze" + nodes.append(oh.make_node("Squeeze", inputs=[current, ax_name], outputs=[out])) + current = out + 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).""" + if not axes: + return input_name + ax_name = f"{prefix}_squeeze_axes" + initializers.append(onh.from_array(np.array(sorted(axes), dtype=np.int64), name=ax_name)) + 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 = f"{prefix}_unsqueeze_axes" + initializers.append(onh.from_array(np.array(axes, dtype=np.int64), name=ax_name)) + out = f"{prefix}_unsqueeze" + nodes.append(oh.make_node("Unsqueeze", inputs=[input_name, ax_name], outputs=[out])) + return out diff --git a/src/pquant/core/keras/onnx/layers.py b/src/pquant/core/keras/onnx/layers.py new file mode 100644 index 0000000..1cc40a2 --- /dev/null +++ b/src/pquant/core/keras/onnx/layers.py @@ -0,0 +1,579 @@ +"""Per-layer ONNX graph builders (Dense/Conv/BN/Pool/Softmax/MHA) for the PQuant Keras converter.""" + +import numpy as np +import onnx.helper as oh +import onnx.numpy_helper as onh +from onnx import TensorProto + +from pquant.core.keras.layers import PQBatchNormalization +from pquant.core.keras.onnx.helpers import ( + add_transpose, + bn_transpose_info, + channels_last, + 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 # [out, in] + out_units = kernel_np.shape[0] + + q_weight = emit_param( + prefix, "weight", kernel_np, layer.weight_quantizer, nodes, initializers, use_qonnx, store_integer_weights, out_units + ) + + gemm_inputs = [current, q_weight] + + if layer._bias is not None: + bias_np = to_np(layer._bias) + q_bias = emit_param( + prefix, "bias", bias_np, layer.bias_quantizer, nodes, initializers, use_qonnx, store_integer_weights + ) + 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): + 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 = 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] + + out_channels = kernel_onnx.shape[0] + + q_weight = emit_param( + prefix, + "weight", + kernel_onnx, + layer.weight_quantizer, + nodes, + initializers, + use_qonnx, + store_integer_weights, + out_channels, + ) + + conv_inputs = [current, q_weight] + + if layer._bias is not None: + bias_np = to_np(layer._bias) + q_bias = emit_param( + prefix, "bias", bias_np, layer.bias_quantizer, nodes, initializers, use_qonnx, store_integer_weights + ) + 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] * ndim + pads = p + p # ONNX format: [begin_0, begin_1, ..., end_0, end_1, ...] + auto_pad = "NOTSET" + + 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 = to_np(layer._kernel) # [kH, kW, in, depth_mult] + in_ch, depth_mult = kernel_np.shape[2], kernel_np.shape[3] + 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] + + q_weight = emit_param( + prefix, + "weight", + kernel_onnx, + layer.weight_quantizer, + nodes, + initializers, + use_qonnx, + store_integer_weights, + out_channels, + ) + + conv_inputs = [current, q_weight] + + if layer._bias is not None: + bias_np = to_np(layer._bias) + q_bias = emit_param( + prefix, "bias", bias_np, layer.bias_quantizer, nodes, initializers, use_qonnx, store_integer_weights + ) + 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" + + 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 = to_np(layer.gamma) if layer.gamma is not None else None + beta_np = to_np(layer.beta) if layer.beta is not None else None + + if gamma_np is None: + # scale=False: use ones + n_ch = to_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 = to_np(layer.moving_mean).shape[0] + beta_np = np.zeros(n_ch, dtype=np.float32) + + qonnx_p = use_qonnx and is_pq + intstore_p = store_integer_weights and is_pq + q_gamma = emit_param( + prefix, "gamma", gamma_np, layer.weight_quantizer if is_pq else None, nodes, initializers, qonnx_p, intstore_p + ) + q_beta = emit_param( + prefix, "beta", beta_np, layer.bias_quantizer if is_pq else None, nodes, initializers, qonnx_p, intstore_p + ) + + mean_name = f"{prefix}_running_mean" + var_name = f"{prefix}_running_var" + initializers.append(onh.from_array(to_np(layer.moving_mean), name=mean_name)) + initializers.append(onh.from_array(to_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): + current = maybe_quant_input(layer, prefix, current, nodes, initializers, quant_fn) + + kernel_np = to_np(layer._kernel).T # [out, in] + out_units = kernel_np.shape[0] + + q_weight = emit_param( + prefix, "weight", kernel_np, layer.weight_quantizer, nodes, initializers, use_qonnx, store_integer_weights, out_units + ) + + # 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 = to_np(layer._bias) + q_bias = emit_param( + prefix, "bias", bias_np, 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 + + current = maybe_quant_output(layer, prefix, current, nodes, initializers, quant_fn) + return current + + +def add_quantized_softmax(sm, prefix, current, nodes, initializers, quant_fn, kpm_mask=None): + enable = sm.enable_quantization + scaler = float(sm.input_scaler) + stable = bool(sm.stable) + eps = float(sm.epsilon) + + def qdq(q, pfx, x): + k, i, f = q.get_quantization_bits() + q_nodes, out = quant_fn(pfx, x, q.round_mode, to_np(k), to_np(i), to_np(f), initializers, overflow_mode=q.overflow) + nodes.extend(q_nodes) + return out + + # 1) Softmax input quantizer. + if sm.quantize_input and enable: + current = qdq(sm.input_quantizer, f"{prefix}_sm_in_q", current) + + # 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 (only when quantize_input==stable), + # Exp of (-scaler * x) for the stable branch (+scaler otherwise), output QDQ. + exp_t = sm.exp_table + if exp_t.quantize_input and enable: + exp_in = qdq(exp_t.input_quantizer, f"{prefix}_sm_exp_in_q", exp_in) + coeff = -scaler if stable else scaler + exp_arg = exp_in + if coeff != 1.0: + coeff_name = f"{prefix}_sm_exp_coeff" + initializers.append(onh.from_array(np.array(coeff, dtype=np.float32), name=coeff_name)) + exp_arg = f"{prefix}_sm_exp_arg" + nodes.append(oh.make_node("Mul", inputs=[exp_in, coeff_name], outputs=[exp_arg])) + exp_inp = f"{prefix}_sm_exp" + nodes.append(oh.make_node("Exp", inputs=[exp_arg], outputs=[exp_inp])) + if exp_t.quantize_output and enable: + exp_inp = qdq(exp_t.output_quantizer, f"{prefix}_sm_exp_out_q", exp_inp) + + # 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=TensorProto.FLOAT)) + masked = f"{prefix}_sm_masked" + nodes.append(oh.make_node("Mul", inputs=[kpm_f, exp_inp], outputs=[masked])) + exp_inp = masked + + # 4) Sum over the last axis (ReduceSum takes axes as an input from opset 13). + sum_axes = f"{prefix}_sm_sum_axes" + initializers.append(onh.from_array(np.array([-1], dtype=np.int64), name=sum_axes)) + sums = f"{prefix}_sm_sum" + nodes.append(oh.make_node("ReduceSum", inputs=[exp_inp, sum_axes], outputs=[sums], keepdims=1)) + + # 5) Quantized reciprocal table: input QDQ, 1/(x+eps), output QDQ. + inv_t = sm.inv_table + inv_in = sums + if inv_t.quantize_input and enable: + inv_in = qdq(inv_t.input_quantizer, f"{prefix}_sm_inv_in_q", inv_in) + eps_name = f"{prefix}_sm_eps" + initializers.append(onh.from_array(np.array(eps, dtype=np.float32), name=eps_name)) + 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_t.quantize_output and enable: + divisor = qdq(inv_t.output_quantizer, f"{prefix}_sm_inv_out_q", divisor) + + # 6) Multiply numerator by reciprocal. + out = f"{prefix}_sm_out" + nodes.append(oh.make_node("Mul", inputs=[exp_inp, divisor], outputs=[out])) + current = out + + # 7) Softmax output quantizer. + if sm.quantize_output and enable: + current = qdq(sm.output_quantizer, f"{prefix}_sm_out_q", current) + 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, +): + 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 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 = f"{prefix}_kpm_axes" + initializers.append(onh.from_array(np.array([1, 2], dtype=np.int64), name=kpm_axes)) + 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])) + + current = add_quantized_softmax( + layer.softmax, f"{prefix}_attn", current, nodes, initializers, quant_fn, kpm_mask=kpm_mult + ) + attn_w_name = current # softmax output = attention weights (also averaged over heads below) + + 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) + + 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: + 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"): + 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 + + act = layer.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])) + 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 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 69a0cc1..0000000 --- a/src/pquant/core/torch/convert_to_onnx.py +++ /dev/null @@ -1,1701 +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"): - k_val = int(k.item()) - if 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 - 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()) - 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 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 _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 _maybe_quant_input(module, prefix, current, nodes, initializers, quant_fn): - # input_quantizer is created conditionally, so guard it; the bool flags are always present. - if getattr(module, "input_quantizer", None) is not None and module.quantize_input and module.enable_quantization: - 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=q.overflow) - 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 module.quantize_output and module.enable_quantization: - 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=q.overflow - ) - nodes.extend(new_nodes) - return current - - -def _emit_param(prefix, name, arr, quantizer, nodes, initializers, use_qonnx, store_integer_weights): - if use_qonnx: - fp_name = f"{prefix}_{name}_fp" - initializers.append(onh.from_array(arr, name=fp_name)) - 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, i, f = quantizer.get_quantization_bits() - q_nodes, out = _int_weight_node(f"{prefix}_{name}", arr, k, i, f, initializers) - nodes.extend(q_nodes) - return out - out = f"{prefix}_{name}" - initializers.append(onh.from_array(arr, name=out)) - return out - - -# --------------------------------------------------------------------------- -# per-layer graph builders -# --------------------------------------------------------------------------- - - -def _add_dense_integer(module, prefix, current, nodes, initializers): - 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") - - # --- 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()) # get_quantization_bits() always returns tensors - 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 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): - 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 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" - 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) - q_bias = _emit_param( - prefix, "bias", bias_np, 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 - - 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) - q_weight = _emit_param( - prefix, "weight", weight_np, module.weight_quantizer, nodes, initializers, use_qonnx, store_integer_weights - ) - - gemm_inputs = [current, q_weight] - - if module._bias is not None: - bias_np = module._bias.detach().cpu().numpy().astype(np.float32) - q_bias = _emit_param( - prefix, "bias", bias_np, module.bias_quantizer, nodes, initializers, use_qonnx, store_integer_weights - ) - 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) - q_weight = _emit_param( - prefix, "weight", weight_np, module.weight_quantizer, nodes, initializers, use_qonnx, store_integer_weights - ) - - conv_inputs = [current, q_weight] - - if module._bias is not None: - bias_np = module._bias.detach().cpu().numpy().astype(np.float32) - q_bias = _emit_param( - prefix, "bias", bias_np, module.bias_quantizer, nodes, initializers, use_qonnx, store_integer_weights - ) - 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) - - 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) - - q_gamma = _emit_param( - prefix, "gamma", gamma_np, module.weight_quantizer, nodes, initializers, use_qonnx, store_integer_weights - ) - q_beta = _emit_param( - prefix, "beta", beta_np, module.bias_quantizer, nodes, initializers, use_qonnx, store_integer_weights - ) - - 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): - 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 - - qonnx_p = use_qonnx and has_weight - intstore_p = store_integer_weights and has_weight - q_gamma = _emit_param( - prefix, - "gamma", - gamma_np, - module.weight_quantizer if has_weight else None, - nodes, - initializers, - qonnx_p, - intstore_p, - ) - if has_bias: - q_beta = _emit_param( - prefix, - "beta", - beta_np, - module.bias_quantizer if has_weight else None, - nodes, - initializers, - qonnx_p, - intstore_p, - ) - - 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) - - 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_quantized_softmax(sm, prefix, current, nodes, initializers, quant_fn, kpm_mask=None): - enable = sm.enable_quantization - scaler = float(sm.input_scaler) - stable = bool(sm.stable) - eps = float(sm.epsilon) - - def qdq(q, pfx, x): - k, i, f = q.get_quantization_bits() - q_nodes, out = quant_fn(pfx, x, q.round_mode, k, i, f, initializers, overflow_mode=q.overflow) - nodes.extend(q_nodes) - return out - - if sm.quantize_input and enable: - current = qdq(sm.input_quantizer, f"{prefix}_sm_in_q", current) - - 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 - - exp_t = sm.exp_table - if exp_t.quantize_input and enable: - exp_in = qdq(exp_t.input_quantizer, f"{prefix}_sm_exp_in_q", exp_in) - coeff = -scaler if stable else scaler - exp_arg = exp_in - if coeff != 1.0: - coeff_name = f"{prefix}_sm_exp_coeff" - initializers.append(onh.from_array(np.array(coeff, dtype=np.float32), name=coeff_name)) - exp_arg = f"{prefix}_sm_exp_arg" - nodes.append(oh.make_node("Mul", inputs=[exp_in, coeff_name], outputs=[exp_arg])) - exp_inp = f"{prefix}_sm_exp" - nodes.append(oh.make_node("Exp", inputs=[exp_arg], outputs=[exp_inp])) - if exp_t.quantize_output and enable: - exp_inp = qdq(exp_t.output_quantizer, f"{prefix}_sm_exp_out_q", exp_inp) - - 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=TensorProto.FLOAT)) - masked = f"{prefix}_sm_masked" - nodes.append(oh.make_node("Mul", inputs=[kpm_f, exp_inp], outputs=[masked])) - exp_inp = masked - - sum_axes = f"{prefix}_sm_sum_axes" - initializers.append(onh.from_array(np.array([-1], dtype=np.int64), name=sum_axes)) - sums = f"{prefix}_sm_sum" - nodes.append(oh.make_node("ReduceSum", inputs=[exp_inp, sum_axes], outputs=[sums], keepdims=1)) - - inv_t = sm.inv_table - inv_in = sums - if inv_t.quantize_input and enable: - inv_in = qdq(inv_t.input_quantizer, f"{prefix}_sm_inv_in_q", inv_in) - eps_name = f"{prefix}_sm_eps" - initializers.append(onh.from_array(np.array(eps, dtype=np.float32), name=eps_name)) - 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_t.quantize_output and enable: - divisor = qdq(inv_t.output_quantizer, f"{prefix}_sm_inv_out_q", divisor) - - out = f"{prefix}_sm_out" - nodes.append(oh.make_node("Mul", inputs=[exp_inp, divisor], outputs=[out])) - current = out - - if sm.quantize_output and enable: - current = qdq(sm.output_quantizer, f"{prefix}_sm_out_q", current) - return current - - -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, -): - 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 - - 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 = f"{prefix}_kpm_axes" - initializers.append(onh.from_array(np.array([1, 2], dtype=np.int64), name=kpm_axes)) - 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])) - - current = _add_quantized_softmax( - module.softmax, f"{prefix}_attn", current, nodes, initializers, quant_fn, kpm_mask=kpm_mult - ) - attn_w_name = current # softmax output = attention weights (also averaged over heads below) - - 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" # 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=module.overflow) - 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. - - Note: - This is a thin wrapper over convert_to_onnx_fx(). An ``nn.Sequential`` is a - plain linear chain, so torch.fx always traces it successfully; routing through - the FX converter keeps a single code path for both linear and branched models. - """ - return convert_to_onnx_fx( - model, - input_shape, - output_path=output_path, - opset=opset, - use_qonnx=use_qonnx, - store_integer_weights=store_integer_weights, - integer_ops=integer_ops, - include_clip=include_clip, - batch_size=batch_size, - ) - - -# --------------------------------------------------------------------------- -# 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: - # ----- 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): - _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 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, - concrete_args: dict | None = None, - input_dtypes=None, - batch_size: int | None = None, -) -> 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). - - 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: - 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. - - Remaining args match the per-layer quantization behaviour described in the module docstring. - """ - 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)) - - graph = _PQTracer().trace(model, concrete_args=concrete_args) - gm = _fx.GraphModule(model, graph) - - 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() - - tensor_phs = list(gm.graph.find_nodes(op="placeholder")) - if len(tensor_phs) != len(input_shapes): - raise ValueError( - f"FX export: model.forward has {len(tensor_phs)} tensor input(s) but " - f"input_shape describes {len(input_shapes)}. Specialize non-tensor " - f"arguments via concrete_args={{...}}." - ) - input_names = ["input"] if len(tensor_phs) == 1 else [str(p.target) for p in tensor_phs] - ph_to_name = {p: n for p, n in zip(tensor_phs, input_names)} - - from torch.fx.passes.shape_prop import ShapeProp - - 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) - - 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]: - 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] = ph_to_name[node] - - elif node.op == "get_attr": - 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): - # forward(query, key, value, key_padding_mask=None, attn_mask=None, ...) - 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 - - def _mask_name(pos, kw, node=node): - 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 node_to_name[arg] - - kpm_name = _mask_name(3, "key_padding_mask") - attn_mask_name = _mask_name(4, "attn_mask") - 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, - key_padding_mask=kpm_name, - attn_mask=attn_mask_name, - ) - # 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 torch._assert or getattr(fn, "__name__", "") == "_assert" or fn is _operator.eq: - continue - - 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) - - graph_input_names = set(input_names) - for idx, nm in enumerate(output_names): - if nm in graph_input_names: - ident = f"{nm}_identity_out{idx}" - onnx_nodes.append(oh.make_node("Identity", inputs=[nm], outputs=[ident])) - output_names[idx] = ident - - 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) - ] - - onnx_graph = oh.make_graph( - nodes=onnx_nodes, - name="pquant_onnx_fx", - inputs=input_vis, - 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 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..77fe135 --- /dev/null +++ b/src/pquant/core/torch/onnx/convert_to_onnx.py @@ -0,0 +1,854 @@ +""" +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 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.onnx.helpers import ( # noqa: E402 + emit_getitem, + emit_squeeze, + emit_unsqueeze, + maybe_quant_input, + maybe_quant_output, + qdq_node, + quant_node, +) +from pquant.core.torch.onnx.layers import ( # noqa: E402 + add_avgpool, + add_batchnorm, + add_conv, + add_dense, + add_layernorm, + add_mha, +) +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): + """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.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): + out, _ = add_mha( + module, prefix, current, current, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights + ) + return out + if isinstance(module, Quantizer): + k, i, f = module.get_quantization_bits() + new_nodes, out = quant_fn(prefix, current, module.round_mode, k, i, f, initializers, overflow_mode=module.overflow) + nodes.extend(new_nodes) + return out + raise TypeError(f"Unsupported module type for ONNX export: {type(module).__name__}") + + +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: + # ----- 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 + + +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 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)) + + graph = PQTracer().trace(model, concrete_args=concrete_args) + gm = fx.GraphModule(model, graph) + + 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() + + tensor_phs = list(gm.graph.find_nodes(op="placeholder")) + if len(tensor_phs) != len(input_shapes): + raise ValueError( + f"FX export: model.forward has {len(tensor_phs)} tensor input(s) but " + f"input_shape describes {len(input_shapes)}. Specialize non-tensor " + f"arguments via concrete_args={{...}}." + ) + input_names = ["input"] if len(tensor_phs) == 1 else [str(p.target) for p in tensor_phs] + ph_to_name = {p: n for p, n in zip(tensor_phs, input_names)} + + from torch.fx.passes.shape_prop import ShapeProp + + 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) + + 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]: + 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 node_shape(n: fx.Node) -> tuple: + 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 tuple(meta.shape) + + def node_rank(n: fx.Node) -> int: + return len(node_shape(n)) + + def squeeze_axes_for(node: fx.Node) -> list[int]: + """Resolve the axes a torch squeeze()/​.squeeze() call removes.""" + in_shape = 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 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] = ph_to_name[node] + + elif node.op == "get_attr": + 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): + # forward(query, key, value, key_padding_mask=None, attn_mask=None, ...) + 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 + + def mask_name(pos, kw, node=node): + 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 node_to_name[arg] + + kpm_name = mask_name(3, "key_padding_mask") + attn_mask_name = mask_name(4, "attn_mask") + 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, + key_padding_mask=kpm_name, + attn_mask=attn_mask_name, + ) + 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 torch._assert or getattr(fn, "__name__", "") == "_assert" or fn is _operator.eq: + continue + + if fn is _operator.getitem: + container = node_to_name[node.args[0]] + if isinstance(container, tuple): + # Unpack a tuple output (e.g. from PQMultiheadAttention). + node_to_name[node] = container[node.args[1]] + else: + # Tensor slicing: x[:, 0], x[..., :4], ... → Slice (+ Squeeze) + rank = node_rank(node.args[0]) + node_to_name[node] = emit_getitem(node.name, container, node.args[1], rank, onnx_nodes, initializers) + 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 = node_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 = node_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 + + elif fn is torch.squeeze: + node_to_name[node] = emit_squeeze( + node.name, res(node.args[0]), squeeze_axes_for(node), onnx_nodes, initializers + ) + + elif fn is torch.unsqueeze: + dim = int(node.args[1]) if len(node.args) > 1 else int(node.kwargs["dim"]) + axes = [dim % (node_rank(node.args[0]) + 1)] + node_to_name[node] = emit_unsqueeze(node.name, res(node.args[0]), axes, onnx_nodes, initializers) + + 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 = node_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 = node_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 + + elif node.target == "squeeze": + node_to_name[node] = emit_squeeze(node.name, x, squeeze_axes_for(node), onnx_nodes, initializers) + + elif node.target == "unsqueeze": + dim = int(node.args[1]) if len(node.args) > 1 else int(node.kwargs["dim"]) + axes = [dim % (node_rank(node.args[0]) + 1)] + node_to_name[node] = emit_unsqueeze(node.name, x, axes, onnx_nodes, initializers) + + 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) + + graph_input_names = set(input_names) + for idx, nm in enumerate(output_names): + if nm in graph_input_names: + ident = f"{nm}_identity_out{idx}" + onnx_nodes.append(oh.make_node("Identity", inputs=[nm], outputs=[ident])) + output_names[idx] = ident + + 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) + ] + + onnx_graph = oh.make_graph( + nodes=onnx_nodes, + name="pquant_onnx_fx", + inputs=input_vis, + 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 diff --git a/src/pquant/core/torch/onnx/helpers.py b/src/pquant/core/torch/onnx/helpers.py new file mode 100644 index 0000000..37510e0 --- /dev/null +++ b/src/pquant/core/torch/onnx/helpers.py @@ -0,0 +1,304 @@ +""" +Low-level ONNX node emitters and small utilities shared by the PQuant +torch → ONNX converter. + +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). +""" + +import numpy as np +import onnx.helper as oh +import onnx.numpy_helper as onh + +# --------------------------------------------------------------------------- +# 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"): + k_val = int(k.item()) + if 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 + + +def qdq_node( + name_prefix, input_name, rounding_mode, k, i, f, initializers, overflow_mode="SAT", include_clip=True +): # noqa: ARG001 + 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 + + +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()) + 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 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 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 maybe_quant_input(module, prefix, current, nodes, initializers, quant_fn): + # input_quantizer is created conditionally, so guard it; the bool flags are always present. + if getattr(module, "input_quantizer", None) is not None and module.quantize_input and module.enable_quantization: + 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=q.overflow) + 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 module.quantize_output and module.enable_quantization: + 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=q.overflow + ) + nodes.extend(new_nodes) + return current + + +def emit_param(prefix, name, arr, quantizer, nodes, initializers, use_qonnx, store_integer_weights): + if use_qonnx: + fp_name = f"{prefix}_{name}_fp" + initializers.append(onh.from_array(arr, name=fp_name)) + 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, i, f = quantizer.get_quantization_bits() + q_nodes, out = int_weight_node(f"{prefix}_{name}", arr, k, i, f, initializers) + nodes.extend(q_nodes) + return out + out = f"{prefix}_{name}" + initializers.append(onh.from_array(arr, name=out)) + return out + + +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)): + name = f"{prefix}_slice_{part}" + initializers.append(onh.from_array(np.array(vals, dtype=np.int64), name=name)) + slice_inputs.append(name) + current = f"{prefix}_slice" + nodes.append(oh.make_node("Slice", inputs=slice_inputs, outputs=[current])) + if squeeze_axes: + # Squeeze takes axes as an input tensor from opset 13 on (the converter minimum). + ax_name = f"{prefix}_squeeze_axes" + initializers.append(onh.from_array(np.array(squeeze_axes, dtype=np.int64), name=ax_name)) + out = f"{prefix}_squeeze" + nodes.append(oh.make_node("Squeeze", inputs=[current, ax_name], outputs=[out])) + current = out + 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).""" + if not axes: + return input_name + ax_name = f"{prefix}_squeeze_axes" + initializers.append(onh.from_array(np.array(sorted(axes), dtype=np.int64), name=ax_name)) + 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 = f"{prefix}_unsqueeze_axes" + initializers.append(onh.from_array(np.array(axes, dtype=np.int64), name=ax_name)) + out = f"{prefix}_unsqueeze" + nodes.append(oh.make_node("Unsqueeze", inputs=[input_name, ax_name], outputs=[out])) + return out diff --git a/src/pquant/core/torch/onnx/layers.py b/src/pquant/core/torch/onnx/layers.py new file mode 100644 index 0000000..b7117bc --- /dev/null +++ b/src/pquant/core/torch/onnx/layers.py @@ -0,0 +1,580 @@ +"""Per-layer ONNX graph builders (Dense/Conv/BN/LN/AvgPool/Softmax/MHA) for the PQuant torch converter.""" + +import numpy as np +import onnx.helper as oh +import onnx.numpy_helper as onh +from onnx import TensorProto + +from pquant.core.torch.onnx.helpers import ( + emit_param, + maybe_quant_input, + maybe_quant_output, + qdq_node, + to_list, + torch_padding_to_onnx, +) + + +def add_dense_integer(module, prefix, current, nodes, initializers): + 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") + + # --- 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()) # get_quantization_bits() always returns tensors + 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 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], + ) + ) + + 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): + 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 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" + 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) + q_bias = emit_param( + prefix, "bias", bias_np, 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 + + 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) + q_weight = emit_param( + prefix, "weight", weight_np, module.weight_quantizer, nodes, initializers, use_qonnx, store_integer_weights + ) + + gemm_inputs = [current, q_weight] + + if module._bias is not None: + bias_np = module._bias.detach().cpu().numpy().astype(np.float32) + q_bias = emit_param( + prefix, "bias", bias_np, module.bias_quantizer, nodes, initializers, use_qonnx, store_integer_weights + ) + 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) + q_weight = emit_param( + prefix, "weight", weight_np, module.weight_quantizer, nodes, initializers, use_qonnx, store_integer_weights + ) + + conv_inputs = [current, q_weight] + + if module._bias is not None: + bias_np = module._bias.detach().cpu().numpy().astype(np.float32) + q_bias = emit_param( + prefix, "bias", bias_np, module.bias_quantizer, nodes, initializers, use_qonnx, store_integer_weights + ) + 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) + + 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) + + q_gamma = emit_param( + prefix, "gamma", gamma_np, module.weight_quantizer, nodes, initializers, use_qonnx, store_integer_weights + ) + q_beta = emit_param( + prefix, "beta", beta_np, module.bias_quantizer, nodes, initializers, use_qonnx, store_integer_weights + ) + + 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): + 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 + + qonnx_p = use_qonnx and has_weight + intstore_p = store_integer_weights and has_weight + q_gamma = emit_param( + prefix, + "gamma", + gamma_np, + module.weight_quantizer if has_weight else None, + nodes, + initializers, + qonnx_p, + intstore_p, + ) + if has_bias: + q_beta = emit_param( + prefix, + "beta", + beta_np, + module.bias_quantizer if has_weight else None, + nodes, + initializers, + qonnx_p, + intstore_p, + ) + + 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) + + 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 + + +def add_quantized_softmax(sm, prefix, current, nodes, initializers, quant_fn, kpm_mask=None): + enable = sm.enable_quantization + scaler = float(sm.input_scaler) + stable = bool(sm.stable) + eps = float(sm.epsilon) + + def qdq(q, pfx, x): + k, i, f = q.get_quantization_bits() + q_nodes, out = quant_fn(pfx, x, q.round_mode, k, i, f, initializers, overflow_mode=q.overflow) + nodes.extend(q_nodes) + return out + + if sm.quantize_input and enable: + current = qdq(sm.input_quantizer, f"{prefix}_sm_in_q", current) + + 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 + + exp_t = sm.exp_table + if exp_t.quantize_input and enable: + exp_in = qdq(exp_t.input_quantizer, f"{prefix}_sm_exp_in_q", exp_in) + coeff = -scaler if stable else scaler + exp_arg = exp_in + if coeff != 1.0: + coeff_name = f"{prefix}_sm_exp_coeff" + initializers.append(onh.from_array(np.array(coeff, dtype=np.float32), name=coeff_name)) + exp_arg = f"{prefix}_sm_exp_arg" + nodes.append(oh.make_node("Mul", inputs=[exp_in, coeff_name], outputs=[exp_arg])) + exp_inp = f"{prefix}_sm_exp" + nodes.append(oh.make_node("Exp", inputs=[exp_arg], outputs=[exp_inp])) + if exp_t.quantize_output and enable: + exp_inp = qdq(exp_t.output_quantizer, f"{prefix}_sm_exp_out_q", exp_inp) + + 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=TensorProto.FLOAT)) + masked = f"{prefix}_sm_masked" + nodes.append(oh.make_node("Mul", inputs=[kpm_f, exp_inp], outputs=[masked])) + exp_inp = masked + + sum_axes = f"{prefix}_sm_sum_axes" + initializers.append(onh.from_array(np.array([-1], dtype=np.int64), name=sum_axes)) + sums = f"{prefix}_sm_sum" + nodes.append(oh.make_node("ReduceSum", inputs=[exp_inp, sum_axes], outputs=[sums], keepdims=1)) + + inv_t = sm.inv_table + inv_in = sums + if inv_t.quantize_input and enable: + inv_in = qdq(inv_t.input_quantizer, f"{prefix}_sm_inv_in_q", inv_in) + eps_name = f"{prefix}_sm_eps" + initializers.append(onh.from_array(np.array(eps, dtype=np.float32), name=eps_name)) + 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_t.quantize_output and enable: + divisor = qdq(inv_t.output_quantizer, f"{prefix}_sm_inv_out_q", divisor) + + out = f"{prefix}_sm_out" + nodes.append(oh.make_node("Mul", inputs=[exp_inp, divisor], outputs=[out])) + current = out + + if sm.quantize_output and enable: + current = qdq(sm.output_quantizer, f"{prefix}_sm_out_q", current) + return current + + +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, +): + 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 + + 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 = f"{prefix}_kpm_axes" + initializers.append(onh.from_array(np.array([1, 2], dtype=np.int64), name=kpm_axes)) + 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])) + + current = add_quantized_softmax( + module.softmax, f"{prefix}_attn", current, nodes, initializers, quant_fn, kpm_mask=kpm_mult + ) + attn_w_name = current # softmax output = attention weights (also averaged over heads below) + + 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" # 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 diff --git a/tests/test_keras_onnx_converter.py b/tests/test_keras_onnx_converter.py index b7e840e..6cb3c0d 100644 --- a/tests/test_keras_onnx_converter.py +++ b/tests/test_keras_onnx_converter.py @@ -10,10 +10,10 @@ 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, @@ -24,13 +24,9 @@ PQMultiheadAttention, apply_final_compression, ) - -ort = pytest.importorskip("onnxruntime", reason="onnxruntime not installed") +from pquant.core.keras.onnx import convert_to_onnx ATOL = 1e-4 -# When quantization is enabled, torch/keras fake-quant and ONNX QuantizeLinear can round -# a few values to opposite sides of a 0.5 boundary (op/accumulation-order ULP differences), -# so allow ~1 quantization level of slack for graphs that re-quantize intermediates. QUANT_ATOL = 5e-3 @@ -40,9 +36,6 @@ def _atol(cfg): @pytest.fixture(params=[False, True], ids=["float", "quant"]) def cfg(request): - # Run every cfg-based test twice: float path and quantization-enabled, so the - # emitted Quantize/DequantizeLinear nodes are actually exercised against onnxruntime - # (they are skipped entirely when enable_quantization is False). c = pquant.cs_config() c.quantization_parameters.enable_quantization = request.param return c @@ -66,113 +59,49 @@ def _onnx_run(model, x: np.ndarray, input_shape: tuple, tmp_path) -> np.ndarray: return sess.run(None, {in_name: x})[0] -@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) - model = keras.Model(inputs, x) - - dummy = np.zeros((1, IN), dtype=np.float32) - model(dummy) - 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(cfg), err_msg=f"PQDense bias={bias}: keras vs ONNX mismatch") - - -@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) - - 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) - - 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(cfg), err_msg=f"PQConv2d bias={bias}: 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) - - 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) - - 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(cfg), err_msg=f"PQConv1d bias={bias}: 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 +# (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) inputs = keras.Input(shape=input_shape) - x = PQBatchNormalization(cfg, axis=bn_axis)(inputs) + x = make_layer(cfg)(inputs) model = keras.Model(inputs, x) - dummy = np.zeros((1, *input_shape), dtype=np.float32) - model(dummy, training=True) # warm up running stats + model(np.zeros((1, *input_shape), dtype=np.float32), **warmup_kwargs) 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(cfg), err_msg="PQBatchNormalization: keras vs ONNX mismatch") - - -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) - - inputs = keras.Input(shape=input_shape) - x = PQDepthwiseConv2d(cfg, kernel_size=3, padding="same")(inputs) - model = keras.Model(inputs, x) - - 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(cfg), err_msg="PQDepthwiseConv2d: keras vs ONNX mismatch") + np.testing.assert_allclose(keras_out, onnx_out, atol=_atol(cfg), err_msg="keras vs ONNX mismatch") @pytest.mark.parametrize("activation", ["relu", "tanh", "hard_tanh"]) @@ -305,3 +234,58 @@ def test_mha_key_padding_mask_onnx(cfg, tmp_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_out, 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)) + + dummy = np.zeros((1, IN), dtype=np.float32) + model(dummy) + 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(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, IN), dtype=np.float32) + model(dummy) + 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) + assert keras_out.shape == onnx_out.shape + np.testing.assert_allclose(keras_out, onnx_out, atol=_atol(cfg), err_msg="squeeze/expand_dims: keras vs ONNX mismatch") diff --git a/tests/test_torch_onnx_converter.py b/tests/test_torch_onnx_converter.py index 5edeadf..fd9a501 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,9 @@ 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.core.torch.onnx.convert_to_onnx import export_qdq_layernorm # noqa: E402 from pquant.layers import ( # noqa: E402 PQActivation, PQAvgPool1d, @@ -37,13 +35,7 @@ PQMultiheadAttention, ) -ort = pytest.importorskip("onnxruntime", reason="onnxruntime not installed") - -ATOL = 1e-4 # float32 Gemm/Conv can differ by ~1 ULP; keep some slack -# When quantization is enabled, torch fake-quant and ONNX QuantizeLinear can round a -# few values to opposite sides of a 0.5 boundary (the rounding inputs differ by float -# ULPs from differing op/accumulation order), so allow ~1 quantization level of slack -# for graphs that re-quantize intermediate activations. +ATOL = 1e-4 QUANT_ATOL = 5e-3 @@ -51,17 +43,8 @@ def _atol(cfg): return QUANT_ATOL if cfg.quantization_parameters.enable_quantization else ATOL -# --------------------------------------------------------------------------- -# fixtures -# --------------------------------------------------------------------------- - - @pytest.fixture(params=[False, True], ids=["float", "quant"]) def cfg(request): - # Run every cfg-based test twice: once with the plain float path and once with - # quantization enabled so the emitted Quantize/DequantizeLinear nodes are - # actually exercised against onnxruntime (they are skipped entirely when - # enable_quantization is False). c = pquant.cs_config() c.quantization_parameters.enable_quantization = request.param return c @@ -69,7 +52,6 @@ def cfg(request): @pytest.fixture def cfg_quant(): - # Quantization-enabled config for tests that specifically target the QDQ path. c = pquant.cs_config() c.quantization_parameters.enable_quantization = True return c @@ -93,7 +75,7 @@ def _onnx_run(model: nn.Module, x: torch.Tensor, input_shape: tuple, tmp_path) - 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] @@ -105,119 +87,70 @@ def _torch_out(model: nn.Module, x: torch.Tensor) -> np.ndarray: return model(x).cpu().numpy() -@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") - - -@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) - with torch.no_grad(): - model(x) - _apply_compression(model) - - 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") - - -@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) - with torch.no_grad(): - model(x) - _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") - - -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) - with torch.no_grad(): - model(x) - _apply_compression(model) - model.eval() # switch BN to use running stats - - 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") - - -def test_batchnorm1d_onnx(cfg, tmp_path): - C, L = 8, 16 - model = nn.Sequential( - PQBatchNorm1d(cfg, num_features=C), - nn.ReLU(), - ) - x = torch.randn(4, C, L) - with torch.no_grad(): - model(x) - _apply_compression(model) - model.eval() - - 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") - - -def test_avgpool2d_onnx(cfg, tmp_path): - C, H, W = 8, 8, 8 - model = nn.Sequential( - PQAvgPool2d(cfg, kernel_size=2, stride=2), - ) - x = torch.randn(2, C, H, W) - with torch.no_grad(): - model(x) - _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") - - -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) +# (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) + model(x) # warm-up in train mode (initialises any running stats) _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_out = _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_out, onnx_out, atol=ATOL, err_msg="torch vs ONNX mismatch") class _SelfAttnModel(nn.Module): @@ -316,7 +249,7 @@ def test_mha_key_padding_mask_onnx(cfg, bias, tmp_path): torch_out = model(x, key_padding_mask).cpu().numpy() path = str(tmp_path / "mha_kpm.onnx") - proto = convert_to_onnx_fx(model, input_shape=[(T, E), (T,)], output_path=path, input_dtypes=["float32", "bool"]) + proto = convert_to_onnx(model, input_shape=[(T, E), (T,)], output_path=path, input_dtypes=["float32", "bool"]) # 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") @@ -435,7 +368,7 @@ def test_two_input_onnx(cfg, bias, tmp_path): torch_out = model(a, b).cpu().numpy() path = str(tmp_path / "two_input.onnx") - model_proto = convert_to_onnx_fx(model, input_shape=[(IN_A,), (IN_B,)], output_path=path) + 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. @@ -461,7 +394,7 @@ def test_two_input_shape_count_mismatch(cfg, tmp_path): path = str(tmp_path / "bad_count.onnx") with pytest.raises(ValueError, match="tensor input"): - convert_to_onnx_fx(model, input_shape=(16,), output_path=path) # only one shape + convert_to_onnx(model, input_shape=(16,), output_path=path) # only one shape class _FlaggedModel(nn.Module): @@ -493,7 +426,7 @@ def test_concrete_args_specialization(cfg, scale_up, tmp_path): torch_out = model(x, scale_up).cpu().numpy() path = str(tmp_path / f"flag_{scale_up}.onnx") - model_proto = convert_to_onnx_fx(model, input_shape=(IN,), output_path=path, concrete_args={"scale_up": scale_up}) + 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. @@ -538,7 +471,7 @@ def test_residual_concat_onnx(cfg_quant, tmp_path): torch_out = model(x).cpu().numpy() path = str(tmp_path / "residual_concat.onnx") - model_proto = convert_to_onnx_fx(model, input_shape=(DIM,), output_path=path) + 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 @@ -626,7 +559,7 @@ def test_cnn_flatten_to_dense_onnx(cfg_quant, use_reshape, tmp_path): torch_out = model(x).cpu().numpy() path = str(tmp_path / f"cnn_flatten_{use_reshape}.onnx") - convert_to_onnx_fx(model, input_shape=(IN_C, HW, HW), output_path=path) + 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( @@ -661,7 +594,7 @@ def test_scalar_ops_onnx(cfg_quant, tmp_path): torch_out = model(x).cpu().numpy() path = str(tmp_path / "scalar_ops.onnx") - model_proto = convert_to_onnx_fx(model, input_shape=(DIM,), output_path=path) + model_proto = convert_to_onnx(model, input_shape=(DIM,), output_path=path) op_types = [n.op_type for n in model_proto.graph.node] for expected in ("Mul", "Sub", "Div", "Sigmoid"): assert expected in op_types, f"missing {expected} node" @@ -695,7 +628,7 @@ def test_multi_output_onnx(cfg_quant, tmp_path): t0, t1 = (t.cpu().numpy() for t in model(x)) path = str(tmp_path / "multi_output.onnx") - model_proto = convert_to_onnx_fx(model, input_shape=(DIM,), output_path=path) + model_proto = convert_to_onnx(model, input_shape=(DIM,), output_path=path) assert len(model_proto.graph.output) == 2 sess = ort.InferenceSession(path) @@ -737,9 +670,8 @@ def test_pqlayernorm_onnx(cfg_quant, tmp_path): (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), - (lambda: nn.Sequential(nn.BatchNorm2d(3)), (3, 8, 8), 2), ], - ids=["leaky_relu", "maxpool2d", "upsample", "dropout", "batchnorm2d"], + ids=["leaky_relu", "maxpool2d", "upsample", "dropout"], ) def test_plain_passthrough_layers_onnx(make_model, input_shape, batch, tmp_path): model = make_model() @@ -834,3 +766,68 @@ def test_qdq_layernorm_validation(tmp_path): # eps_q0 < 1 with pytest.raises(ValueError, match="eps_q0"): export_qdq_layernorm(path, (4, D), gamma, beta, -7, -6, eps_q0=0) + + +@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_out = _torch_out(model, x) + onnx_out = _onnx_run(model, x, input_shape=(16,), tmp_path=tmp_path) + np.testing.assert_allclose(torch_out, 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_out = _torch_out(model, x) + onnx_out = _onnx_run(model, x, input_shape=(16,), tmp_path=tmp_path) + assert torch_out.shape == onnx_out.shape + np.testing.assert_allclose(torch_out, onnx_out, atol=_atol(cfg), err_msg="squeeze/unsqueeze: torch vs ONNX mismatch") From 7b787b4204a75996581c8f5087d1a598974d90a7 Mon Sep 17 00:00:00 2001 From: Roope Niemi Date: Fri, 10 Jul 2026 16:37:11 +0200 Subject: [PATCH 4/8] rename onnx layers to layer builders to avoid confusion with core layers. cleaned more code --- src/pquant/core/keras/onnx/convert_to_onnx.py | 14 +------------- src/pquant/core/keras/onnx/helpers.py | 7 +++---- .../keras/onnx/{layers.py => layer_builders.py} | 3 ++- src/pquant/core/torch/onnx/convert_to_onnx.py | 7 +------ src/pquant/core/torch/onnx/helpers.py | 4 ---- .../torch/onnx/{layers.py => layer_builders.py} | 9 --------- 6 files changed, 7 insertions(+), 37 deletions(-) rename src/pquant/core/keras/onnx/{layers.py => layer_builders.py} (99%) rename src/pquant/core/torch/onnx/{layers.py => layer_builders.py} (97%) diff --git a/src/pquant/core/keras/onnx/convert_to_onnx.py b/src/pquant/core/keras/onnx/convert_to_onnx.py index 0813d46..6fe6df4 100644 --- a/src/pquant/core/keras/onnx/convert_to_onnx.py +++ b/src/pquant/core/keras/onnx/convert_to_onnx.py @@ -43,7 +43,7 @@ quant_node, to_np, ) -from pquant.core.keras.onnx.layers import ( +from pquant.core.keras.onnx.layer_builders import ( add_avgpool, add_batchnorm, add_conv, @@ -143,8 +143,6 @@ def emit_layer( if isinstance(layer, PQBatchNormalization): return add_batchnorm(layer, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights) - # --- Standard Keras layers (weightless/structural only; weighted layers - # must be PQ variants — plain Conv/Dense/BatchNorm are not supported) --- if type(layer).__name__ == "GetItem": # keras.ops GetItem operation recorded by ``x[...]`` KerasTensor syntax. node = layer._inbound_nodes[0] @@ -241,11 +239,6 @@ def emit_layer( raise TypeError(f"Unsupported Keras layer type for ONNX export: {type(layer).__name__!r}") -# --------------------------------------------------------------------------- -# Keras functional model graph traversal -# --------------------------------------------------------------------------- - - def build_tensor_onnx_map(model): tensor_to_onnx = {} for i, inp in enumerate(model.inputs): @@ -288,11 +281,6 @@ def register_layer_output(layer, onnx_name, tensor_to_onnx): tensor_to_onnx[id(out_tensors[0])] = onnx_name -# --------------------------------------------------------------------------- -# main conversion -# --------------------------------------------------------------------------- - - def convert_to_onnx( model: keras.Model, input_shape: tuple, diff --git a/src/pquant/core/keras/onnx/helpers.py b/src/pquant/core/keras/onnx/helpers.py index abaf8ef..0e348e0 100644 --- a/src/pquant/core/keras/onnx/helpers.py +++ b/src/pquant/core/keras/onnx/helpers.py @@ -205,13 +205,12 @@ 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. + ONNX BN always normalises on axis 1 (NCHW; true in every opset). We assume Keras uses channels_last format, + so we must insert 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 4 + 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: diff --git a/src/pquant/core/keras/onnx/layers.py b/src/pquant/core/keras/onnx/layer_builders.py similarity index 99% rename from src/pquant/core/keras/onnx/layers.py rename to src/pquant/core/keras/onnx/layer_builders.py index 1cc40a2..cf6609b 100644 --- a/src/pquant/core/keras/onnx/layers.py +++ b/src/pquant/core/keras/onnx/layer_builders.py @@ -186,7 +186,8 @@ def add_depthwise_conv(layer, prefix, current, nodes, initializers, quant_fn, us def add_batchnorm(layer, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights): - """PQBatchNormalization / standard BatchNormalization.""" + """PQBatchNormalization (also handles plain keras BatchNormalization, + but emit_layer currently only dispatches the PQ variant here).""" need_tr, perm_to_nchw, perm_to_nhwx = bn_transpose_info(layer) if need_tr: diff --git a/src/pquant/core/torch/onnx/convert_to_onnx.py b/src/pquant/core/torch/onnx/convert_to_onnx.py index 77fe135..9f683c3 100644 --- a/src/pquant/core/torch/onnx/convert_to_onnx.py +++ b/src/pquant/core/torch/onnx/convert_to_onnx.py @@ -45,7 +45,7 @@ qdq_node, quant_node, ) -from pquant.core.torch.onnx.layers import ( # noqa: E402 +from pquant.core.torch.onnx.layer_builders import ( # noqa: E402 add_avgpool, add_batchnorm, add_conv, @@ -291,7 +291,6 @@ def check_q_int16(arr: np.ndarray, frac_bits: int, name: str) -> None: 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) @@ -305,7 +304,6 @@ def check_q_int16(arr: np.ndarray, frac_bits: int, name: str) -> None: 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"), @@ -315,7 +313,6 @@ def check_q_int16(arr: np.ndarray, frac_bits: int, name: str) -> None: onh.from_array(beta.astype(np.float32), name="beta"), ] - # ----- build nodes ----- nodes = [ oh.make_node( "DequantizeLinear", @@ -345,7 +342,6 @@ def check_q_int16(arr: np.ndarray, frac_bits: int, name: str) -> None: ), ] - # ----- 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)) @@ -360,7 +356,6 @@ def check_q_int16(arr: np.ndarray, frac_bits: int, name: str) -> None: 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[:] diff --git a/src/pquant/core/torch/onnx/helpers.py b/src/pquant/core/torch/onnx/helpers.py index 37510e0..a870c7d 100644 --- a/src/pquant/core/torch/onnx/helpers.py +++ b/src/pquant/core/torch/onnx/helpers.py @@ -22,10 +22,6 @@ import onnx.helper as oh import onnx.numpy_helper as onh -# --------------------------------------------------------------------------- -# QONNX Quant node -# --------------------------------------------------------------------------- - ROUND_MODE_MAP = { "TRN": "FLOOR", "RND": "ROUND", diff --git a/src/pquant/core/torch/onnx/layers.py b/src/pquant/core/torch/onnx/layer_builders.py similarity index 97% rename from src/pquant/core/torch/onnx/layers.py rename to src/pquant/core/torch/onnx/layer_builders.py index b7117bc..774aefb 100644 --- a/src/pquant/core/torch/onnx/layers.py +++ b/src/pquant/core/torch/onnx/layer_builders.py @@ -436,7 +436,6 @@ def add_mha( 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" @@ -457,7 +456,6 @@ def add_mha( 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" @@ -497,11 +495,9 @@ def split_heads(x_name, pfx): 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" @@ -561,17 +557,12 @@ def split_heads(x_name, pfx): 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])) From a5a9f61427d8d6248f0205e6bf26e5bef88906a6 Mon Sep 17 00:00:00 2001 From: Roope Niemi Date: Wed, 15 Jul 2026 12:55:13 +0200 Subject: [PATCH 5/8] clean up more --- src/pquant/core/keras/onnx/convert_to_onnx.py | 285 +++--- src/pquant/core/keras/onnx/helpers.py | 370 +------ src/pquant/core/keras/onnx/layer_builders.py | 533 +++------- src/pquant/core/onnx_common.py | 519 ++++++++++ src/pquant/core/torch/onnx/convert_to_onnx.py | 962 +++++++----------- src/pquant/core/torch/onnx/helpers.py | 300 ------ src/pquant/core/torch/onnx/layer_builders.py | 666 +++++------- tests/test_torch_onnx_converter.py | 109 -- 8 files changed, 1422 insertions(+), 2322 deletions(-) create mode 100644 src/pquant/core/onnx_common.py delete mode 100644 src/pquant/core/torch/onnx/helpers.py diff --git a/src/pquant/core/keras/onnx/convert_to_onnx.py b/src/pquant/core/keras/onnx/convert_to_onnx.py index 6fe6df4..9be8740 100644 --- a/src/pquant/core/keras/onnx/convert_to_onnx.py +++ b/src/pquant/core/keras/onnx/convert_to_onnx.py @@ -21,7 +21,6 @@ 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 @@ -34,15 +33,7 @@ PQDepthwiseConv2d, PQMultiheadAttention, ) -from pquant.core.keras.onnx.helpers import ( - emit_getitem, - emit_squeeze, - emit_unsqueeze, - keras_dtype_to_tp, - qdq_node, - quant_node, - to_np, -) +from pquant.core.keras.onnx.helpers import keras_dtype_to_tp from pquant.core.keras.onnx.layer_builders import ( add_avgpool, add_batchnorm, @@ -53,17 +44,104 @@ 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)] - arr = np.asarray(to_np(mask)) - name = f"{prefix}_{kind}_const" - initializers.append(onh.from_array(arr, name=name)) - return name + 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 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( @@ -81,28 +159,8 @@ def emit_layer( """Emit ONNX nodes for a single Keras layer. Returns the ONNX output name.""" if isinstance(layer, PQMultiheadAttention): - 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] - kwargs = layer._inbound_nodes[0].arguments.kwargs if layer._inbound_nodes 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_integer_weights, - key_padding_mask=kpm, - attn_mask=attn_mask, + 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): @@ -114,27 +172,15 @@ def emit_layer( if isinstance(layer, PQDepthwiseConv2d): return add_depthwise_conv(layer, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights) - if isinstance(layer, PQConv2d): + if isinstance(layer, (PQConv2d, PQConv1d)): + ndim = 2 if isinstance(layer, PQConv2d) else 1 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, + ndim=ndim, quant_fn=quant_fn, use_qonnx=use_qonnx, store_integer_weights=store_integer_weights, @@ -144,47 +190,16 @@ def emit_layer( return add_batchnorm(layer, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights) if type(layer).__name__ == "GetItem": - # keras.ops GetItem operation recorded by ``x[...]`` KerasTensor syntax. - node = layer._inbound_nodes[0] - args = node.arguments.args - spec = args[1] if len(args) > 1 else node.arguments.kwargs["key"] - rank = len(args[0].shape) - return emit_getitem(prefix, current, spec, rank, nodes, initializers) + return add_getitem_op(layer, prefix, current, nodes, initializers) if type(layer).__name__ == "ExpandDims": - # keras.ops.expand_dims operation; the axis is stored on the op. - rank = len(layer._inbound_nodes[0].arguments.args[0].shape) - return emit_unsqueeze(prefix, current, [int(layer.axis) % (rank + 1)], nodes, initializers) + return add_expand_dims_op(layer, prefix, current, nodes, initializers) if type(layer).__name__ == "Squeeze": - # 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 = layer._inbound_nodes[0].arguments.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) + return add_squeeze_op(layer, prefix, current, nodes, initializers) 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 + return add_standard_activation(layer, prefix, current, nodes) if isinstance(layer, keras.layers.Flatten): out = f"{prefix}_flatten" @@ -192,12 +207,9 @@ def emit_layer( 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" + 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" - 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 @@ -207,71 +219,66 @@ def emit_layer( 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.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.GlobalAveragePooling2D): - return add_global_avgpool(layer, prefix, current, nodes, ndim=2) + 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.GlobalAveragePooling1D): - return add_global_avgpool(layer, prefix, current, nodes, ndim=1) + 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,)): + 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): - 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 + """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 [] - node = layer._inbound_nodes[0] - input_tensors = node.input_tensors + 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: - key = id(t) - if key not in tensor_to_onnx: + 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[key]) + 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 - node = layer._inbound_nodes[0] - out_tensors = node.output_tensors + 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)): @@ -327,7 +334,6 @@ def convert_to_onnx( onnx_nodes: list[onnx.NodeProto] = [] initializers: list[onnx.TensorProto] = [] - tensor_to_onnx = build_tensor_onnx_map(model) last_output_name: str = "" @@ -339,13 +345,11 @@ def convert_to_onnx( if not input_onnx_names: continue - current = input_onnx_names[0] prefix = layer.name.replace("/", "_").replace(":", "_") - output_name = emit_layer( layer, prefix, - current, + input_onnx_names[0], onnx_nodes, initializers, quant_fn, @@ -358,26 +362,23 @@ def convert_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 - n_in = len(model.inputs) - if n_in == 1: - input_names = ["input"] + input_names = model_input_names(model) + if len(model.inputs) == 1: input_shapes = [tuple(input_shape)] else: - input_names = [f"input_{i}" for i in range(n_in)] 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 n_in == 1 else dummies, training=False) + 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 - output_shape = [batch_dim] + list(dummy_out_np.shape[1:]) + 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, output_shape) + 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, @@ -386,20 +387,6 @@ def convert_to_onnx( 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 - - _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) + 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 index 0e348e0..0541608 100644 --- a/src/pquant/core/keras/onnx/helpers.py +++ b/src/pquant/core/keras/onnx/helpers.py @@ -1,189 +1,11 @@ -""" -Low-level ONNX node emitters and small utilities shared by the PQuant -Keras → ONNX converter. - -Fixed-point (k, i, f) mapping ------------------------------- -QONNX: - scale = 2^(-f) - zero_point = 0 - bit_width = k + i + f - signed = int(k) +"""Keras-specific utilities for the PQuant Keras → ONNX converter. -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). +The backend-agnostic node emitters live in ``pquant.core.onnx_common``. """ import keras -import numpy as np -import onnx.helper as oh -import onnx.numpy_helper as onh from onnx import TensorProto -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 - - -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). - """ - 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 - - -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 - def keras_dtype_to_tp(dtype): """Map a Keras/numpy dtype string to an ONNX TensorProto dtype (default float32).""" @@ -197,16 +19,23 @@ def keras_dtype_to_tp(dtype): }.get(str(dtype), TensorProto.FLOAT) -def to_np(tensor): - return np.array(tensor, dtype=np.float32) +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). We assume Keras uses channels_last format, - so we must insert Transpose nodes around the BN op. + 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) @@ -225,180 +54,7 @@ def bn_transpose_info(layer): # 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 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 emit_param(prefix, name, arr, quantizer, nodes, initializers, use_qonnx, store_integer_weights, out_channels=None): - """Emit the ONNX value for a learnable parameter (kernel/bias/gamma/beta) and return its name""" - if use_qonnx: - fp_name = f"{prefix}_{name}_fp" - initializers.append(onh.from_array(arr, name=fp_name)) - k, i, f = quantizer.get_quantization_bits() - q_nodes, out = quant_node( - f"{prefix}_{name}", - fp_name, - quantizer.round_mode, - to_np(k), - to_np(i), - to_np(f), - initializers, - overflow_mode=quantizer.overflow, - ) - nodes.extend(q_nodes) - return out - if store_integer_weights: - k, i, f = quantizer.get_quantization_bits() - if out_channels is not None: - k_a = weight_f_for_onnx(to_np(k), out_channels) - i_a = weight_f_for_onnx(to_np(i), out_channels) - f_a = weight_f_for_onnx(to_np(f), out_channels) - else: - k_a, i_a, f_a = to_np(k), to_np(i), to_np(f) - q_nodes, out = int_weight_node(f"{prefix}_{name}", arr, k_a, i_a, f_a, initializers) - nodes.extend(q_nodes) - return out - out = f"{prefix}_{name}" - initializers.append(onh.from_array(arr, name=out)) - return out - - -def maybe_quant_input(layer, prefix, current, nodes, initializers, quant_fn): - if getattr(layer, "input_quantizer", None) is not None and layer.quantize_input and layer.enable_quantization: - q = layer.input_quantizer - k, i, f = q.get_quantization_bits() - new_nodes, current = quant_fn( - f"{prefix}_in", - current, - q.round_mode, - to_np(k), - to_np(i), - to_np(f), - initializers, - overflow_mode=q.overflow, - ) - 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 layer.quantize_output and layer.enable_quantization: - q = layer.output_quantizer - k, i, f = q.get_quantization_bits() - new_nodes, current = quant_fn( - f"{prefix}_out", - current, - q.round_mode, - to_np(k), - to_np(i), - to_np(f), - initializers, - overflow_mode=q.overflow, - ) - 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", keras.config.image_data_format()) == "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) - - -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)): - name = f"{prefix}_slice_{part}" - initializers.append(onh.from_array(np.array(vals, dtype=np.int64), name=name)) - slice_inputs.append(name) - current = f"{prefix}_slice" - nodes.append(oh.make_node("Slice", inputs=slice_inputs, outputs=[current])) - if squeeze_axes: - # Squeeze takes axes as an input tensor from opset 13 on (the converter minimum). - ax_name = f"{prefix}_squeeze_axes" - initializers.append(onh.from_array(np.array(squeeze_axes, dtype=np.int64), name=ax_name)) - out = f"{prefix}_squeeze" - nodes.append(oh.make_node("Squeeze", inputs=[current, ax_name], outputs=[out])) - current = out - 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).""" - if not axes: - return input_name - ax_name = f"{prefix}_squeeze_axes" - initializers.append(onh.from_array(np.array(sorted(axes), dtype=np.int64), name=ax_name)) - 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 = f"{prefix}_unsqueeze_axes" - initializers.append(onh.from_array(np.array(axes, dtype=np.int64), name=ax_name)) - out = f"{prefix}_unsqueeze" - nodes.append(oh.make_node("Unsqueeze", inputs=[input_name, ax_name], outputs=[out])) - return out diff --git a/src/pquant/core/keras/onnx/layer_builders.py b/src/pquant/core/keras/onnx/layer_builders.py index cf6609b..806118e 100644 --- a/src/pquant/core/keras/onnx/layer_builders.py +++ b/src/pquant/core/keras/onnx/layer_builders.py @@ -1,15 +1,16 @@ -"""Per-layer ONNX graph builders (Dense/Conv/BN/Pool/Softmax/MHA) for the PQuant Keras converter.""" +"""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 -import onnx.numpy_helper as onh -from onnx import TensorProto from pquant.core.keras.layers import PQBatchNormalization -from pquant.core.keras.onnx.helpers import ( +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, - bn_transpose_info, - channels_last, + conv_padding_attrs, + emit_mha_core, emit_param, maybe_quant_input, maybe_quant_output, @@ -21,84 +22,71 @@ 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 # [out, in] - out_units = kernel_np.shape[0] - + 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, out_units + 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: - bias_np = to_np(layer._bias) - q_bias = emit_param( - prefix, "bias", bias_np, layer.bias_quantizer, nodes, initializers, use_qonnx, store_integer_weights + gemm_inputs.append( + emit_param( + prefix, + "bias", + to_np(layer._bias), + layer.bias_quantizer, + nodes, + initializers, + use_qonnx, + store_integer_weights, + ) ) - 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): - cl = channels_last(layer) + return maybe_quant_output(layer, prefix, gemm_out, nodes, initializers, quant_fn) - 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) +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) - # 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] + 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: - kernel_onnx = np.transpose(kernel_np, (2, 1, 0)) # [kL,in,out] → [out,in,kL] - - out_channels = kernel_onnx.shape[0] - - q_weight = emit_param( - prefix, - "weight", - kernel_onnx, - layer.weight_quantizer, - nodes, - initializers, - use_qonnx, - store_integer_weights, - out_channels, - ) + q_weight = f"{prefix}_weight_t" + add_initializer(initializers, q_weight, kernel_np.T) # pre-transposed [in, out] - conv_inputs = [current, q_weight] + 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: - bias_np = to_np(layer._bias) q_bias = emit_param( - prefix, "bias", bias_np, layer.bias_quantizer, nodes, initializers, use_qonnx, store_integer_weights + prefix, "bias", to_np(layer._bias), layer.bias_quantizer, nodes, initializers, use_qonnx, store_integer_weights ) - conv_inputs.append(q_bias) + 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) - 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" +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=getattr(layer, "groups", 1), + group=groups, auto_pad=auto_pad, ) if pads is not None: @@ -106,122 +94,80 @@ def add_conv(layer, prefix, current, nodes, initializers, ndim, quant_fn, use_qo conv_out = f"{prefix}_conv" nodes.append(oh.make_node("Conv", inputs=conv_inputs, outputs=[conv_out], **conv_attrs)) - current = conv_out + return 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) +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) - kernel_np = to_np(layer._kernel) # [kH, kW, in, depth_mult] - in_ch, depth_mult = kernel_np.shape[2], kernel_np.shape[3] - 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] - - q_weight = emit_param( - prefix, - "weight", - kernel_onnx, - layer.weight_quantizer, - nodes, - initializers, - use_qonnx, - store_integer_weights, - out_channels, - ) - + 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: - bias_np = to_np(layer._bias) - q_bias = emit_param( - prefix, "bias", bias_np, layer.bias_quantizer, nodes, initializers, use_qonnx, store_integer_weights + conv_inputs.append( + emit_param(prefix, "bias", to_np(layer._bias), layer.bias_quantizer, nodes, initializers, use_qonnx, store_int) ) - 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" + current = add_conv_node(layer, prefix, conv_inputs, groups, ndim, nodes) + current = maybe_quant_output(layer, prefix, current, nodes, initializers, quant_fn) - 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 + if is_channels_last: + current = add_transpose(f"{prefix}_post", current, perm_to_nhwx, nodes) + return current - 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) +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 + ) - if cl: - current = add_transpose(f"{prefix}_post", current, [0, 2, 3, 1], nodes) - return current + +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_tr, perm_to_nchw, perm_to_nhwx = bn_transpose_info(layer) - - if need_tr: + 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) - is_pq = isinstance(layer, PQBatchNormalization) + 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 - gamma_np = to_np(layer.gamma) if layer.gamma is not None else None - beta_np = to_np(layer.beta) if layer.beta is not None else None - - if gamma_np is None: - # scale=False: use ones - n_ch = to_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 = to_np(layer.moving_mean).shape[0] - beta_np = np.zeros(n_ch, dtype=np.float32) - - qonnx_p = use_qonnx and is_pq - intstore_p = store_integer_weights and is_pq - q_gamma = emit_param( - prefix, "gamma", gamma_np, layer.weight_quantizer if is_pq else None, nodes, initializers, qonnx_p, intstore_p - ) - q_beta = emit_param( - prefix, "beta", beta_np, layer.bias_quantizer if is_pq else None, nodes, initializers, qonnx_p, intstore_p - ) + # 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 - mean_name = f"{prefix}_running_mean" - var_name = f"{prefix}_running_var" - initializers.append(onh.from_array(to_np(layer.moving_mean), name=mean_name)) - initializers.append(onh.from_array(to_np(layer.moving_variance), name=var_name)) + 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( @@ -234,123 +180,11 @@ def add_batchnorm(layer, prefix, current, nodes, initializers, quant_fn, use_qon ) current = bn_out - if need_tr: + if need_transpose: 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): - current = maybe_quant_input(layer, prefix, current, nodes, initializers, quant_fn) - - kernel_np = to_np(layer._kernel).T # [out, in] - out_units = kernel_np.shape[0] - - q_weight = emit_param( - prefix, "weight", kernel_np, layer.weight_quantizer, nodes, initializers, use_qonnx, store_integer_weights, out_units - ) - - # 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 = to_np(layer._bias) - q_bias = emit_param( - prefix, "bias", bias_np, 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 - - current = maybe_quant_output(layer, prefix, current, nodes, initializers, quant_fn) - return current - - -def add_quantized_softmax(sm, prefix, current, nodes, initializers, quant_fn, kpm_mask=None): - enable = sm.enable_quantization - scaler = float(sm.input_scaler) - stable = bool(sm.stable) - eps = float(sm.epsilon) - - def qdq(q, pfx, x): - k, i, f = q.get_quantization_bits() - q_nodes, out = quant_fn(pfx, x, q.round_mode, to_np(k), to_np(i), to_np(f), initializers, overflow_mode=q.overflow) - nodes.extend(q_nodes) - return out - - # 1) Softmax input quantizer. - if sm.quantize_input and enable: - current = qdq(sm.input_quantizer, f"{prefix}_sm_in_q", current) - - # 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 (only when quantize_input==stable), - # Exp of (-scaler * x) for the stable branch (+scaler otherwise), output QDQ. - exp_t = sm.exp_table - if exp_t.quantize_input and enable: - exp_in = qdq(exp_t.input_quantizer, f"{prefix}_sm_exp_in_q", exp_in) - coeff = -scaler if stable else scaler - exp_arg = exp_in - if coeff != 1.0: - coeff_name = f"{prefix}_sm_exp_coeff" - initializers.append(onh.from_array(np.array(coeff, dtype=np.float32), name=coeff_name)) - exp_arg = f"{prefix}_sm_exp_arg" - nodes.append(oh.make_node("Mul", inputs=[exp_in, coeff_name], outputs=[exp_arg])) - exp_inp = f"{prefix}_sm_exp" - nodes.append(oh.make_node("Exp", inputs=[exp_arg], outputs=[exp_inp])) - if exp_t.quantize_output and enable: - exp_inp = qdq(exp_t.output_quantizer, f"{prefix}_sm_exp_out_q", exp_inp) - - # 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=TensorProto.FLOAT)) - masked = f"{prefix}_sm_masked" - nodes.append(oh.make_node("Mul", inputs=[kpm_f, exp_inp], outputs=[masked])) - exp_inp = masked - - # 4) Sum over the last axis (ReduceSum takes axes as an input from opset 13). - sum_axes = f"{prefix}_sm_sum_axes" - initializers.append(onh.from_array(np.array([-1], dtype=np.int64), name=sum_axes)) - sums = f"{prefix}_sm_sum" - nodes.append(oh.make_node("ReduceSum", inputs=[exp_inp, sum_axes], outputs=[sums], keepdims=1)) - - # 5) Quantized reciprocal table: input QDQ, 1/(x+eps), output QDQ. - inv_t = sm.inv_table - inv_in = sums - if inv_t.quantize_input and enable: - inv_in = qdq(inv_t.input_quantizer, f"{prefix}_sm_inv_in_q", inv_in) - eps_name = f"{prefix}_sm_eps" - initializers.append(onh.from_array(np.array(eps, dtype=np.float32), name=eps_name)) - 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_t.quantize_output and enable: - divisor = qdq(inv_t.output_quantizer, f"{prefix}_sm_inv_out_q", divisor) - - # 6) Multiply numerator by reciprocal. - out = f"{prefix}_sm_out" - nodes.append(oh.make_node("Mul", inputs=[exp_inp, divisor], outputs=[out])) - current = out - - # 7) Softmax output quantizer. - if sm.quantize_output and enable: - current = qdq(sm.output_quantizer, f"{prefix}_sm_out_q", current) - return current - - def add_mha( layer, prefix, @@ -365,12 +199,7 @@ def add_mha( key_padding_mask=None, attn_mask=None, ): - 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 / 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 ) @@ -381,126 +210,21 @@ def add_mha( 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 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 = f"{prefix}_kpm_axes" - initializers.append(onh.from_array(np.array([1, 2], dtype=np.int64), name=kpm_axes)) - 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])) - - current = add_quantized_softmax( - layer.softmax, f"{prefix}_attn", current, nodes, initializers, quant_fn, kpm_mask=kpm_mult + 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 ) - attn_w_name = current # softmax output = attention weights (also averaged over heads below) - - 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) --- + + # 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 + layer.out_proj, f"{prefix}_out_proj", context, 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] + 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) @@ -517,31 +241,27 @@ def add_avgpool(layer, prefix, current, nodes, initializers, ndim, quant_fn): count_include_pad=0, ) ) - current = pool_out - - current = maybe_quant_output(layer, prefix, current, nodes, initializers, quant_fn) + current = maybe_quant_output(layer, prefix, pool_out, nodes, initializers, quant_fn) - if cl: + 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): - cl = channels_last(layer) - - if cl: - perm_to_nchw = [0, 3, 1, 2] if ndim == 2 else [0, 2, 1] + 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 cl: + 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 @@ -549,32 +269,23 @@ 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"): - m_val = float(np.array(layer.multiplier).ravel()[0]) - scale = float(2.0 ** round(m_val)) - scale_name = f"{prefix}_mul_scale" + 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" - 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 - act = layer.activation_name + activation = layer.activation_name act_out = f"{prefix}_act" - if act == "relu": + if activation == "relu": nodes.append(oh.make_node("Relu", inputs=[current], outputs=[act_out])) - elif act == "tanh": + elif activation == "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), - ] + 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 {act!r} for ONNX export") - current = act_out + raise TypeError(f"PQActivation: unsupported activation {activation!r} for ONNX export") - # --- optional output quantization --- - current = maybe_quant_output(layer, prefix, current, nodes, initializers, quant_fn) - return current + 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/onnx/convert_to_onnx.py b/src/pquant/core/torch/onnx/convert_to_onnx.py index 9f683c3..4b41631 100644 --- a/src/pquant/core/torch/onnx/convert_to_onnx.py +++ b/src/pquant/core/torch/onnx/convert_to_onnx.py @@ -9,21 +9,32 @@ import functools import logging -import operator as _operator +import 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 +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, @@ -36,22 +47,16 @@ PQLayerNorm, PQMultiheadAttention, ) -from pquant.core.torch.onnx.helpers import ( # noqa: E402 - emit_getitem, - emit_squeeze, - emit_unsqueeze, - maybe_quant_input, - maybe_quant_output, - qdq_node, - quant_node, -) from pquant.core.torch.onnx.layer_builders import ( # noqa: E402 + add_activation, add_avgpool, add_batchnorm, add_conv, add_dense, add_layernorm, + add_maxpool, add_mha, + add_upsample, ) from pquant.core.torch.quantizer import Quantizer # noqa: E402 @@ -62,26 +67,15 @@ def emit_module(module, prefix, current, nodes, initializers, quant_fn, use_qonn return add_dense( module, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights, integer_ops ) - if isinstance(module, PQConv2d): + if isinstance(module, (PQConv2d, PQConv1d)): + ndim = 2 if isinstance(module, PQConv2d) else 1 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, + ndim=ndim, quant_fn=quant_fn, use_qonnx=use_qonnx, store_integer_weights=store_integer_weights, @@ -90,282 +84,34 @@ def emit_module(module, prefix, current, nodes, initializers, quant_fn, use_qonn 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, (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.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.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]], - ) - ) + 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): - # 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): - out, _ = add_mha( - module, prefix, current, current, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights - ) - return out - if isinstance(module, Quantizer): - k, i, f = module.get_quantization_bits() - new_nodes, out = quant_fn(prefix, current, module.round_mode, k, i, f, initializers, overflow_mode=module.overflow) - nodes.extend(new_nodes) - return out + 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__}") -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: - # ----- 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") - - 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 - - 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"), - ] - - 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", - ), - ] - - 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 - - _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 - - class PQTracer(fx.Tracer): _LEAF_TYPES = ( PQDense, @@ -427,6 +173,311 @@ def normalize_input_dtypes(input_dtypes, n: int): 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, + ) + + 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, @@ -502,317 +553,19 @@ def convert_to_onnx( input_shapes = normalize_input_shapes(input_shape) input_torch_dtypes, input_tp_dtypes = normalize_input_dtypes(input_dtypes, len(input_shapes)) - graph = PQTracer().trace(model, concrete_args=concrete_args) - gm = fx.GraphModule(model, graph) - - 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() - - tensor_phs = list(gm.graph.find_nodes(op="placeholder")) - if len(tensor_phs) != len(input_shapes): - raise ValueError( - f"FX export: model.forward has {len(tensor_phs)} tensor input(s) but " - f"input_shape describes {len(input_shapes)}. Specialize non-tensor " - f"arguments via concrete_args={{...}}." - ) - input_names = ["input"] if len(tensor_phs) == 1 else [str(p.target) for p in tensor_phs] - ph_to_name = {p: n for p, n in zip(tensor_phs, input_names)} - - from torch.fx.passes.shape_prop import ShapeProp + 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) - 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]: - 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 node_shape(n: fx.Node) -> tuple: - 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 tuple(meta.shape) - - def node_rank(n: fx.Node) -> int: - return len(node_shape(n)) - - def squeeze_axes_for(node: fx.Node) -> list[int]: - """Resolve the axes a torch squeeze()/​.squeeze() call removes.""" - in_shape = 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 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] = ph_to_name[node] - - elif node.op == "get_attr": - 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): - # forward(query, key, value, key_padding_mask=None, attn_mask=None, ...) - 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 - - def mask_name(pos, kw, node=node): - 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 node_to_name[arg] - - kpm_name = mask_name(3, "key_padding_mask") - attn_mask_name = mask_name(4, "attn_mask") - 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, - key_padding_mask=kpm_name, - attn_mask=attn_mask_name, - ) - 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 torch._assert or getattr(fn, "__name__", "") == "_assert" or fn is _operator.eq: - continue - - if fn is _operator.getitem: - container = node_to_name[node.args[0]] - if isinstance(container, tuple): - # Unpack a tuple output (e.g. from PQMultiheadAttention). - node_to_name[node] = container[node.args[1]] - else: - # Tensor slicing: x[:, 0], x[..., :4], ... → Slice (+ Squeeze) - rank = node_rank(node.args[0]) - node_to_name[node] = emit_getitem(node.name, container, node.args[1], rank, onnx_nodes, initializers) - 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 = node_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 = node_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 - - elif fn is torch.squeeze: - node_to_name[node] = emit_squeeze( - node.name, res(node.args[0]), squeeze_axes_for(node), onnx_nodes, initializers - ) - - elif fn is torch.unsqueeze: - dim = int(node.args[1]) if len(node.args) > 1 else int(node.kwargs["dim"]) - axes = [dim % (node_rank(node.args[0]) + 1)] - node_to_name[node] = emit_unsqueeze(node.name, res(node.args[0]), axes, onnx_nodes, initializers) - - 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 = node_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 = node_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 - - elif node.target == "squeeze": - node_to_name[node] = emit_squeeze(node.name, x, squeeze_axes_for(node), onnx_nodes, initializers) - - elif node.target == "unsqueeze": - dim = int(node.args[1]) if len(node.args) > 1 else int(node.kwargs["dim"]) - axes = [dim % (node_rank(node.args[0]) + 1)] - node_to_name[node] = emit_unsqueeze(node.name, x, axes, onnx_nodes, initializers) - - 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) - - graph_input_names = set(input_names) - for idx, nm in enumerate(output_names): - if nm in graph_input_names: - ident = f"{nm}_identity_out{idx}" - onnx_nodes.append(oh.make_node("Identity", inputs=[nm], outputs=[ident])) - output_names[idx] = ident + 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 {})) @@ -828,22 +581,13 @@ def mask_name(pos, kw, node=node): for name, t in zip(output_names, dummy_outs) ] - onnx_graph = oh.make_graph( - nodes=onnx_nodes, + graph = oh.make_graph( + nodes=emitter.nodes, name="pquant_onnx_fx", inputs=input_vis, outputs=output_vis, - initializer=initializers, + initializer=emitter.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) + 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/helpers.py b/src/pquant/core/torch/onnx/helpers.py deleted file mode 100644 index a870c7d..0000000 --- a/src/pquant/core/torch/onnx/helpers.py +++ /dev/null @@ -1,300 +0,0 @@ -""" -Low-level ONNX node emitters and small utilities shared by the PQuant -torch → ONNX converter. - -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). -""" - -import numpy as np -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 quant_node(name_prefix, input_name, rounding_mode, k, i, f, initializers, overflow_mode="SAT"): - k_val = int(k.item()) - if 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 - - -def qdq_node( - name_prefix, input_name, rounding_mode, k, i, f, initializers, overflow_mode="SAT", include_clip=True -): # noqa: ARG001 - 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 - - -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()) - 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 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 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 maybe_quant_input(module, prefix, current, nodes, initializers, quant_fn): - # input_quantizer is created conditionally, so guard it; the bool flags are always present. - if getattr(module, "input_quantizer", None) is not None and module.quantize_input and module.enable_quantization: - 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=q.overflow) - 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 module.quantize_output and module.enable_quantization: - 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=q.overflow - ) - nodes.extend(new_nodes) - return current - - -def emit_param(prefix, name, arr, quantizer, nodes, initializers, use_qonnx, store_integer_weights): - if use_qonnx: - fp_name = f"{prefix}_{name}_fp" - initializers.append(onh.from_array(arr, name=fp_name)) - 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, i, f = quantizer.get_quantization_bits() - q_nodes, out = int_weight_node(f"{prefix}_{name}", arr, k, i, f, initializers) - nodes.extend(q_nodes) - return out - out = f"{prefix}_{name}" - initializers.append(onh.from_array(arr, name=out)) - return out - - -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)): - name = f"{prefix}_slice_{part}" - initializers.append(onh.from_array(np.array(vals, dtype=np.int64), name=name)) - slice_inputs.append(name) - current = f"{prefix}_slice" - nodes.append(oh.make_node("Slice", inputs=slice_inputs, outputs=[current])) - if squeeze_axes: - # Squeeze takes axes as an input tensor from opset 13 on (the converter minimum). - ax_name = f"{prefix}_squeeze_axes" - initializers.append(onh.from_array(np.array(squeeze_axes, dtype=np.int64), name=ax_name)) - out = f"{prefix}_squeeze" - nodes.append(oh.make_node("Squeeze", inputs=[current, ax_name], outputs=[out])) - current = out - 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).""" - if not axes: - return input_name - ax_name = f"{prefix}_squeeze_axes" - initializers.append(onh.from_array(np.array(sorted(axes), dtype=np.int64), name=ax_name)) - 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 = f"{prefix}_unsqueeze_axes" - initializers.append(onh.from_array(np.array(axes, dtype=np.int64), name=ax_name)) - out = f"{prefix}_unsqueeze" - nodes.append(oh.make_node("Unsqueeze", inputs=[input_name, ax_name], outputs=[out])) - return out diff --git a/src/pquant/core/torch/onnx/layer_builders.py b/src/pquant/core/torch/onnx/layer_builders.py index 774aefb..868d0cd 100644 --- a/src/pquant/core/torch/onnx/layer_builders.py +++ b/src/pquant/core/torch/onnx/layer_builders.py @@ -1,142 +1,130 @@ -"""Per-layer ONNX graph builders (Dense/Conv/BN/LN/AvgPool/Softmax/MHA) for the PQuant torch converter.""" +"""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 -import onnx.numpy_helper as onh -from onnx import TensorProto -from pquant.core.torch.onnx.helpers import ( +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, - torch_padding_to_onnx, + to_np, ) -def add_dense_integer(module, prefix, current, nodes, initializers): - 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") +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 - # --- 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]), - ] + 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] - # --- 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()) # get_quantization_bits() always returns tensors - 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 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), - ] + 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 - # --- 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], - ) - ) - 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) +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 the output tensor is [batch, out] and out is axis 1. + """ + 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: - combined_scale_np = np.array(float(s_x * s_w_1d[0]), dtype=np.float32) - combined_zp_np = np.array(np.int32(0)) + scale_np = np.array(float(combined_scale_1d[0]), dtype=np.float32) + 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 + 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) - current = maybe_quant_output(module, prefix, current, nodes, initializers, qdq_node) - return current + 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 = module._weight.detach().cpu().numpy().astype(np.float32) # [out, in] + 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( @@ -146,23 +134,21 @@ def add_dense_nd(module, prefix, current, nodes, initializers, quant_fn, use_qon nodes.append(oh.make_node("Transpose", inputs=[q_weight_native], outputs=[q_weight], perm=[1, 0])) else: q_weight = f"{prefix}_weight_T" - initializers.append(onh.from_array(weight_np.T, name=q_weight)) # pre-transposed [in, out] + 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: - bias_np = module._bias.detach().cpu().numpy().astype(np.float32) q_bias = emit_param( - prefix, "bias", bias_np, module.bias_quantizer, nodes, initializers, use_qonnx, store_integer_weights + 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 - current = maybe_quant_output(module, prefix, current, nodes, initializers, quant_fn) - return current + 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): @@ -170,53 +156,68 @@ def add_dense(module, prefix, current, nodes, initializers, quant_fn, 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) q_weight = emit_param( - prefix, "weight", weight_np, module.weight_quantizer, nodes, initializers, use_qonnx, store_integer_weights + 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: - bias_np = module._bias.detach().cpu().numpy().astype(np.float32) - q_bias = emit_param( - prefix, "bias", bias_np, module.bias_quantizer, nodes, initializers, use_qonnx, store_integer_weights + gemm_inputs.append( + emit_param( + prefix, + "bias", + to_np(module._bias), + module.bias_quantizer, + nodes, + initializers, + use_qonnx, + store_integer_weights, + ) ) - 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 + 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) - weight_np = module._weight.detach().cpu().numpy().astype(np.float32) q_weight = emit_param( - prefix, "weight", weight_np, module.weight_quantizer, nodes, initializers, use_qonnx, store_integer_weights + 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: - bias_np = module._bias.detach().cpu().numpy().astype(np.float32) - q_bias = emit_param( - prefix, "bias", bias_np, module.bias_quantizer, nodes, initializers, use_qonnx, store_integer_weights + conv_inputs.append( + emit_param( + prefix, + "bias", + to_np(module._bias), + module.bias_quantizer, + nodes, + initializers, + use_qonnx, + store_integer_weights, + ) ) - 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) + 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), @@ -229,29 +230,28 @@ def add_conv(module, prefix, current, nodes, initializers, ndim, quant_fn, use_q 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 + 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) - gamma_np = module._weight.detach().cpu().numpy().astype(np.float32) - beta_np = module._bias.detach().cpu().numpy().astype(np.float32) - q_gamma = emit_param( - prefix, "gamma", gamma_np, module.weight_quantizer, nodes, initializers, use_qonnx, store_integer_weights + prefix, + "gamma", + to_np(module._weight), + module.weight_quantizer, + nodes, + initializers, + use_qonnx, + store_integer_weights, ) q_beta = emit_param( - prefix, "beta", beta_np, module.bias_quantizer, nodes, initializers, use_qonnx, store_integer_weights + prefix, "beta", to_np(module._bias), module.bias_quantizer, nodes, initializers, use_qonnx, store_integer_weights ) - - 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)) + 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( @@ -268,59 +268,32 @@ def add_batchnorm(module, prefix, current, nodes, initializers, quant_fn, use_qo 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) - ns = ( - tuple(int(d) for d in module.normalized_shape) - if hasattr(module.normalized_shape, "__iter__") - else (int(module.normalized_shape),) - ) - axis = -len(ns) + 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 = 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 + 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) - qonnx_p = use_qonnx and has_weight - intstore_p = store_integer_weights and has_weight - q_gamma = emit_param( - prefix, - "gamma", - gamma_np, - module.weight_quantizer if has_weight else None, - nodes, - initializers, - qonnx_p, - intstore_p, - ) + ln_inputs = [current, q_gamma] if has_bias: + bias_quantizer = module.bias_quantizer if has_weight else None q_beta = emit_param( - prefix, - "beta", - beta_np, - module.bias_quantizer if has_weight else None, - nodes, - initializers, - qonnx_p, - intstore_p, + prefix, "beta", to_np(module._bias), bias_quantizer, nodes, initializers, use_qonnx, store_integer_weights ) - - 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), - ) + 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 + return maybe_quant_output(module, prefix, ln_out, nodes, initializers, quant_fn) def add_avgpool(module, prefix, current, nodes, initializers, ndim, quant_fn): @@ -334,87 +307,118 @@ def add_avgpool(module, prefix, current, nodes, initializers, ndim, quant_fn): 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), + pads=symmetric_pads(module.padding, ndim), ceil_mode=int(module.ceil_mode), count_include_pad=int(module.count_include_pad), ) ) - current = pool_out + return maybe_quant_output(module, prefix, pool_out, nodes, initializers, quant_fn) - current = maybe_quant_output(module, prefix, current, nodes, initializers, quant_fn) - return current +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_quantized_softmax(sm, prefix, current, nodes, initializers, quant_fn, kpm_mask=None): - enable = sm.enable_quantization - scaler = float(sm.input_scaler) - stable = bool(sm.stable) - eps = float(sm.epsilon) - def qdq(q, pfx, x): - k, i, f = q.get_quantization_bits() - q_nodes, out = quant_fn(pfx, x, q.round_mode, k, i, f, initializers, overflow_mode=q.overflow) - nodes.extend(q_nodes) - 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" - if sm.quantize_input and enable: - current = qdq(sm.input_quantizer, f"{prefix}_sm_in_q", current) + 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 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])) + +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: - exp_in = current - - exp_t = sm.exp_table - if exp_t.quantize_input and enable: - exp_in = qdq(exp_t.input_quantizer, f"{prefix}_sm_exp_in_q", exp_in) - coeff = -scaler if stable else scaler - exp_arg = exp_in - if coeff != 1.0: - coeff_name = f"{prefix}_sm_exp_coeff" - initializers.append(onh.from_array(np.array(coeff, dtype=np.float32), name=coeff_name)) - exp_arg = f"{prefix}_sm_exp_arg" - nodes.append(oh.make_node("Mul", inputs=[exp_in, coeff_name], outputs=[exp_arg])) - exp_inp = f"{prefix}_sm_exp" - nodes.append(oh.make_node("Exp", inputs=[exp_arg], outputs=[exp_inp])) - if exp_t.quantize_output and enable: - exp_inp = qdq(exp_t.output_quantizer, f"{prefix}_sm_exp_out_q", exp_inp) - - 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=TensorProto.FLOAT)) - masked = f"{prefix}_sm_masked" - nodes.append(oh.make_node("Mul", inputs=[kpm_f, exp_inp], outputs=[masked])) - exp_inp = masked - - sum_axes = f"{prefix}_sm_sum_axes" - initializers.append(onh.from_array(np.array([-1], dtype=np.int64), name=sum_axes)) - sums = f"{prefix}_sm_sum" - nodes.append(oh.make_node("ReduceSum", inputs=[exp_inp, sum_axes], outputs=[sums], keepdims=1)) - - inv_t = sm.inv_table - inv_in = sums - if inv_t.quantize_input and enable: - inv_in = qdq(inv_t.input_quantizer, f"{prefix}_sm_inv_in_q", inv_in) - eps_name = f"{prefix}_sm_eps" - initializers.append(onh.from_array(np.array(eps, dtype=np.float32), name=eps_name)) - 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_t.quantize_output and enable: - divisor = qdq(inv_t.output_quantizer, f"{prefix}_sm_inv_out_q", divisor) - - out = f"{prefix}_sm_out" - nodes.append(oh.make_node("Mul", inputs=[exp_inp, divisor], outputs=[out])) - current = out - - if sm.quantize_output and enable: - current = qdq(sm.output_quantizer, f"{prefix}_sm_out_q", current) - return current + # 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( @@ -431,21 +435,12 @@ def add_mha( key_padding_mask=None, attn_mask=None, ): - H = module.num_heads - head_dim = module.head_dim - E = module.embed_dim - scale_val = float(module.scale) - 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_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 ) @@ -456,116 +451,13 @@ def add_mha( module.v_proj, f"{prefix}_v_proj", v_input, nodes, initializers, quant_fn, use_qonnx, store_integer_weights ) - 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 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 = f"{prefix}_kpm_axes" - initializers.append(onh.from_array(np.array([1, 2], dtype=np.int64), name=kpm_axes)) - 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])) - - current = add_quantized_softmax( - module.softmax, f"{prefix}_attn", current, nodes, initializers, quant_fn, kpm_mask=kpm_mult + 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 ) - attn_w_name = current # softmax output = attention weights (also averaged over heads below) - - 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" # 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])) - out = add_dense_nd( - module.out_proj, f"{prefix}_out_proj", ctx_merged, nodes, initializers, quant_fn, use_qonnx, store_integer_weights + module.out_proj, f"{prefix}_out_proj", context, nodes, initializers, quant_fn, use_qonnx, store_integer_weights ) - avg_attn = f"{prefix}_avg_attn_weights" - nodes.append(oh.make_node("ReduceMean", inputs=[attn_w_name], outputs=[avg_attn], axes=[1], keepdims=0)) 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 - + out = add_transpose(f"{prefix}_out_seq_first", out, [1, 0, 2], nodes) return out, avg_attn diff --git a/tests/test_torch_onnx_converter.py b/tests/test_torch_onnx_converter.py index fd9a501..2c26dc4 100644 --- a/tests/test_torch_onnx_converter.py +++ b/tests/test_torch_onnx_converter.py @@ -21,7 +21,6 @@ import pquant # noqa: E402 from pquant.core.torch.layers import Quantizer # noqa: E402 from pquant.core.torch.onnx import convert_to_onnx # noqa: E402 -from pquant.core.torch.onnx.convert_to_onnx import export_qdq_layernorm # noqa: E402 from pquant.layers import ( # noqa: E402 PQActivation, PQAvgPool1d, @@ -263,83 +262,6 @@ def test_mha_key_padding_mask_onnx(cfg, bias, tmp_path): ) -@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, - ) - - # ----- structural checks ----- - 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 ----- - 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) - - class _TwoInputModel(nn.Module): """Two tensor inputs merged by addition after independent Dense layers.""" @@ -737,37 +659,6 @@ def test_qonnx_export_builds(cfg_quant, tmp_path): assert any(n.op_type == "Quant" for n in proto.graph.node) -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) - - @pytest.mark.parametrize( "slicer", [ From 182cf3a36ce64ffa586f4ac902e3112a47406607 Mon Sep 17 00:00:00 2001 From: Roope Niemi Date: Wed, 15 Jul 2026 14:34:52 +0200 Subject: [PATCH 6/8] small cleanups --- src/pquant/core/torch/onnx/convert_to_onnx.py | 4 +- tests/test_keras_onnx_converter.py | 68 +++---- tests/test_torch_onnx_converter.py | 176 +++++++++--------- 3 files changed, 127 insertions(+), 121 deletions(-) diff --git a/src/pquant/core/torch/onnx/convert_to_onnx.py b/src/pquant/core/torch/onnx/convert_to_onnx.py index 4b41631..d777e59 100644 --- a/src/pquant/core/torch/onnx/convert_to_onnx.py +++ b/src/pquant/core/torch/onnx/convert_to_onnx.py @@ -186,7 +186,7 @@ def resolve_perm_dims(args, rank: int) -> list[int]: return [int(d) % rank for d in dims] -class _FxGraphEmitter: +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 @@ -563,7 +563,7 @@ def convert_to_onnx( with torch.no_grad(): ShapeProp(gm).propagate(*probes) - emitter = _FxGraphEmitter(gm, ph_to_name, quant_fn, use_qonnx, store_integer_weights, integer_ops) + 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) diff --git a/tests/test_keras_onnx_converter.py b/tests/test_keras_onnx_converter.py index 6cb3c0d..0f6be41 100644 --- a/tests/test_keras_onnx_converter.py +++ b/tests/test_keras_onnx_converter.py @@ -30,7 +30,7 @@ QUANT_ATOL = 5e-3 -def _atol(cfg): +def atol(cfg): return QUANT_ATOL if cfg.quantization_parameters.enable_quantization else ATOL @@ -41,17 +41,17 @@ def cfg(request): return c -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) @@ -76,7 +76,7 @@ def _onnx_run(model, x: np.ndarray, input_shape: tuple, tmp_path) -> np.ndarray: 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), + lambda cfg: PQBatchNormalization(cfg, axis=1 if channels_first() else -1), 8, (4, 4), 4, @@ -89,7 +89,7 @@ def _onnx_run(model, x: np.ndarray, input_shape: tuple, tmp_path) -> np.ndarray: @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) + input_shape = (channels, *spatial) if channels_first() else (*spatial, channels) x_np = np.random.randn(batch, *input_shape).astype(np.float32) inputs = keras.Input(shape=input_shape) @@ -99,9 +99,9 @@ def test_single_layer_onnx(cfg, make_layer, channels, spatial, batch, warmup_kwa model(np.zeros((1, *input_shape), dtype=np.float32), **warmup_kwargs) 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(cfg), err_msg="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") @pytest.mark.parametrize("activation", ["relu", "tanh", "hard_tanh"]) @@ -115,10 +115,10 @@ def test_pqactivation_onnx(cfg, activation, tmp_path): apply_final_compression(model) x_np = np.random.randn(4, DIM).astype(np.float32) - keras_out = _keras_out(model, x_np) - onnx_out = _onnx_run(model, x_np, input_shape=(DIM,), tmp_path=tmp_path) + 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_out, onnx_out, atol=_atol(cfg), err_msg=f"PQActivation {activation}: keras vs ONNX mismatch" + keras_output, onnx_out, atol=atol(cfg), err_msg=f"PQActivation {activation}: keras vs ONNX mismatch" ) @@ -136,9 +136,9 @@ def test_residual_concat_onnx(cfg, tmp_path): apply_final_compression(model) x_np = np.random.randn(4, DIM).astype(np.float32) - keras_out = _keras_out(model, x_np) - onnx_out = _onnx_run(model, x_np, input_shape=(DIM,), tmp_path=tmp_path) - np.testing.assert_allclose(keras_out, onnx_out, atol=_atol(cfg), err_msg="residual+concat: keras vs ONNX mismatch") + 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") def test_two_input_onnx(cfg, tmp_path): @@ -155,7 +155,7 @@ def test_two_input_onnx(cfg, tmp_path): xa = np.random.randn(3, IN_A).astype(np.float32) xb = np.random.randn(3, IN_B).astype(np.float32) - keras_out = _keras_out(model, [xa, xb]) + 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) @@ -166,7 +166,7 @@ def test_two_input_onnx(cfg, tmp_path): 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_out, onnx_out, atol=_atol(cfg), err_msg="two-input: keras vs ONNX mismatch") + 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]) @@ -181,10 +181,10 @@ def test_mha_onnx(cfg, bias, tmp_path): apply_final_compression(model) x_np = np.random.randn(2, T, E).astype(np.float32) - keras_out = _keras_out(model, x_np) - onnx_out = _onnx_run(model, x_np, input_shape=(T, E), tmp_path=tmp_path) + 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_out, onnx_out, atol=_atol(cfg), err_msg=f"PQMultiheadAttention bias={bias}: keras vs ONNX mismatch" + keras_output, onnx_out, atol=atol(cfg), err_msg=f"PQMultiheadAttention bias={bias}: keras vs ONNX mismatch" ) @@ -201,9 +201,11 @@ def test_mha_causal_attn_mask_onnx(cfg, tmp_path): apply_final_compression(model) x_np = np.random.randn(2, T, E).astype(np.float32) - keras_out = _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_out, onnx_out, atol=_atol(cfg), err_msg="MHA causal attn_mask: keras vs ONNX mismatch") + 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_mha_key_padding_mask_onnx(cfg, tmp_path): @@ -222,7 +224,7 @@ def test_mha_key_padding_mask_onnx(cfg, tmp_path): 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_out = _keras_out(model, [x_np, mask_np]) + 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) @@ -233,7 +235,9 @@ def test_mha_key_padding_mask_onnx(cfg, tmp_path): 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_out, onnx_out, atol=_atol(cfg), err_msg="MHA key_padding_mask: keras vs ONNX mismatch") + np.testing.assert_allclose( + keras_output, onnx_out, atol=atol(cfg), err_msg="MHA key_padding_mask: keras vs ONNX mismatch" + ) @pytest.mark.parametrize( @@ -258,9 +262,9 @@ def test_tensor_slicing_onnx(cfg, slicer, tmp_path): 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(cfg), err_msg="tensor slicing: keras vs ONNX mismatch") + 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( @@ -285,7 +289,7 @@ def test_squeeze_unsqueeze_onnx(cfg, reshaper, tmp_path): 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) - assert keras_out.shape == onnx_out.shape - np.testing.assert_allclose(keras_out, onnx_out, atol=_atol(cfg), err_msg="squeeze/expand_dims: keras vs ONNX mismatch") + 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") diff --git a/tests/test_torch_onnx_converter.py b/tests/test_torch_onnx_converter.py index 2c26dc4..445e189 100644 --- a/tests/test_torch_onnx_converter.py +++ b/tests/test_torch_onnx_converter.py @@ -38,7 +38,7 @@ QUANT_ATOL = 5e-3 -def _atol(cfg): +def atol(cfg): return QUANT_ATOL if cfg.quantization_parameters.enable_quantization else ATOL @@ -56,13 +56,13 @@ def cfg_quant(): 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,7 +71,7 @@ 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(model, input_shape=input_shape, output_path=path) @@ -80,7 +80,7 @@ def _onnx_run_fx(model: nn.Module, x: torch.Tensor, input_shape: tuple, tmp_path 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() @@ -145,14 +145,14 @@ def test_single_layer_onnx(cfg, make_model, input_shape, batch, tmp_path): x = torch.randn(batch, *input_shape) with torch.no_grad(): model(x) # warm-up in train mode (initialises any running stats) - _apply_compression(model) + apply_compression(model) - torch_out = _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_out, onnx_out, atol=ATOL, err_msg="torch vs ONNX mismatch") + 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") -class _SelfAttnModel(nn.Module): +class SelfAttnModel(nn.Module): """Thin wrapper so FX tracing sees a single-input model.""" def __init__(self, mha: PQMultiheadAttention): @@ -168,23 +168,23 @@ def forward(self, x): 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) + 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_out = _torch_out(model, x) - onnx_out = _onnx_run_fx(model, x, input_shape=(T, E), tmp_path=tmp_path) + # 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_out, onnx_out, atol=_atol(cfg), err_msg=f"PQMultiheadAttention bias={bias}: torch vs ONNX mismatch" + torch_output, onnx_out, atol=atol(cfg), err_msg=f"PQMultiheadAttention bias={bias}: torch vs ONNX mismatch" ) -class _CausalSelfAttnModel(nn.Module): +class CausalSelfAttnModel(nn.Module): """Self-attention with a constant additive causal mask (the decoder-inference case).""" def __init__(self, mha: PQMultiheadAttention, seq_len: int): @@ -202,21 +202,21 @@ def forward(self, x): 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) + 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_fx(model, x, input_shape=(T, E), tmp_path=tmp_path) + 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_out, onnx_out, atol=_atol(cfg), err_msg=f"MHA causal attn_mask bias={bias}: torch vs ONNX mismatch" + torch_output, onnx_out, atol=atol(cfg), err_msg=f"MHA causal attn_mask bias={bias}: torch vs ONNX mismatch" ) -class _PaddedSelfAttnModel(nn.Module): +class PaddedSelfAttnModel(nn.Module): """Self-attention with a runtime bool key_padding_mask input (True == padding).""" def __init__(self, mha: PQMultiheadAttention): @@ -234,18 +234,18 @@ def test_mha_key_padding_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 = _PaddedSelfAttnModel(mha) + 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, key_padding_mask) - _apply_compression(model) + apply_compression(model) model.eval() with torch.no_grad(): - torch_out = model(x, key_padding_mask).cpu().numpy() + 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"]) @@ -258,11 +258,11 @@ def test_mha_key_padding_mask_onnx(cfg, bias, tmp_path): onnx_out = sess.run(None, {"x": x.cpu().numpy(), "key_padding_mask": key_padding_mask.cpu().numpy()})[0] np.testing.assert_allclose( - torch_out, onnx_out, atol=_atol(cfg), err_msg=f"MHA key_padding_mask bias={bias}: torch vs ONNX mismatch" + torch_output, onnx_out, atol=atol(cfg), err_msg=f"MHA key_padding_mask bias={bias}: torch vs ONNX mismatch" ) -class _TwoInputModel(nn.Module): +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): @@ -277,17 +277,17 @@ def forward(self, a, 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) + 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(a, b) # warm-up - _apply_compression(model) + apply_compression(model) model.eval() with torch.no_grad(): - torch_out = model(a, b).cpu().numpy() + 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) @@ -304,22 +304,22 @@ def test_two_input_onnx(cfg, bias, tmp_path): assert set(names) == {"a", "b"} onnx_out = sess.run(None, {"a": a.cpu().numpy(), "b": b.cpu().numpy()})[0] - np.testing.assert_allclose(torch_out, onnx_out, atol=ATOL, err_msg=f"two-input bias={bias}: 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") 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) + model = TwoInputModel(cfg, 16, 4, 8, bias=True) with torch.no_grad(): model(torch.randn(2, 16), torch.randn(2, 4)) - _apply_compression(model) + 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): +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): @@ -336,16 +336,16 @@ def forward(self, x, scale_up: bool = False): @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) + model = FlaggedModel(cfg, IN, OUT) x = torch.randn(3, IN) with torch.no_grad(): model(x, scale_up) - _apply_compression(model) + apply_compression(model) model.eval() with torch.no_grad(): - torch_out = model(x, scale_up).cpu().numpy() + 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}) @@ -359,11 +359,11 @@ def test_concrete_args_specialization(cfg, scale_up, tmp_path): onnx_out = sess.run(None, {"input": x.cpu().numpy()})[0] np.testing.assert_allclose( - torch_out, onnx_out, atol=ATOL, err_msg=f"concrete_args scale_up={scale_up}: torch vs ONNX mismatch" + torch_output, onnx_out, atol=ATOL, err_msg=f"concrete_args scale_up={scale_up}: torch vs ONNX mismatch" ) -class _ResidualConcatModel(nn.Module): +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): @@ -381,16 +381,16 @@ def forward(self, x): def test_residual_concat_onnx(cfg_quant, tmp_path): DIM, OUT = 16, 8 - model = _ResidualConcatModel(cfg_quant, DIM, OUT) + model = ResidualConcatModel(cfg_quant, DIM, OUT) x = torch.randn(4, DIM) with torch.no_grad(): model(x) - _apply_compression(model) + apply_compression(model) model.eval() with torch.no_grad(): - torch_out = model(x).cpu().numpy() + 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) @@ -400,7 +400,7 @@ def test_residual_concat_onnx(cfg_quant, tmp_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=QUANT_ATOL, err_msg="residual+concat: torch vs ONNX mismatch") + np.testing.assert_allclose(torch_output, onnx_out, atol=QUANT_ATOL, err_msg="residual+concat: torch vs ONNX mismatch") @pytest.mark.parametrize("activation", ["relu", "tanh", "hard_tanh", "leaky_relu", "gelu"]) @@ -411,11 +411,13 @@ def test_pqactivation_onnx(cfg_quant, activation, tmp_path): 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=(DIM,), tmp_path=tmp_path) - np.testing.assert_allclose(torch_out, onnx_out, atol=ATOL, err_msg=f"PQActivation {activation}: 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): @@ -436,7 +438,7 @@ def test_standalone_quantizer_onnx(cfg_quant, tmp_path): x = torch.randn(4, 16) with torch.no_grad(): model(x) - _apply_compression(model) + apply_compression(model) # A standalone quantizer must emit a Quantize/Dequantize pair. path = str(tmp_path / "quantizer.onnx") @@ -444,13 +446,13 @@ def test_standalone_quantizer_onnx(cfg_quant, tmp_path): op_types = [n.op_type for n in model_proto.graph.node] assert "QuantizeLinear" in op_types and "DequantizeLinear" in op_types - torch_out = _torch_out(model, x) + 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_out, onnx_out, atol=ATOL, err_msg="standalone Quantizer: torch vs ONNX mismatch") + np.testing.assert_allclose(torch_output, onnx_out, atol=ATOL, err_msg="standalone Quantizer: torch vs ONNX mismatch") -class _CNNFlattenModel(nn.Module): +class CNNFlattenModel(nn.Module): def __init__(self, cfg, in_c: int, hw: int, out: int, use_reshape: bool): super().__init__() self.conv = PQConv2d(cfg, in_channels=in_c, out_channels=4, kernel_size=3, padding=1) @@ -469,27 +471,27 @@ def forward(self, x): @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) + model = CNNFlattenModel(cfg_quant, IN_C, HW, OUT, use_reshape) x = torch.randn(2, IN_C, HW, HW) with torch.no_grad(): model(x) - _apply_compression(model) + apply_compression(model) model.eval() with torch.no_grad(): - torch_out = model(x).cpu().numpy() + 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=QUANT_ATOL, err_msg=f"CNN→Dense reshape={use_reshape}: torch vs ONNX mismatch" + torch_output, onnx_out, atol=QUANT_ATOL, err_msg=f"CNN→Dense reshape={use_reshape}: torch vs ONNX mismatch" ) -class _ScalarOpsModel(nn.Module): +class ScalarOpsModel(nn.Module): def __init__(self, cfg, dim: int): super().__init__() self.d = PQDense(cfg, in_features=dim, out_features=dim) @@ -504,16 +506,16 @@ def forward(self, x): def test_scalar_ops_onnx(cfg_quant, tmp_path): DIM = 16 - model = _ScalarOpsModel(cfg_quant, DIM) + model = ScalarOpsModel(cfg_quant, DIM) x = torch.randn(4, DIM) with torch.no_grad(): model(x) - _apply_compression(model) + apply_compression(model) model.eval() with torch.no_grad(): - torch_out = model(x).cpu().numpy() + 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) @@ -523,10 +525,10 @@ def test_scalar_ops_onnx(cfg_quant, tmp_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="scalar ops+sigmoid: torch vs ONNX mismatch") + np.testing.assert_allclose(torch_output, onnx_out, atol=ATOL, err_msg="scalar ops+sigmoid: torch vs ONNX mismatch") -class _MultiOutputModel(nn.Module): +class MultiOutputModel(nn.Module): def __init__(self, cfg, dim: int): super().__init__() self.a = PQDense(cfg, in_features=dim, out_features=8) @@ -538,12 +540,12 @@ def forward(self, x): def test_multi_output_onnx(cfg_quant, tmp_path): DIM = 16 - model = _MultiOutputModel(cfg_quant, DIM) + model = MultiOutputModel(cfg_quant, DIM) x = torch.randn(3, DIM) with torch.no_grad(): model(x) - _apply_compression(model) + apply_compression(model) model.eval() with torch.no_grad(): @@ -571,18 +573,18 @@ def test_pqlayernorm_onnx(cfg_quant, tmp_path): x = torch.randn(4, DIM) with torch.no_grad(): model(x) - _apply_compression(model) + apply_compression(model) model.eval() with torch.no_grad(): - torch_out = model(x).cpu().numpy() + 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_out, onnx_out, atol=ATOL, err_msg="PQLayerNorm: torch vs ONNX mismatch") + np.testing.assert_allclose(torch_output, onnx_out, atol=ATOL, err_msg="PQLayerNorm: torch vs ONNX mismatch") @pytest.mark.parametrize( @@ -600,37 +602,37 @@ def test_plain_passthrough_layers_onnx(make_model, input_shape, batch, tmp_path) model.eval() # Dropout/BatchNorm must be in eval mode for a deterministic compare x = torch.randn(batch, *input_shape) - torch_out = _torch_out(model, x) - onnx_out = _onnx_run(model, x, input_shape=input_shape, tmp_path=tmp_path) - np.testing.assert_allclose(torch_out, onnx_out, atol=ATOL, err_msg="plain passthrough layer: torch vs ONNX mismatch") + 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): +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) + 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_out = _torch_out(model, x) + 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_out, onnx_out, atol=ATOL, err_msg=f"integer_ops={integer_ops}: mismatch") + np.testing.assert_allclose(torch_output, onnx_out, atol=ATOL, err_msg=f"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) + 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) @@ -640,7 +642,7 @@ def test_include_clip_toggle_structure(cfg_quant, tmp_path): 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) + 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] @@ -652,7 +654,7 @@ 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) + 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)) @@ -684,11 +686,11 @@ def forward(self, x): x = torch.randn(4, 16) 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=(16,), tmp_path=tmp_path) - np.testing.assert_allclose(torch_out, onnx_out, atol=_atol(cfg), err_msg="tensor slicing: torch vs ONNX mismatch") + 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( @@ -716,9 +718,9 @@ def forward(self, x): x = torch.randn(4, 16) 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=(16,), tmp_path=tmp_path) - assert torch_out.shape == onnx_out.shape - np.testing.assert_allclose(torch_out, onnx_out, atol=_atol(cfg), err_msg="squeeze/unsqueeze: torch vs ONNX mismatch") + 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") From 4aaf8a25b2dfa62c69361d25e590656849897e7c Mon Sep 17 00:00:00 2001 From: Roope Niemi Date: Thu, 16 Jul 2026 10:57:30 +0200 Subject: [PATCH 7/8] add private function names --- src/pquant/__init__.py | 4 + src/pquant/core/keras/onnx/convert_to_onnx.py | 94 ++++----- src/pquant/core/keras/onnx/helpers.py | 8 +- src/pquant/core/keras/onnx/layer_builders.py | 58 +++--- src/pquant/core/torch/onnx/convert_to_onnx.py | 196 +++++++++--------- src/pquant/core/torch/onnx/layer_builders.py | 48 ++--- 6 files changed, 209 insertions(+), 199 deletions(-) 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/onnx/convert_to_onnx.py b/src/pquant/core/keras/onnx/convert_to_onnx.py index 9be8740..bcbeaf5 100644 --- a/src/pquant/core/keras/onnx/convert_to_onnx.py +++ b/src/pquant/core/keras/onnx/convert_to_onnx.py @@ -33,16 +33,16 @@ PQDepthwiseConv2d, PQMultiheadAttention, ) -from pquant.core.keras.onnx.helpers import keras_dtype_to_tp +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_depthwise_conv, - add_global_avgpool, - add_mha, - add_pq_activation, + _add_avgpool, + _add_batchnorm, + _add_conv, + _add_dense, + _add_depthwise_conv, + _add_global_avgpool, + _add_mha, + _add_pq_activation, ) from pquant.core.onnx_common import ( add_initializer, @@ -59,7 +59,7 @@ _ACTIVATION_OPS = {"relu": "Relu", "sigmoid": "Sigmoid", "tanh": "Tanh"} -def resolve_mask_arg(mask, prefix, kind, tensor_to_onnx, initializers): +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 @@ -68,12 +68,12 @@ def resolve_mask_arg(mask, prefix, kind, tensor_to_onnx, initializers): return add_initializer(initializers, f"{prefix}_{kind}_const", np.asarray(to_np(mask))) -def call_arguments(layer): +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 add_mha_layer(layer, prefix, input_onnx_names, nodes, initializers, quant_fn, use_qonnx, store_int, tensor_to_onnx): +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: @@ -81,11 +81,11 @@ def add_mha_layer(layer, prefix, input_onnx_names, nodes, initializers, quant_fn else: q_in = k_in = v_in = input_onnx_names[0] - arguments = call_arguments(layer) + 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( + 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, @@ -101,24 +101,24 @@ def add_mha_layer(layer, prefix, input_onnx_names, nodes, initializers, quant_fn ) -def add_getitem_op(layer, prefix, current, nodes, initializers): +def _add_getitem_op(layer, prefix, current, nodes, initializers): """keras.ops GetItem operation recorded by ``x[...]`` KerasTensor syntax.""" - arguments = call_arguments(layer) + 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): +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) + 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): +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 + 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] @@ -128,7 +128,7 @@ def add_squeeze_op(layer, prefix, current, nodes, initializers): return emit_squeeze(prefix, current, axes, nodes, initializers) -def add_standard_activation(layer, prefix, current, nodes): +def _add_standard_activation(layer, prefix, current, nodes): """keras.layers.ReLU or keras.layers.Activation with a supported activation.""" activation = ( layer.activation.__name__ @@ -144,7 +144,7 @@ def add_standard_activation(layer, prefix, current, nodes): return out -def emit_layer( +def _emit_layer( layer, prefix, current, @@ -159,22 +159,22 @@ def emit_layer( """Emit ONNX nodes for a single Keras layer. Returns the ONNX output name.""" if isinstance(layer, PQMultiheadAttention): - return add_mha_layer( + 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) + 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) + 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) + 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( + return _add_conv( layer, prefix, current, @@ -187,19 +187,19 @@ def emit_layer( ) if isinstance(layer, PQBatchNormalization): - return add_batchnorm(layer, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights) + 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) + return _add_getitem_op(layer, prefix, current, nodes, initializers) if type(layer).__name__ == "ExpandDims": - return add_expand_dims_op(layer, prefix, current, nodes, initializers) + return _add_expand_dims_op(layer, prefix, current, nodes, initializers) if type(layer).__name__ == "Squeeze": - return add_squeeze_op(layer, prefix, current, nodes, initializers) + 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) + return _add_standard_activation(layer, prefix, current, nodes) if isinstance(layer, keras.layers.Flatten): out = f"{prefix}_flatten" @@ -234,11 +234,11 @@ def emit_layer( 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) + 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) + return _add_global_avgpool(layer, prefix, current, nodes, ndim=ndim) if isinstance(layer, keras.layers.Dropout): return current # identity at inference @@ -246,18 +246,18 @@ def emit_layer( raise TypeError(f"Unsupported Keras layer type for ONNX export: {type(layer).__name__!r}") -def build_tensor_onnx_map(model): +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))} + return {id(inp): name for inp, name in zip(model.inputs, _model_input_names(model))} -def 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): +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 [] @@ -275,7 +275,7 @@ def inbound_input_names(layer, tensor_to_onnx): return result -def register_layer_output(layer, onnx_name, tensor_to_onnx): +def _register_layer_output(layer, onnx_name, tensor_to_onnx): if not layer._inbound_nodes: return out_tensors = layer._inbound_nodes[0].output_tensors @@ -334,19 +334,19 @@ def convert_to_onnx( onnx_nodes: list[onnx.NodeProto] = [] initializers: list[onnx.TensorProto] = [] - tensor_to_onnx = build_tensor_onnx_map(model) + 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) + 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( + output_name = _emit_layer( layer, prefix, input_onnx_names[0], @@ -359,16 +359,16 @@ def convert_to_onnx( tensor_to_onnx=tensor_to_onnx, ) - register_layer_output(layer, output_name, 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) + 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] + 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) diff --git a/src/pquant/core/keras/onnx/helpers.py b/src/pquant/core/keras/onnx/helpers.py index 0541608..9f43c0f 100644 --- a/src/pquant/core/keras/onnx/helpers.py +++ b/src/pquant/core/keras/onnx/helpers.py @@ -7,7 +7,7 @@ from onnx import TensorProto -def keras_dtype_to_tp(dtype): +def _keras_dtype_to_tp(dtype): """Map a Keras/numpy dtype string to an ONNX TensorProto dtype (default float32).""" return { "float32": TensorProto.FLOAT, @@ -19,18 +19,18 @@ def keras_dtype_to_tp(dtype): }.get(str(dtype), TensorProto.FLOAT) -def channels_last(layer): +def _channels_last(layer): return getattr(layer, "data_format", keras.config.image_data_format()) == "channels_last" -def nchw_perms(ndim): +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): +def _bn_transpose_info(layer): """ Return (need_transpose, perm_fwd, perm_bwd) for a BatchNormalization layer. diff --git a/src/pquant/core/keras/onnx/layer_builders.py b/src/pquant/core/keras/onnx/layer_builders.py index 806118e..5e07932 100644 --- a/src/pquant/core/keras/onnx/layer_builders.py +++ b/src/pquant/core/keras/onnx/layer_builders.py @@ -4,7 +4,11 @@ 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.keras.onnx.helpers import ( + _bn_transpose_info, + _channels_last, + _nchw_perms, +) from pquant.core.onnx_common import ( add_float_scalar, add_initializer, @@ -19,7 +23,7 @@ ) -def add_dense(layer, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights): +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) @@ -48,7 +52,7 @@ def add_dense(layer, prefix, current, nodes, initializers, quant_fn, use_qonnx, 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): +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) @@ -79,7 +83,7 @@ def add_dense_nd(layer, prefix, current, nodes, initializers, quant_fn, use_qonn return maybe_quant_output(layer, prefix, current, nodes, initializers, quant_fn) -def add_conv_node(layer, prefix, conv_inputs, groups, ndim, nodes): +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( @@ -97,11 +101,11 @@ def add_conv_node(layer, prefix, conv_inputs, groups, ndim, nodes): return conv_out -def add_conv_common(layer, prefix, current, kernel_onnx, groups, ndim, nodes, initializers, quant_fn, use_qonnx, store_int): +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) + is_channels_last = _channels_last(layer) if is_channels_last: - perm_to_nchw, perm_to_nhwx = nchw_perms(ndim) + 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) @@ -113,7 +117,7 @@ def add_conv_common(layer, prefix, current, kernel_onnx, groups, ndim, nodes, in 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 = _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: @@ -121,7 +125,7 @@ def add_conv_common(layer, prefix, current, kernel_onnx, groups, ndim, nodes, in return current -def add_conv(layer, prefix, current, nodes, initializers, ndim, quant_fn, use_qonnx, store_integer_weights): +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: @@ -129,25 +133,25 @@ def add_conv(layer, prefix, current, nodes, initializers, ndim, quant_fn, use_qo 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( + 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): +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( + 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): +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) + 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) @@ -185,7 +189,7 @@ def add_batchnorm(layer, prefix, current, nodes, initializers, quant_fn, use_qon return current -def add_mha( +def _add_mha( layer, prefix, q_input, @@ -200,13 +204,13 @@ def add_mha( attn_mask=None, ): # Q / K / V projections: (B, L, E) → (B, L, E) - q_proj_out = add_dense_nd( + 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( + 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( + v_proj_out = _add_dense_nd( layer.v_proj, f"{prefix}_v_proj", v_input, nodes, initializers, quant_fn, use_qonnx, store_integer_weights ) @@ -215,16 +219,16 @@ def add_mha( ) # Output projection: (B, T, E) → (B, T, E) - out = add_dense_nd( + 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) +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) + 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) @@ -248,10 +252,10 @@ def add_avgpool(layer, prefix, current, nodes, initializers, ndim, quant_fn): return current -def add_global_avgpool(layer, prefix, current, nodes, ndim): - is_channels_last = channels_last(layer) +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) + perm_to_nchw, _ = _nchw_perms(ndim) current = add_transpose(f"{prefix}_pre", current, perm_to_nchw, nodes) pool_out = f"{prefix}_global_pool" @@ -265,7 +269,7 @@ def add_global_avgpool(layer, prefix, current, nodes, ndim): return current -def add_pq_activation(layer, prefix, current, nodes, initializers, quant_fn): +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"): diff --git a/src/pquant/core/torch/onnx/convert_to_onnx.py b/src/pquant/core/torch/onnx/convert_to_onnx.py index d777e59..688c2c8 100644 --- a/src/pquant/core/torch/onnx/convert_to_onnx.py +++ b/src/pquant/core/torch/onnx/convert_to_onnx.py @@ -48,28 +48,30 @@ PQMultiheadAttention, ) from pquant.core.torch.onnx.layer_builders import ( # noqa: E402 - add_activation, - add_avgpool, - add_batchnorm, - add_conv, - add_dense, - add_layernorm, - add_maxpool, - add_mha, - add_upsample, + _add_activation, + _add_avgpool, + _add_batchnorm, + _add_conv, + _add_dense, + _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): +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( + 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( + return _add_conv( module, prefix, current, @@ -81,14 +83,14 @@ def emit_module(module, prefix, current, nodes, initializers, quant_fn, use_qonn 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) + 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) + 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) + 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) + 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): @@ -104,9 +106,9 @@ def emit_module(module, prefix, current, nodes, initializers, quant_fn, use_qonn 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) + return _add_maxpool(module, prefix, current, nodes) if isinstance(module, nn.Upsample): - return add_upsample(module, prefix, current, nodes, initializers) + 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__}") @@ -131,14 +133,14 @@ 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]: +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): +def _normalize_input_dtypes(input_dtypes, n: int): torch_map = { "float32": torch.float32, "float": torch.float32, @@ -173,14 +175,14 @@ def normalize_input_dtypes(input_dtypes, n: int): return torch_dtypes, tp_dtypes -def swap_perm(rank: int, d0: int, d1: int) -> list[int]: +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]: +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] @@ -228,28 +230,28 @@ def __init__(self, gm, ph_to_name, quant_fn, use_qonnx, store_integer_weights, i self.node_to_name: dict[fx.Node, str] = {} self.output_names: list[str] = [] - def run(self) -> 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) + self._emit_get_attr(node) elif node.op == "call_module": - self.emit_call_module(node) + self._emit_call_module(node) elif node.op == "call_function": - self.emit_call_function(node) + self._emit_call_function(node) elif node.op == "call_method": - self.emit_call_method(node) + self._emit_call_method(node) elif node.op == "output": - self.collect_outputs(node) + self._collect_outputs(node) return self.output_names - def name_of(self, arg) -> str: + 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]: + 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): @@ -260,16 +262,16 @@ def binop_inputs(self, node: fx.Node) -> list[str]: raise TypeError(f"FX export: unsupported binary-op arg type {type(arg).__name__}") return names - def node_shape(self, node: fx.Node) -> tuple: + 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 _node_rank(self, node: fx.Node) -> int: + return len(self._node_shape(node)) - def emit_get_attr(self, node): + def _emit_get_attr(self, node): obj = self.gm for part in node.target.split("."): obj = getattr(obj, part) @@ -277,16 +279,16 @@ def emit_get_attr(self, node): add_initializer(self.initializers, node.name, obj.detach().cpu().numpy()) self.node_to_name[node] = node.name - def emit_call_module(self, node): + 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) + self.node_to_name[node] = self._emit_mha_module(module, node, prefix) return - self.node_to_name[node] = emit_module( + self.node_to_name[node] = _emit_module( module, prefix, - self.name_of(node.args[0]), + self._name_of(node.args[0]), self.nodes, self.initializers, self.quant_fn, @@ -295,14 +297,14 @@ def emit_call_module(self, node): self.integer_ops, ) - def emit_mha_module(self, module, node, prefix) -> tuple: + 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( + 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, @@ -317,7 +319,7 @@ def emit_mha_module(self, module, node, prefix) -> tuple: attn_mask=attn_mask_name, ) - def mask_name(self, node, pos, kw): + 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 @@ -325,55 +327,55 @@ def mask_name(self, node, pos, kw): 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): + 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) + self._emit_getitem(node) elif fn in self._BINARY_OPS: - self.add_simple_node(node, self._BINARY_OPS[fn], self.binop_inputs(node)) + 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])]) + self._add_simple_node(node, self._UNARY_OPS[fn], [self._name_of(node.args[0])]) elif fn is torch.transpose: - self.emit_transpose(node) + self._emit_transpose(node) elif fn is torch.permute: - self.emit_permute(node) + self._emit_permute(node) elif fn is torch.cat: - tensors = [self.name_of(a) for a in node.args[0]] + 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)) + self._add_simple_node(node, "Concat", tensors, axis=int(dim)) elif fn is torch.flatten: - self.emit_flatten(node, default_start_dim=0) + self._emit_flatten(node, default_start_dim=0) elif fn is torch.squeeze: - self.emit_squeeze(node) + self._emit_squeeze(node) elif fn is torch.unsqueeze: - self.emit_unsqueeze(node) + self._emit_unsqueeze(node) else: raise TypeError(f"Unsupported call_function for FX ONNX export: {fn}") - def emit_call_method(self, node): + def _emit_call_method(self, node): method = node.target if method == "relu": - self.add_simple_node(node, "Relu", [self.name_of(node.args[0])]) + self._add_simple_node(node, "Relu", [self._name_of(node.args[0])]) elif method == "flatten": - self.emit_flatten(node, default_start_dim=1) + self._emit_flatten(node, default_start_dim=1) elif method in ("view", "reshape"): - self.emit_reshape(node) + self._emit_reshape(node) elif method == "transpose": - self.emit_transpose(node) + self._emit_transpose(node) elif method == "permute": - self.emit_permute(node) + self._emit_permute(node) elif method == "matmul": - self.add_simple_node(node, "MatMul", self.binop_inputs(node)) + self._add_simple_node(node, "MatMul", self._binop_inputs(node)) elif method == "squeeze": - self.emit_squeeze(node) + self._emit_squeeze(node) elif method == "unsqueeze": - self.emit_unsqueeze(node) + self._emit_unsqueeze(node) else: raise TypeError(f"Unsupported call_method for FX ONNX export: {node.target!r}") - def collect_outputs(self, node): + def _collect_outputs(self, node): ret = node.args[0] rets = list(ret) if isinstance(ret, (tuple, list)) else [ret] for r in rets: @@ -383,63 +385,63 @@ def collect_outputs(self, node): # 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): + 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): + 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]) + 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): + 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) + 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_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): + 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)) + self._add_simple_node(node, "Flatten", [self._name_of(node.args[0])], axis=int(start_dim)) - def emit_reshape(self, node): + 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]) + 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 _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]: + def _squeeze_axes(self, node) -> list[int]: """Resolve the axes a torch squeeze()/.squeeze() call removes.""" - in_shape = self.node_shape(node.args[0]) + 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): + 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) + 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): +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))): @@ -453,7 +455,7 @@ def prune_untranslatable_nodes(gm): gm.recompile() -def graph_input_names(gm, n_expected: int) -> dict: +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: @@ -468,7 +470,7 @@ def graph_input_names(gm, n_expected: int) -> dict: return dict(zip(placeholders, names)) -def route_input_passthrough_outputs(output_names, input_names, nodes): +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): @@ -550,12 +552,12 @@ def convert_to_onnx( 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)) + 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)) + _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) @@ -564,8 +566,8 @@ def convert_to_onnx( 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) + 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 {})) diff --git a/src/pquant/core/torch/onnx/layer_builders.py b/src/pquant/core/torch/onnx/layer_builders.py index 868d0cd..6094d04 100644 --- a/src/pquant/core/torch/onnx/layer_builders.py +++ b/src/pquant/core/torch/onnx/layer_builders.py @@ -20,7 +20,7 @@ ) -def quantize_input_to_int(quantizer, prefix, current, nodes, initializers): +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). @@ -44,7 +44,7 @@ def quantize_input_to_int(quantizer, prefix, current, nodes, initializers): return int_name, zp_name, scale -def integer_weights_transposed(module, prefix, initializers): +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. @@ -73,7 +73,7 @@ def integer_weights_transposed(module, prefix, initializers): return weight_name, zp_name, scale_1d, per_channel -def dequantize_accumulator(prefix, current, combined_scale_1d, per_channel, nodes, initializers): +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 the output tensor is [batch, out] and out is axis 1. @@ -94,13 +94,13 @@ def dequantize_accumulator(prefix, current, combined_scale_1d, per_channel, node return out -def add_dense_integer(module, prefix, current, nodes, initializers): +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) + 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] @@ -114,13 +114,13 @@ def add_dense_integer(module, prefix, current, nodes, initializers): 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) + 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): +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) @@ -151,9 +151,9 @@ def add_dense_nd(module, prefix, current, nodes, initializers, quant_fn, use_qon 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): +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) + return _add_dense_integer(module, prefix, current, nodes, initializers) current = maybe_quant_input(module, prefix, current, nodes, initializers, quant_fn) q_weight = emit_param( @@ -188,7 +188,7 @@ def add_dense(module, prefix, current, nodes, initializers, quant_fn, use_qonnx, 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): +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( @@ -234,7 +234,7 @@ def add_conv(module, prefix, current, nodes, initializers, ndim, quant_fn, use_q 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): +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( @@ -265,7 +265,7 @@ def add_batchnorm(module, prefix, current, nodes, initializers, quant_fn, use_qo return bn_out -def add_layernorm(module, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights): +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)) @@ -296,7 +296,7 @@ def add_layernorm(module, prefix, current, nodes, initializers, quant_fn, use_qo return maybe_quant_output(module, prefix, ln_out, nodes, initializers, quant_fn) -def add_avgpool(module, prefix, current, nodes, initializers, ndim, 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" @@ -315,7 +315,7 @@ def add_avgpool(module, prefix, current, nodes, initializers, ndim, quant_fn): return maybe_quant_output(module, prefix, pool_out, nodes, initializers, quant_fn) -def add_maxpool(module, prefix, current, nodes): +def _add_maxpool(module, prefix, current, nodes): out = f"{prefix}_maxpool" nodes.append( oh.make_node( @@ -330,7 +330,7 @@ def add_maxpool(module, prefix, current, nodes): return out -def add_upsample(module, prefix, current, nodes, initializers): +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 @@ -353,7 +353,7 @@ def add_upsample(module, prefix, current, nodes, initializers): return out -def add_activation(module, prefix, current, nodes, initializers, quant_fn): +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) @@ -371,14 +371,14 @@ def add_activation(module, prefix, current, nodes, initializers, quant_fn): 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) + _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): +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) @@ -421,7 +421,7 @@ def add_gelu(module, prefix, current, act_out, nodes, initializers): ] -def add_mha( +def _add_mha( module, prefix, q_input, @@ -441,20 +441,20 @@ def add_mha( 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( + 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( + 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( + 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( + out = _add_dense_nd( module.out_proj, f"{prefix}_out_proj", context, nodes, initializers, quant_fn, use_qonnx, store_integer_weights ) From 10d1759d4735d470d9db85e1852dd8a49e9ece3c Mon Sep 17 00:00:00 2001 From: Roope Niemi Date: Tue, 21 Jul 2026 16:23:22 +0200 Subject: [PATCH 8/8] use the add_dense_nd based on rank --- src/pquant/core/keras/onnx/convert_to_onnx.py | 12 +++++- src/pquant/core/torch/onnx/convert_to_onnx.py | 17 ++++++++- src/pquant/core/torch/onnx/layer_builders.py | 5 ++- tests/test_keras_onnx_converter.py | 37 +++++++++++++++++++ tests/test_torch_onnx_converter.py | 33 +++++++++++++++++ 5 files changed, 100 insertions(+), 4 deletions(-) diff --git a/src/pquant/core/keras/onnx/convert_to_onnx.py b/src/pquant/core/keras/onnx/convert_to_onnx.py index bcbeaf5..c0b7a6f 100644 --- a/src/pquant/core/keras/onnx/convert_to_onnx.py +++ b/src/pquant/core/keras/onnx/convert_to_onnx.py @@ -39,6 +39,7 @@ _add_batchnorm, _add_conv, _add_dense, + _add_dense_nd, _add_depthwise_conv, _add_global_avgpool, _add_mha, @@ -73,6 +74,13 @@ def _call_arguments(layer): 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] @@ -167,7 +175,9 @@ def _emit_layer( 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) + # 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) diff --git a/src/pquant/core/torch/onnx/convert_to_onnx.py b/src/pquant/core/torch/onnx/convert_to_onnx.py index 688c2c8..3d06a68 100644 --- a/src/pquant/core/torch/onnx/convert_to_onnx.py +++ b/src/pquant/core/torch/onnx/convert_to_onnx.py @@ -53,6 +53,7 @@ _add_batchnorm, _add_conv, _add_dense, + _add_dense_nd, _add_layernorm, _add_maxpool, _add_mha, @@ -62,10 +63,23 @@ def _emit_module( - module, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights, integer_ops=False + 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 ) @@ -295,6 +309,7 @@ def _emit_call_module(self, node): 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: diff --git a/src/pquant/core/torch/onnx/layer_builders.py b/src/pquant/core/torch/onnx/layer_builders.py index 6094d04..fb33b7a 100644 --- a/src/pquant/core/torch/onnx/layer_builders.py +++ b/src/pquant/core/torch/onnx/layer_builders.py @@ -76,12 +76,13 @@ def _integer_weights_transposed(module, prefix, initializers): 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 the output tensor is [batch, out] and out is axis 1. + 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} + dql_kwargs = {"axis": -1} else: scale_np = np.array(float(combined_scale_1d[0]), dtype=np.float32) zp_np = np.array(np.int32(0)) diff --git a/tests/test_keras_onnx_converter.py b/tests/test_keras_onnx_converter.py index 0f6be41..01a562e 100644 --- a/tests/test_keras_onnx_converter.py +++ b/tests/test_keras_onnx_converter.py @@ -293,3 +293,40 @@ def test_squeeze_unsqueeze_onnx(cfg, reshaper, tmp_path): 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 445e189..4927915 100644 --- a/tests/test_torch_onnx_converter.py +++ b/tests/test_torch_onnx_converter.py @@ -630,6 +630,39 @@ def test_integer_weight_storage_onnx(cfg_quant, integer_ops, tmp_path): 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)