22#
33# SPDX-License-Identifier: Apache-2.0
44
5+ from __future__ import annotations
6+
57import io
68import os
79import unittest
@@ -52,7 +54,7 @@ def _compare(
5254 model_opt : onnx .ModelProto ,
5355 model_ori : onnx .ModelProto ,
5456 n_times : int = 5 ,
55- input_shapes : Optional [ TensorShapes ] = None ,
57+ input_shapes : TensorShapes | None = None ,
5658 verbose = True ,
5759 ) -> bool :
5860 """
@@ -65,7 +67,7 @@ def _compare(
6567 def get_shape_from_value_info_proto (v : onnx .ValueInfoProto ) -> list [int ]:
6668 return [dim .dim_value for dim in v .type .tensor_type .shape .dim ]
6769
68- def get_value_info_all (m : onnx .ModelProto , name : str ) -> Optional [ onnx .ValueInfoProto ] :
70+ def get_value_info_all (m : onnx .ModelProto , name : str ) -> onnx .ValueInfoProto | None :
6971 for v in m .graph .value_info :
7072 if v .name == name :
7173 return v
@@ -89,7 +91,7 @@ def get_shape(m: onnx.ModelProto, name: str) -> TensorShape:
8991 return get_shape_from_value_info_proto (v )
9092 raise RuntimeError (f'Cannot get shape of "{ name } "' )
9193
92- def get_elem_type (m : onnx .ModelProto , name : str ) -> Optional [ int ] :
94+ def get_elem_type (m : onnx .ModelProto , name : str ) -> int | None :
9395 v = get_value_info_all (m , name )
9496 if v is not None :
9597 return v .type .tensor_type .elem_type
@@ -127,7 +129,7 @@ def get_input_names(model: onnx.ModelProto) -> list[str]:
127129 )
128130 return input_names
129131
130- def generate_rand_input (model , input_shapes : Optional [ TensorShapes ] = None ):
132+ def generate_rand_input (model , input_shapes : TensorShapes | None = None ):
131133 if input_shapes is None :
132134 input_shapes = {}
133135 input_names = get_input_names (model )
@@ -151,7 +153,7 @@ def generate_rand_input(model, input_shapes: Optional[TensorShapes] = None):
151153 return inputs
152154
153155 def forward (
154- model , inputs = None , input_shapes : Optional [ TensorShapes ] = None
156+ model , inputs = None , input_shapes : TensorShapes | None = None
155157 ) -> dict [str , np .ndarray ]:
156158 if input_shapes is None :
157159 input_shapes = {}
@@ -1571,38 +1573,123 @@ def test_fuse_add_bias_into_conv_with_non_constant_bias(self):
15711573
15721574 # type: () -> None
15731575 def test_fuse_add_bias_into_conv_with_quanted_bias (self ):
1574- nodes = [helper .make_node ("Conv" , ["X" , "Y" ], ["Z" ]),
1575- helper .make_node ("QuantizeLinear" , ["A" , "scale" , "zero_point" ], ["B" ], axis = 0 ),
1576- helper .make_node ("DequantizeLinear" , ["B" , "scale" , "zero_point" ], ["C" ], axis = 0 ),
1577- helper .make_node ("Add" , ["Z" , "C" ], ["D" ])]
1576+ nodes = [
1577+ helper .make_node ("Conv" , ["X" , "Y" ], ["Z" ]),
1578+ helper .make_node ("QuantizeLinear" , ["A" , "scale" , "zero_point" ], ["B" ], axis = 0 ),
1579+ helper .make_node ("DequantizeLinear" , ["B" , "scale" , "zero_point" ], ["C" ], axis = 0 ),
1580+ helper .make_node ("Add" , ["Z" , "C" ], ["D" ]),
1581+ ]
15781582 graph = helper .make_graph (
15791583 nodes ,
15801584 "test" ,
1581- [helper .make_tensor_value_info ("X" , TensorProto .FLOAT , (1 , 5 , 3 , 3 )),
1582- helper .make_tensor_value_info (
1583- "Y" , TensorProto .FLOAT , (16 , 5 , 3 , 3 )),
1584- helper .make_tensor_value_info ("A" , TensorProto .FLOAT , (16 , 1 , 1 ))],
1585- [helper .make_tensor_value_info (
1586- "D" , TensorProto .FLOAT , (1 , 16 , 1 , 1 ))],
1587- [helper .make_tensor ("scale" , TensorProto .FLOAT ,
1588- dims = (16 ,),
1589- vals = np .random .rand (16 ).astype (np .float32 ).tobytes (),
1590- raw = True ),
1591- helper .make_tensor ("zero_point" , TensorProto .INT8 ,
1592- dims = (16 ,),
1593- vals = np .zeros ([16 ]).astype (np .int8 ).tobytes (),
1594- raw = True )],
1595- value_info = [helper .make_tensor_value_info (
1596- "C" , TensorProto .FLOAT , (16 , 1 , 1 ))]
1585+ [
1586+ helper .make_tensor_value_info ("X" , TensorProto .FLOAT , (1 , 5 , 3 , 3 )),
1587+ helper .make_tensor_value_info ("Y" , TensorProto .FLOAT , (16 , 5 , 3 , 3 )),
1588+ helper .make_tensor_value_info ("A" , TensorProto .FLOAT , (16 , 1 , 1 )),
1589+ ],
1590+ [helper .make_tensor_value_info ("D" , TensorProto .FLOAT , (1 , 16 , 1 , 1 ))],
1591+ [
1592+ helper .make_tensor (
1593+ "scale" ,
1594+ TensorProto .FLOAT ,
1595+ dims = (16 ,),
1596+ vals = np .random .rand (16 ).astype (np .float32 ).tobytes (),
1597+ raw = True ,
1598+ ),
1599+ helper .make_tensor (
1600+ "zero_point" ,
1601+ TensorProto .INT8 ,
1602+ dims = (16 ,),
1603+ vals = np .zeros ([16 ]).astype (np .int8 ).tobytes (),
1604+ raw = True ,
1605+ ),
1606+ ],
1607+ value_info = [helper .make_tensor_value_info ("C" , TensorProto .FLOAT , (16 , 1 , 1 ))],
1608+ )
1609+ optimized_model = self ._optimized (
1610+ graph , ["fuse_add_bias_into_conv" ], opset_imports = [helper .make_opsetid ("" , 13 )]
15971611 )
1598- optimized_model = self ._optimized (graph , ["fuse_add_bias_into_conv" ], opset_imports = [helper .make_opsetid ("" , 13 )])
15991612
16001613 assert len (list (optimized_model .graph .node )) == 4
1601- assert optimized_model .graph .node [0 ].op_type == 'Squeeze'
1602- assert optimized_model .graph .node [1 ].op_type == 'QuantizeLinear'
1603- assert optimized_model .graph .node [2 ].op_type == 'DequantizeLinear'
1604- assert optimized_model .graph .node [3 ].op_type == 'Conv'
1605- assert optimized_model .graph .output [0 ].name == 'D'
1614+ assert optimized_model .graph .node [0 ].op_type == "QuantizeLinear"
1615+ assert optimized_model .graph .node [1 ].op_type == "DequantizeLinear"
1616+ assert optimized_model .graph .node [2 ].op_type == "Squeeze"
1617+ assert optimized_model .graph .node [3 ].op_type == "Conv"
1618+ assert optimized_model .graph .output [0 ].name == "D"
1619+
1620+ def test_fuse_add_bias_into_conv_dependency_chain (self ) -> None :
1621+ for bias_shape in [[1 ], [2 , 1 , 1 ]]:
1622+ with self .subTest (bias_shape = bias_shape ):
1623+ graph = helper .make_graph (
1624+ [
1625+ helper .make_node ("Conv" , ["X" , "W" ], ["conv" ]),
1626+ helper .make_node ("Sin" , ["A" ], ["sin" ]),
1627+ helper .make_node ("Cos" , ["sin" ], ["bias" ]),
1628+ helper .make_node ("Add" , ["conv" , "bias" ], ["Y" ]),
1629+ ],
1630+ "bias_dependency_chain" ,
1631+ [
1632+ helper .make_tensor_value_info ("X" , TensorProto .FLOAT , [1 , 1 , 2 , 2 ]),
1633+ helper .make_tensor_value_info ("A" , TensorProto .FLOAT , bias_shape ),
1634+ ],
1635+ [helper .make_tensor_value_info ("Y" , TensorProto .FLOAT , [1 , 2 , 2 , 2 ])],
1636+ initializer = [
1637+ numpy_helper .from_array (np .ones ([2 , 1 , 1 , 1 ], np .float32 ), "W" )
1638+ ],
1639+ value_info = [
1640+ helper .make_tensor_value_info ("bias" , TensorProto .FLOAT , bias_shape )
1641+ ],
1642+ )
1643+ optimized = self ._optimized (graph , ["fuse_add_bias_into_conv" ])
1644+ assert all (n .op_type != "Add" for n in optimized .graph .node )
1645+ assert optimized .graph .node [- 1 ].op_type == "Conv"
1646+
1647+ def test_fuse_add_bias_into_conv_per_channel_qdq (self ) -> None :
1648+ for axis in [1 , - 3 ]:
1649+ with self .subTest (axis = axis ):
1650+ graph = helper .make_graph (
1651+ [
1652+ helper .make_node ("Conv" , ["X" , "W" ], ["conv" ]),
1653+ helper .make_node (
1654+ "QuantizeLinear" , ["A" , "scale" , "zero" ], ["quantized" ], axis = axis
1655+ ),
1656+ helper .make_node (
1657+ "DequantizeLinear" ,
1658+ ["quantized" , "scale" , "zero" ],
1659+ ["bias" ],
1660+ axis = axis ,
1661+ ),
1662+ helper .make_node ("Add" , ["conv" , "bias" ], ["Y" ]),
1663+ ],
1664+ "per_channel_bias" ,
1665+ [
1666+ helper .make_tensor_value_info ("X" , TensorProto .FLOAT , [1 , 1 , 2 , 2 ]),
1667+ helper .make_tensor_value_info ("A" , TensorProto .FLOAT , [1 , 2 , 1 , 1 ]),
1668+ ],
1669+ [
1670+ helper .make_tensor_value_info ("Y" , TensorProto .FLOAT , [1 , 2 , 2 , 2 ]),
1671+ helper .make_tensor_value_info (
1672+ "quantized" , TensorProto .INT8 , [1 , 2 , 1 , 1 ]
1673+ ),
1674+ ],
1675+ initializer = [
1676+ numpy_helper .from_array (np .ones ([2 , 1 , 1 , 1 ], np .float32 ), "W" ),
1677+ numpy_helper .from_array (np .array ([0.1 , 0.2 ], np .float32 ), "scale" ),
1678+ numpy_helper .from_array (np .zeros (2 , np .int8 ), "zero" ),
1679+ ],
1680+ value_info = [
1681+ helper .make_tensor_value_info (
1682+ "bias" , TensorProto .FLOAT , [1 , 2 , 1 , 1 ]
1683+ )
1684+ ],
1685+ )
1686+ optimized = self ._optimized (graph , ["fuse_add_bias_into_conv" ])
1687+ assert [n .op_type for n in optimized .graph .node ] == [
1688+ "QuantizeLinear" , "DequantizeLinear" , "Squeeze" , "Conv"
1689+ ]
1690+ assert next (
1691+ a .i for a in optimized .graph .node [0 ].attribute if a .name == "axis"
1692+ ) == axis
16061693
16071694 def test_fuse_matmul_add_bias_into_gemm (self ): # type: () -> None
16081695 matmul = helper .make_node ("MatMul" , ["X" , "Y" ], ["Z" ])
0 commit comments