(with_bias, with_relu)
| 223 | return data |
| 224 | |
| 225 | def _optimize(with_bias, with_relu): |
| 226 | pattern, first_key = _get_matmul_pattern(with_bias, with_relu) |
| 227 | ptm_list = util.get_matched_pattern(simple_graph, pattern, first_key) |
| 228 | ptm_list = [ptm for ptm in ptm_list if ptm['matmul'] not in update_dict] |
| 229 | if opt_config: |
| 230 | ptm_list = [ptm for ptm in ptm_list if ptm['matmul'] in opt_config] |
| 231 | for ptm in ptm_list: |
| 232 | if ptm['matmul'] in update_dict: |
| 233 | continue |
| 234 | w_data = _get_weight_data(ptm['matmul']) |
| 235 | if with_bias: |
| 236 | bias_data = _get_weight_data(ptm['bias_add']) |
| 237 | if w_data is None or (with_bias and bias_data is None): |
| 238 | continue |
| 239 | node = util.get_node_by_name(graph_def, simple_graph, ptm['matmul']) |
| 240 | opt_dtype = opt_config.get(node.name) if opt_config else data_type |
| 241 | print(f'Optimize dense op to {opt_dtype}: {node.name}') |
| 242 | update_dict[node.name] = [opt_dtype] |
| 243 | pref = node.name |
| 244 | dense_op = session.graph.get_operation_by_name(node.name) |
| 245 | if opt_dtype in [BF16, FP16]: |
| 246 | tf_dtype = tf.bfloat16 if opt_dtype == BF16 else tf.float16 |
| 247 | w_f16_ts = tf.constant( |
| 248 | value=tf.cast(w_data, tf_dtype).eval(), |
| 249 | dtype=tf_dtype, |
| 250 | name=f'{pref}/{opt_dtype.lower()}_weight', |
| 251 | ) |
| 252 | in_f16_ts = tf.cast( |
| 253 | dense_op.inputs[0], tf_dtype, name=f'{pref}/{opt_dtype.lower()}' |
| 254 | ) |
| 255 | out_f16_ts = tf.matmul( |
| 256 | a=in_f16_ts, |
| 257 | b=w_f16_ts, |
| 258 | transpose_a=dense_op.get_attr('transpose_a'), |
| 259 | transpose_b=dense_op.get_attr('transpose_b'), |
| 260 | name=f'{pref}/{opt_dtype.lower()}_matmul', |
| 261 | ) |
| 262 | if with_bias: |
| 263 | bias_f16_ts = tf.constant( |
| 264 | value=tf.cast(bias_data, tf_dtype).eval(), |
| 265 | dtype=tf_dtype, |
| 266 | name=f'{pref}/{opt_dtype.lower()}_bias', |
| 267 | ) |
| 268 | out_f16_ts = tf.nn.bias_add(out_f16_ts, bias_f16_ts) |
| 269 | out_f16_ts = tf.nn.relu(out_f16_ts) if with_relu else out_f16_ts |
| 270 | out_fp32_ts = tf.cast(out_f16_ts, tf.float32, name=f'{pref}/fp32') |
| 271 | update_op_inputs(session.graph, {ptm[first_key]: out_fp32_ts}) |
| 272 | continue |
| 273 | elif opt_dtype != INT8: |
| 274 | raise Exception(f'Unsupported data type: {opt_dtype}') |
| 275 | # Optimize to INT8 |
| 276 | # Update weight |
| 277 | w_max_abs_val = np.max(np.abs(w_data)) |
| 278 | w_scale = np.array(w_max_abs_val / 127.0) |
| 279 | w_int8_data = np.int8(np.round(w_data / w_scale)) |
| 280 | w_min_ts = tf.constant(-1 * w_max_abs_val, tf.float32, name=f'{pref}/w_min') |
| 281 | w_max_ts = tf.constant(w_max_abs_val, tf.float32, name=f'{pref}/w_max') |
| 282 | w_int8_ts = tf.constant(w_int8_data, tf.qint8, name=f'{pref}/int8_weight') |
no test coverage detected