@@ -63,6 +63,8 @@ def apply(self, model):
63
63
stride_w ,
64
64
group ,
65
65
weight_name ,
66
+ conv_weight_inp_name ,
67
+ conv_weight_q_scale_name ,
66
68
W_conv ,
67
69
ifm_ch ,
68
70
ofm_ch ,
@@ -74,12 +76,18 @@ def apply(self, model):
74
76
pad ,
75
77
) = self .extract_conv_params (model , node )
76
78
79
+ if W_conv is None :
80
+ warnings .warn ("Found Conv node with non-initialized weight, skipping" )
81
+ continue
82
+
77
83
# if depthwise conv create sparse matrix and variable "dw"
78
84
# to store as attribute in Im2Col that indicates that the created
79
85
# Im2Col node belongs to a depthwise convolution
80
86
dw = False
81
87
if group == ifm_ch and ofm_ch == ifm_ch :
82
88
W_sparse = np .zeros ((ofm_ch , ifm_ch , k_h , k_w )) # (OFM, IFM, k_H, k_W)
89
+ # TODO: if the convolution is quantized with a non-zero zeropoint we
90
+ # should be using the zeropoint value here instead of np.zeros
83
91
for ch in range (ifm_ch ):
84
92
W_sparse [ch ][ch ] = W_conv [ch ][0 ] # W_conv = [OFM, IFM, k_H, k_W]
85
93
W_conv = W_sparse .astype (np .float32 )
@@ -104,6 +112,21 @@ def apply(self, model):
104
112
# transpose to get ONNX-compatible [k_h*k_w*IFM][OFM] matrix
105
113
W_matmul = W_matmul .T
106
114
model .set_initializer (weight_name , W_matmul )
115
+ if weight_name != conv_weight_inp_name :
116
+ # required for convs with quantized weights
117
+ model .set_tensor_shape (conv_weight_inp_name , W_matmul .shape )
118
+ if conv_weight_q_scale_name is not None :
119
+ # required for convs with quantized weights
120
+ scale_weight_q = model .get_initializer (conv_weight_q_scale_name )
121
+ if scale_weight_q .ndim > 0 :
122
+ # scale shape is originally [OFM, IFM, k_H, k_W]
123
+ # transpose into [OFM, k_H, k_W, IFM]
124
+ scale_weight_q = scale_weight_q .transpose (0 , 2 , 3 , 1 )
125
+ # reshape into [OFM][k_h*k_w*IFM] matrix
126
+ scale_weight_q = scale_weight_q .reshape (ofm_ch , - 1 )
127
+ # transpose to be shape-compatible with weight matrix
128
+ scale_weight_q = scale_weight_q .T
129
+ model .set_initializer (conv_weight_q_scale_name , scale_weight_q )
107
130
108
131
# create new intermediate values
109
132
inp_trans_out = helper .make_tensor_value_info (
@@ -154,7 +177,7 @@ def apply(self, model):
154
177
155
178
matmul_input = im2col_out if need_im2col else inp_trans_out
156
179
# do matmul
157
- matmul_node = helper .make_node ("MatMul" , [matmul_input , weight_name ], [matmul_out ])
180
+ matmul_node = helper .make_node ("MatMul" , [matmul_input , conv_weight_inp_name ], [matmul_out ])
158
181
# NHWC -> NCHW
159
182
out_trans_node = helper .make_node ("Transpose" , [matmul_out ], [cnv_output ], perm = [0 , 3 , 1 , 2 ])
160
183
@@ -178,7 +201,16 @@ def extract_conv_params(self, model, node):
178
201
stride_w = get_by_name (node .attribute , "strides" ).ints [1 ]
179
202
group = get_by_name (node .attribute , "group" ).i
180
203
weight_name = node .input [1 ]
204
+ conv_weight_inp_name = node .input [1 ]
205
+ conv_weight_q_scale_name = None
181
206
W_conv = model .get_initializer (weight_name )
207
+ if W_conv is None :
208
+ # check to see if there is an immediate quantizer node feeding the weight input
209
+ w_producer = model .find_producer (weight_name )
210
+ if not (w_producer is None ) and w_producer .op_type == "Quant" :
211
+ W_conv = model .get_initializer (w_producer .input [0 ])
212
+ weight_name = w_producer .input [0 ]
213
+ conv_weight_q_scale_name = w_producer .input [1 ]
182
214
ifm_ch = model .get_tensor_shape (cnv_input )[1 ] # assume NCHW
183
215
ofm_ch = model .get_tensor_shape (cnv_output )[1 ] # assume NCHW
184
216
ifm_dim_h = model .get_tensor_shape (cnv_input )[2 ] # assume NCHW
@@ -213,6 +245,8 @@ def extract_conv_params(self, model, node):
213
245
stride_w ,
214
246
group ,
215
247
weight_name ,
248
+ conv_weight_inp_name ,
249
+ conv_weight_q_scale_name ,
216
250
W_conv ,
217
251
ifm_ch ,
218
252
ofm_ch ,
0 commit comments