transformer.py 106.2 KB
Newer Older
L
Liangliang He 已提交
1
# Copyright 2018 The MACE Authors. All Rights Reserved.
2 3 4 5 6 7 8 9 10 11 12 13 14
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

L
liutuo 已提交
15 16 17
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
18

Y
yejianwu 已提交
19
import re
20

21 22 23
import numpy as np
import six

L
liutuo 已提交
24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39
from python.py_proto import mace_pb2
from . import base_converter
from .base_converter import ConverterUtil
from .base_converter import DataFormat
from .base_converter import DeviceType
from .base_converter import EltwiseType
from .base_converter import FrameworkType
from .base_converter import MaceKeyword
from .base_converter import MaceOp
from .base_converter import MaceFixedDataFormatOps  # noqa
from .base_converter import MaceTransposableDataFormatOps  # noqa
from .base_converter import PaddingMode
from .base_converter import ReduceType
from .base_converter import TransformerRule
from python.quantize import quantize_util
from python.utils.util import mace_check
40 41 42 43 44 45 46 47 48


class Transformer(base_converter.ConverterInterface):
    """A class for transform naive mace model to optimized model.
    This Transformer should be platform irrelevant. So, do not assume
    tensor name has suffix like ':0".
    """

    def __init__(self, option, model):
49 50
        # Dependencies
        # (TRANSFORM_MATMUL_TO_FC, TRANSFORM_GLOBAL_CONV_TO_FC) -> RESHAPE_FC_WEIGHT  # noqa
51
        self._registered_transformers = {
李寅 已提交
52 53
            TransformerRule.TRANSFORM_FAKE_QUANTIZE:
                self.transform_fake_quantize,
54 55 56
            TransformerRule.REMOVE_IDENTITY_OP: self.remove_identity_op,
            TransformerRule.TRANSFORM_GLOBAL_POOLING:
                self.transform_global_pooling,
Y
yejianwu 已提交
57 58 59 60
            TransformerRule.TRANSFORM_LSTMCELL_ZEROSTATE:
                self.transform_lstmcell_zerostate,
            TransformerRule.TRANSFORM_BASIC_LSTMCELL:
                self.transform_basic_lstmcell,
李寅 已提交
61 62 63
            TransformerRule.FOLD_RESHAPE: self.fold_reshape,
            TransformerRule.TRANSFORM_MATMUL_TO_FC:
                self.transform_matmul_to_fc,
64
            TransformerRule.FOLD_BATCHNORM: self.fold_batchnorm,
65
            TransformerRule.FOLD_BIASADD: self.fold_biasadd,
66 67
            TransformerRule.FOLD_CONV_AND_BN:
                self.fold_conv_and_bn,  # data_format related
L
liutuo 已提交
68 69
            TransformerRule.FOLD_DECONV_AND_BN:
                self.fold_deconv_and_bn,  # data_format related
70 71 72 73
            TransformerRule.FOLD_DEPTHWISE_CONV_AND_BN:
                self.fold_depthwise_conv_and_bn,  # data_format related
            TransformerRule.TRANSFORM_ADD_TO_BIASADD:
                self.transform_add_to_biasadd,
B
Bin Li 已提交
74 75
            TransformerRule.REARRANGE_BATCH_TO_SPACE:
                self.rearrange_batch_to_space,
76
            TransformerRule.FLATTEN_ATROUS_CONV: self.flatten_atrous_conv,
77
            TransformerRule.FOLD_ACTIVATION: self.fold_activation,
78
            TransformerRule.FOLD_SQRDIFF_MEAN: self.fold_squared_diff_mean,
李寅 已提交
79
            TransformerRule.FOLD_EMBEDDING_LOOKUP: self.fold_embedding_lookup,
80
            TransformerRule.TRANSPOSE_FILTERS: self.transpose_filters,
81 82
            TransformerRule.TRANSPOSE_MATMUL_WEIGHT:
                self.transpose_matmul_weight,
Y
yejianwu 已提交
83 84
            TransformerRule.FOLD_FC_RESHAPE:
                self.fold_fc_reshape,
L
liuqi 已提交
85 86
            TransformerRule.ADD_IN_OUT_TENSOR_INFO:
                self.add_in_out_tensor_info,
87
            TransformerRule.ADD_WINOGRAD_ARG: self.add_winograd_arg,
88 89
            TransformerRule.TRANSFORM_GLOBAL_CONV_TO_FC:
                self.transform_global_conv_to_fc,
L
liuqi 已提交
90
            TransformerRule.RESHAPE_FC_WEIGHT: self.reshape_fc_weight,
李寅 已提交
91 92 93 94 95 96
            TransformerRule.QUANTIZE_NODES:
                self.quantize_nodes,
            TransformerRule.ADD_QUANTIZE_TENSOR_RANGE:
                self.add_quantize_tensor_range,
            TransformerRule.QUANTIZE_WEIGHTS:
                self.quantize_weights,
97 98
            TransformerRule.UPDATE_FLOAT_OP_DATA_TYPE:
                self.update_float_op_data_type,
99 100
            TransformerRule.ADD_OPENCL_INFORMATIONS:
                self.add_opencl_informations,
101
            TransformerRule.SORT_BY_EXECUTION: self.sort_by_execution,
102
            TransformerRule.UPDATE_DATA_FORMAT: self.update_data_format,
L
luxuhui 已提交
103 104
            TransformerRule.TRANSPOSE_RESHAPE_AND_FLATTEN:
                self.transform_reshape_and_flatten,
105
            TransformerRule.TRANSPOSE_DATA_FORMAT: self.transpose_data_format,
李寅 已提交
106 107
            TransformerRule.CHECK_QUANTIZE_INFO:
                self.check_quantize_info,
李寅 已提交
108 109
            TransformerRule.TRANSFORM_CHANNEL_SHUFFLE:
                self.transform_channel_shuffle,
110 111
            TransformerRule.QUANTIZE_SPECIFIC_OPS_ONLY:
                self.quantize_specific_ops_only,
112 113
            TransformerRule.FP16_MATMUL_WEIGHT:
                self.fp16_matmul_weight,
Y
yulianfei 已提交
114 115
            TransformerRule.FP16_GATHER_WEIGHT:
                self.fp16_gather_weight,
B
Bin Li 已提交
116 117
            TransformerRule.QUANTIZE_LARGE_WEIGHTS:
                self.quantize_large_weights,
118
        }
119 120 121

        self._option = option
        self._model = model
122
        self._wino_arg = self._option.winograd
123 124 125 126 127

        self._ops = {}
        self._consts = {}
        self._consumers = {}
        self._producer = {}
李寅 已提交
128 129
        self._quantize_activation_info = {}
        self._quantized_tensor = set()
130

B
Bin Li 已提交
131 132 133 134
        self.input_name_map = {}
        self.output_name_map = {}
        self.initialize_name_map()

135
    def run(self):
李寅 已提交
136 137 138
        for key in self._option.transformer_option:
            transformer = self._registered_transformers[key]
            while True:
139
                self.construct_ops_and_consumers(key)
李寅 已提交
140 141
                changed = transformer()
                if not changed:
142
                    break
B
Bin Li 已提交
143
        return self._model, self._quantize_activation_info
144

B
Bin Li 已提交
145 146 147 148 149 150 151 152 153 154 155 156
    def initialize_name_map(self):
        for input_node in self._option.input_nodes.values():
            new_input_name = MaceKeyword.mace_input_node_name \
                             + '_' + input_node.name
            self.input_name_map[input_node.name] = new_input_name

        output_nodes = self._option.check_nodes.values()
        for output_node in output_nodes:
            new_output_name = MaceKeyword.mace_output_node_name \
                              + '_' + output_node.name
            self.output_name_map[output_node.name] = new_output_name

157 158 159 160
    def filter_format(self):
        filter_format_value = ConverterUtil.get_arg(self._model,
                                                    MaceKeyword.mace_filter_format_str).i  # noqa
        filter_format = None
161 162 163 164 165 166
        if filter_format_value == DataFormat.HWIO.value:
            filter_format = DataFormat.HWIO
        elif filter_format_value == DataFormat.OIHW.value:
            filter_format = DataFormat.OIHW
        elif filter_format_value == DataFormat.HWOI.value:
            filter_format = DataFormat.HWOI
167 168 169 170 171 172 173 174 175 176 177
        else:
            mace_check(False, "filter format %d not supported" %
                       filter_format_value)

        return filter_format

    def set_filter_format(self, filter_format):
        arg = ConverterUtil.get_arg(self._model,
                                    MaceKeyword.mace_filter_format_str)
        arg.i = filter_format.value

178
    def construct_ops_and_consumers(self, key):
179 180 181 182 183 184 185 186 187 188 189 190 191 192 193
        self._ops.clear()
        self._consumers.clear()
        self._producer.clear()
        for op in self._model.op:
            self._ops[op.name] = op
        for tensor in self._model.tensors:
            self._consts[tensor.name] = tensor
        for op in self._ops.values():
            for input_tensor in op.input:
                if input_tensor not in self._consumers:
                    self._consumers[input_tensor] = []
                self._consumers[input_tensor].append(op)

            for output_tensor in op.output:
                self._producer[output_tensor] = op
194 195 196 197 198 199 200 201 202 203 204
        if key != TransformerRule.SORT_BY_EXECUTION:
            for input_node in self._option.input_nodes.values():
                input_node_existed = False
                for op in self._model.op:
                    if input_node.name in op.output:
                        input_node_existed = True
                        break
                if not input_node_existed:
                    op = mace_pb2.OperatorDef()
                    op.name = self.normalize_op_name(input_node.name)
                    op.type = "Input"
205 206
                    data_type_arg = op.arg.add()
                    data_type_arg.name = MaceKeyword.mace_op_data_type_str
207
                    data_type_arg.i = input_node.data_type
208 209 210
                    op.output.extend([input_node.name])
                    output_shape = op.output_shape.add()
                    output_shape.dims.extend(input_node.shape)
211
                    if input_node.data_format != DataFormat.NONE:
212
                        if input_node.data_format == DataFormat.NCHW:
213 214
                            self.transpose_shape(output_shape.dims,
                                                 [0, 3, 1, 2])
215
                        ConverterUtil.add_data_format_arg(op,
216
                                                          DataFormat.AUTO)
217 218
                    else:
                        ConverterUtil.add_data_format_arg(op,
219
                                                          DataFormat.NONE)
220
                    self._producer[op.output[0]] = op
221 222 223

    @staticmethod
    def replace(obj_list, source, target):
224
        for i in six.moves.range(len(obj_list)):
225 226 227 228 229 230
            if obj_list[i] == source:
                obj_list[i] = target

    @staticmethod
    def transpose_shape(shape, order):
        transposed_shape = []
231
        for i in six.moves.range(len(order)):
232 233 234 235 236 237 238
            transposed_shape.append(shape[order[i]])
        shape[:] = transposed_shape[:]

    @staticmethod
    def normalize_op_name(name):
        return name.replace(':', '_')

李寅 已提交
239
    def get_tensor_shape(self, tensor):
240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265
        if tensor in self._consts:
            return list(self._consts[tensor].dims)
        elif tensor in self._producer:
            producer = self._producer[tensor]
            for i in six.moves.range(len(producer.output)):
                if producer.output[i] == tensor:
                    return list(producer.output_shape[i].dims)
        else:
            return None

    def get_tensor_data_type(self, tensor):
        if tensor in self._consts:
            return self._consts[tensor].data_type
        elif tensor in self._producer:
            producer = self._producer[tensor]
            for i in six.moves.range(len(producer.output)):
                if producer.output[i] == tensor:
                    if i < len(producer.output_type):
                        return producer.output_type[i]
                    elif ConverterUtil.get_arg(producer, "T") is not None:
                        return ConverterUtil.get_arg(producer, "T").i
                    else:
                        print("No data type filled: ", producer)
                        return None
        else:
            return None
李寅 已提交
266

267 268 269 270 271
    def get_tensor_data_format(self, tensor):
        if tensor in self._producer:
            producer = self._producer[tensor]
            return ConverterUtil.data_format(producer)
        else:
272
            return DataFormat.NONE
273

274 275 276 277 278 279 280 281 282 283 284 285
    def consumer_count(self, tensor_name):
        return len(self._consumers.get(tensor_name, []))

    def is_op_output_node(self, op):
        output_node_tensor_names = [out for out in
                                    self._option.output_nodes]
        for output in op.output:
            if output in output_node_tensor_names:
                return True

        return False

L
liuqi 已提交
286
    def safe_remove_node(self, op, replace_op, remove_input_tensor=False):
李寅 已提交
287 288 289 290 291 292 293 294 295 296
        """remove op.
        1. change the inputs of its consumers to the outputs of replace_op
        2. if the op is output node, change output node to replace op"""

        if replace_op is None:
            # When no replace op specified, we change the inputs of
            # its consumers to the input of the op. This handles the case
            # that the op is identity op and its input is a tensor.
            mace_check(len(op.output) == 1 and len(op.input) == 1,
                       "cannot remove op that w/o replace op specified"
李寅 已提交
297
                       " and input/output length > 1\n" + str(op))
李寅 已提交
298 299 300 301 302 303 304 305 306 307 308

            for consumer_op in self._consumers.get(op.output[0], []):
                self.replace(consumer_op.input, op.output[0], op.input[0])

            mace_check(op.output[0] not in self._option.output_nodes,
                       "cannot remove op that is output node")
        else:
            mace_check(len(op.output) == len(replace_op.output),
                       "cannot remove op since len(op.output) "
                       "!= len(replace_op.output)")

309
            for i in six.moves.range(len(op.output)):
李寅 已提交
310 311 312 313 314 315 316
                for consumer_op in self._consumers.get(op.output[i], []):
                    self.replace(consumer_op.input,
                                 op.output[i],
                                 replace_op.output[i])

            # if the op is output node, change replace_op output name to the op
            # output name
317
            for i in six.moves.range(len(op.output)):
李寅 已提交
318 319 320 321 322 323 324 325
                if op.output[i] in self._option.output_nodes:
                    for consumer in self._consumers.get(
                            replace_op.output[i], []):
                        self.replace(consumer.input,
                                     replace_op.output[i],
                                     op.output[i])
                    replace_op.output[i] = op.output[i]

L
liuqi 已提交
326 327 328 329 330 331
        if remove_input_tensor:
            for input_name in op.input:
                if input_name in self._consts:
                    const_tensor = self._consts[input_name]
                    self._model.tensors.remove(const_tensor)

李寅 已提交
332 333
        self._model.op.remove(op)

334 335 336 337 338
    def add_in_out_tensor_info(self):
        net = self._model
        for input_node in self._option.input_nodes.values():
            input_info = net.input_info.add()
            input_info.name = input_node.name
339
            input_info.data_format = input_node.data_format.value
340
            input_info.dims.extend(input_node.shape)
L
liuqi 已提交
341
            input_info.data_type = input_node.data_type
342

B
Bin Li 已提交
343 344
        output_nodes = self._option.check_nodes.values()
        for output_node in output_nodes:
345 346
            output_info = net.output_info.add()
            output_info.name = output_node.name
347
            output_info.data_format = output_node.data_format.value
348 349
            output_info.dims.extend(
                self._producer[output_node.name].output_shape[0].dims)
L
liuqi 已提交
350
            output_info.data_type = output_node.data_type
351 352 353

        return False

354 355 356 357 358
    def remove_identity_op(self):
        net = self._model
        for op in net.op:
            if op.type == 'Identity':
                print("Remove identity: %s(%s)" % (op.name, op.type))
李寅 已提交
359 360
                self.safe_remove_node(op,
                                      self._producer.get(op.input[0], None))
361
                return True
叶剑武 已提交
362 363 364 365 366 367 368
            if op.type == 'Reshape' and \
                    op.output_shape[0].dims == \
                    self.get_tensor_shape(op.input[0]):
                print("Remove useless reshape: %s(%s)" % (op.name, op.type))
                self.safe_remove_node(op,
                                      self._producer.get(op.input[0], None))
                return True
369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398

        return False

    def transform_global_pooling(self):
        net = self._model
        for op in net.op:
            if op.type == MaceOp.Pooling.name and \
                            ConverterUtil.get_arg(op,
                                                  MaceKeyword.mace_global_pooling_str) is not None:  # noqa
                print("Transform global pooling: %s(%s)" % (op.name, op.type))
                input_shape = self._producer[op.input[0]].output_shape[0].dims
                if ConverterUtil.data_format(op) == DataFormat.NHWC:
                    kernel_shape = input_shape[1:3]
                else:
                    kernel_shape = input_shape[2:4]
                ConverterUtil.get_arg(op,
                                      MaceKeyword.mace_kernel_str).ints[:] \
                    = kernel_shape[:]

        return False

    def fold_batchnorm(self):
        net = self._model
        for op in net.op:
            if (op.type == MaceOp.Eltwise.name
                    and ConverterUtil.get_arg(
                        op, MaceKeyword.mace_element_type_str).i
                    == EltwiseType.PROD.value) \
                    and len(op.input) == 2 \
                    and op.input[1] in self._consts \
399 400
                    and op.output_shape[0].dims[-1:] == \
                    self._consts[op.input[1]].dims \
401 402 403 404 405
                    and self.consumer_count(op.output[0]) == 1 \
                    and not self.is_op_output_node(op):
                consumer_op = self._consumers[op.output[0]][0]
                if (consumer_op.type == MaceOp.Eltwise.name
                    and ConverterUtil.get_arg(
406
                        consumer_op, MaceKeyword.mace_element_type_str).i
407 408 409 410 411 412
                        == EltwiseType.SUM.value
                    or consumer_op.type == MaceOp.BiasAdd.name) \
                        and len(consumer_op.input) == 2 \
                        and consumer_op.input[1] in self._consts \
                        and len(self._consts[consumer_op.input[1]].dims) == 1:
                    print("Fold batchnorm: %s(%s)" % (op.name, op.type))
413
                    consumer_op.type = MaceOp.BatchNorm.name
李寅 已提交
414 415
                    consumer_op.input[:] = [op.input[0], op.input[1],
                                            consumer_op.input[1]]
416
                    net.op.remove(op)
417 418 419
                    return True
        return False

420 421 422 423
    def fold_squared_diff_mean(self):
        net = self._model
        for op in net.op:
            if op.type == MaceOp.Eltwise.name and len(op.input) == 2:
L
liutuo 已提交
424
                elt_type = ConverterUtil.get_arg(
425 426
                    op,
                    MaceKeyword.mace_element_type_str).i
L
liutuo 已提交
427
                if elt_type == EltwiseType.SQR_DIFF.value and\
428 429
                        self.consumer_count(op.output[0]) == 1:
                    consumer_op = self._consumers[op.output[0]][0]
L
liutuo 已提交
430 431 432 433 434 435 436 437 438 439
                    if consumer_op.type == MaceOp.Reduce.name:
                        axis = ConverterUtil.get_arg(
                            consumer_op,
                            MaceKeyword.mace_axis_str).ints
                        keep_dims = ConverterUtil.get_arg(
                            consumer_op,
                            MaceKeyword.mace_keepdims_str).i
                        reduce_type = ConverterUtil.get_arg(
                            consumer_op,
                            MaceKeyword.mace_reduce_type_str).i
L
liutuo 已提交
440
                        if reduce_type == ReduceType.MEAN.value and\
L
liutuo 已提交
441 442 443 444 445 446 447 448 449
                                len(consumer_op.input) == 1 and\
                                axis[0] == 1 and axis[1] == 2 and\
                                keep_dims > 0:
                            print("Fold SquaredDiff Reduce: %s" % op.name)
                            op.type = MaceOp.SqrDiffMean.name
                            op.output[0] = consumer_op.output[0]
                            self.replace_quantize_info(op, consumer_op)
                            self.safe_remove_node(consumer_op, op)
                            return True
450 451 452

        return False

李寅 已提交
453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469
    def fold_embedding_lookup(self):
        net = self._model
        for op in net.op:
            # gather -> mul
            if (op.type == MaceOp.Gather.name and
                    self.consumer_count(op.output[0]) == 1):
                consumer_op = self._consumers[op.output[0]][0]
                if (consumer_op.type == MaceOp.Eltwise.name and
                    ConverterUtil.get_arg(consumer_op,
                                          MaceKeyword.mace_element_type_str).i == EltwiseType.PROD.value and  # noqa
                            len(consumer_op.input) == 1 and
                            op.input[0] in self._consts and
                            self.consumer_count(op.input[0]) == 1):
                    print("Fold Gather and Mul: %s" % op.name)
                    gather_weights = self._consts[op.input[0]]
                    mul_weight = ConverterUtil.get_arg(consumer_op,
                                                       MaceKeyword.mace_scalar_input_str).f  # noqa
L
liukai6 已提交
470
                    gather_weights.float_data[:] = [float_data * mul_weight for float_data in gather_weights.float_data]  # noqa
李寅 已提交
471 472 473
                    self.safe_remove_node(consumer_op, None,
                                          remove_input_tensor=True)

Y
yejianwu 已提交
474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621
    def transform_lstmcell_zerostate(self):
        net = self._model

        zero_state_pattern = \
                re.compile(r'^.*BasicLSTMCellZeroState_?[0-9]*/[a-zA-Z]+_?[0-9]*')  # noqa
        for op in net.op:
            if op.type == MaceOp.Fill.name and \
                    zero_state_pattern.match(op.name):
                print("Transform lstm zerostate")
                concat_op = self._producer[op.input[0]]
                consumer_op = self._consumers[op.output[0]][0]

                dims = [self._consts[concat_op.input[0]].int32_data[0],
                        self._consts[concat_op.input[1]].int32_data[0]]
                tensor_def = net.tensors.add()
                tensor_def.name = op.output[0].replace('/zeros', '/init_const')
                tensor_def.dims.extend(dims)
                tensor_def.data_type = self._consts[op.input[1]].data_type
                tensor_def.float_data.extend(
                        [self._consts[op.input[1]].float_data[0]] *
                        (dims[0] * dims[1]))

                for i in range(len(consumer_op.input)):
                    if zero_state_pattern.match(consumer_op.input[i][:-2]):
                        consumer_op.input[i] = tensor_def.name

                net.tensors.remove(self._consts[op.input[1]])
                net.tensors.remove(self._consts[concat_op.input[0]])
                net.tensors.remove(self._consts[concat_op.input[1]])

                net.op.remove(concat_op)
                net.op.remove(op)

                return True

    def transform_basic_lstmcell(self):
        if self._option.device != DeviceType.GPU.value:
            return False

        net = self._model
        basic_lstm_concat_pattern = \
            re.compile(r'^.*basic_lstm_cell_?[0-9]*/concat_?[0-9]*')
        for op in net.op:
            if op.type == MaceOp.Concat.name and \
                    basic_lstm_concat_pattern.match(op.name):
                print("Transform basic lstmcell")
                ops_to_delete = []
                ops_to_delete.extend([op])

                op_def = net.op.add()
                op_def.name = op.name.replace('/concat', '/folded_lstmcell')
                op_def.type = MaceOp.LSTMCell.name
                op_def.arg.extend(op.arg[:-1])

                # Concat pre output and cur input
                # extend concat inputs
                op_def.input.extend([op_input for op_input in op.input])

                # lstm MatMul in FC of [pre_output, cur_input]
                matmul_op = self._consumers[op.output[0]][0]
                ops_to_delete.extend([matmul_op])
                # extend MatMul weight input
                op_def.input.extend([matmul_op.input[1]])

                # lstm BiasAdd in FC of [pre_output, cur_input]
                biasadd_op = self._consumers[matmul_op.output[0]][0]
                ops_to_delete.extend([biasadd_op])
                # extend BiasAdd bias input
                op_def.input.extend([biasadd_op.input[1]])

                # Split FC output into i, j, f, o
                # i = input_gate, j = new_input, f = forget_gate, o = output_gate  # noqa
                split_op = self._consumers[biasadd_op.output[0]][0]
                ops_to_delete.extend([split_op])

                # input gate activation
                input_gate_op = self._consumers[split_op.output[0]][0]
                ops_to_delete.extend([input_gate_op])
                # new input gate
                new_input_tanh_op = self._consumers[split_op.output[1]][0]
                ops_to_delete.extend([new_input_tanh_op])
                # forget gate add
                forget_add_op = self._consumers[split_op.output[2]][0]
                ops_to_delete.extend([forget_add_op])
                # output gate activation
                output_gate_op = self._consumers[split_op.output[3]][0]
                ops_to_delete.extend([output_gate_op])

                # extend forget add
                mace_check(len(forget_add_op.input) == 1,
                           'Wrong LSTM format in forget gate inputs')
                for arg in forget_add_op.arg:
                    if arg.name == MaceKeyword.mace_scalar_input_str:
                        op_def.arg.extend([arg])

                # state remember
                remember_mul_op = self._consumers[input_gate_op.output[0]][0]
                ops_to_delete.extend([remember_mul_op])
                mace_check(remember_mul_op.name == self._consumers[
                               new_input_tanh_op.output[0]][0].name,
                           'Wrong LSTM format in input sig & input tanh mul')

                # forget gate activation
                forget_gate_op = self._consumers[forget_add_op.output[0]][0]
                ops_to_delete.extend([forget_gate_op])

                # Mul `forget` & `pre cell state`
                forget_mul_op = self._consumers[forget_gate_op.output[0]][0]
                ops_to_delete.extend([forget_mul_op])

                # extend pre cell state input
                op_def.input.extend([forget_mul_op.input[0]])

                # get cur cell state
                # Add `forget gate output` & `remember mul output`
                remember_forget_add_op = \
                    self._consumers[remember_mul_op.output[0]][0]
                ops_to_delete.extend([remember_forget_add_op])
                mace_check(remember_forget_add_op.name ==
                           self._consumers[forget_mul_op.output[0]][0].name,
                           'Wrong LSTM format in add forget gate & remember mul')  # noqa
                op_def.output.extend([remember_forget_add_op.output[0]])
                op_def.output_shape.extend(remember_forget_add_op.output_shape)

                # cell state output tanh
                for consumer in \
                        self._consumers[remember_forget_add_op.output[0]]:
                    if consumer.type == MaceOp.Activation.name and \
                            consumer.name.find('basic_lstm_cell') > 0:
                        cell_tanh_op = consumer
                ops_to_delete.extend([cell_tanh_op])

                # final mul, get output
                final_mul_op = self._consumers[cell_tanh_op.output[0]][0]
                ops_to_delete.extend([final_mul_op])
                mace_check(final_mul_op.name ==
                           self._consumers[output_gate_op.output[0]][0].name,
                           'Wrong LSTM format in final mul')
                op_def.output.extend([final_mul_op.output[0]])
                op_def.output_shape.extend(final_mul_op.output_shape)

                for op_to_del in ops_to_delete:
                    net.op.remove(op_to_del)

                return True

        return False

622 623 624
    def fold_conv_and_bn(self):
        net = self._model
        for op in net.op:
L
liutuo 已提交
625
            if (op.type == MaceOp.Conv2D.name) \
626 627
                    and self.consumer_count(op.output[0]) == 1:
                consumer_op = self._consumers[op.output[0]][0]
628
                input_len = len(op.input)
629 630 631
                if (consumer_op.type == MaceOp.BatchNorm.name
                        and (input_len == 2 or (input_len == 3 and op.input[-1] in self._consts))  # noqa
                        and len(self._consumers[op.input[1]]) == 1):
632 633 634
                    print("Fold conv and bn: %s(%s)" % (op.name, op.type))
                    filter = self._consts[op.input[1]]
                    scale = self._consts[consumer_op.input[1]]
635
                    offset = self._consts[consumer_op.input[2]]
636 637
                    idx = 0
                    filter_format = self.filter_format()
638
                    if filter_format == DataFormat.HWIO:
639 640 641 642
                        for hwi in six.moves.range(filter.dims[0]
                                                   * filter.dims[1]
                                                   * filter.dims[2]):
                            for o in six.moves.range(filter.dims[3]):
643 644
                                filter.float_data[idx] *= scale.float_data[o]
                                idx += 1
645
                    elif filter_format == DataFormat.OIHW:
646 647 648 649
                        for o in six.moves.range(filter.dims[0]):
                            for hwi in six.moves.range(filter.dims[1]
                                                       * filter.dims[2]
                                                       * filter.dims[3]):
650 651 652 653 654 655
                                filter.float_data[idx] *= scale.float_data[o]
                                idx += 1
                    else:
                        mace_check(False, "filter format %s not supported" %
                                   filter_format)

656 657 658 659 660 661 662 663
                    if len(op.input) == 3:
                        conv_bias = self._consts[op.input[2]]
                        for c in six.moves.range(conv_bias.dims[0]):
                            conv_bias.float_data[c] *= scale.float_data[c]
                            conv_bias.float_data[c] += offset.float_data[c]
                        net.tensors.remove(offset)
                    else:
                        op.input.extend([consumer_op.input[2]])
664

665 666
                    # remove bn
                    del consumer_op.input[:]
667
                    net.tensors.remove(scale)
668
                    self.replace_quantize_info(op, consumer_op)
669 670
                    self.safe_remove_node(consumer_op, op)

671 672 673 674
                    return True

        return False

L
liutuo 已提交
675 676 677
    def fold_deconv_and_bn(self):
        net = self._model
        for op in net.op:
L
liutuo 已提交
678
            if (op.type in [MaceOp.Deconv2D.name, MaceOp.DepthwiseDeconv2d]) \
L
liutuo 已提交
679 680
                    and self.consumer_count(op.output[0]) == 1:
                consumer_op = self._consumers[op.output[0]][0]
681 682 683
                framework = ConverterUtil.get_arg(
                        op, MaceKeyword.mace_framework_type_str).i
                input_len = len(op.input)
684
                if (consumer_op.type == MaceOp.BatchNorm.name and (
685
                        (framework == FrameworkType.CAFFE.value and
L
liutuo 已提交
686 687
                         (input_len == 2 or (input_len == 3 and
                                             op.input[-1] in self._consts))) or
688 689
                        (framework == FrameworkType.TENSORFLOW.value and
                         (input_len == 3 or (input_len == 4 and
690 691
                                             op.input[-1] in self._consts))))
                        and len(self._consumers[op.input[1]]) == 1):
L
liutuo 已提交
692 693 694
                    print("Fold deconv and bn: %s(%s)" % (op.name, op.type))
                    filter = self._consts[op.input[1]]
                    scale = self._consts[consumer_op.input[1]]
695
                    offset = self._consts[consumer_op.input[2]]
L
liutuo 已提交
696 697 698
                    idx = 0
                    filter_format = self.filter_format()
                    # in deconv op O and I channel is switched
699
                    if filter_format == DataFormat.HWIO:
L
liutuo 已提交
700 701 702 703 704 705 706
                        for hw in six.moves.range(filter.dims[0]
                                                  * filter.dims[1]):
                            for o in six.moves.range(filter.dims[2]):
                                for i in six.moves.range(filter.dims[3]):
                                    filter.float_data[idx] *=\
                                        scale.float_data[o]
                                    idx += 1
707
                    elif filter_format == DataFormat.OIHW:
L
liutuo 已提交
708 709 710 711 712 713 714 715 716 717 718
                        for i in six.moves.range(filter.dims[0]):
                            for o in six.moves.range(filter.dims[1]):
                                for hw in six.moves.range(filter.dims[2]
                                                          * filter.dims[3]):
                                    filter.float_data[idx] *=\
                                        scale.float_data[o]
                                    idx += 1
                    else:
                        mace_check(False, "filter format %s not supported" %
                                   filter_format)

719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734
                    bias_dim = -1
                    if framework == FrameworkType.CAFFE.value \
                            and len(op.input) == 3:
                        bias_dim = 2
                    if framework == FrameworkType.TENSORFLOW.value \
                            and len(op.input) == 4:
                        bias_dim = 3

                    if bias_dim != -1:
                        conv_bias = self._consts[op.input[bias_dim]]
                        for c in six.moves.range(conv_bias.dims[0]):
                            conv_bias.float_data[c] *= scale.float_data[c]
                            conv_bias.float_data[c] += offset.float_data[c]
                        net.tensors.remove(offset)
                    else:
                        op.input.extend([consumer_op.input[2]])
L
liutuo 已提交
735

736
                    del consumer_op.input[:]
L
liutuo 已提交
737
                    net.tensors.remove(scale)
738
                    self.replace_quantize_info(op, consumer_op)
739 740
                    self.safe_remove_node(consumer_op, op)

L
liutuo 已提交
741 742 743 744
                    return True

        return False

745 746 747 748 749 750
    def fold_depthwise_conv_and_bn(self):
        net = self._model
        for op in net.op:
            if op.type == MaceOp.DepthwiseConv2d.name \
                    and self.consumer_count(op.output[0]) == 1:
                consumer_op = self._consumers[op.output[0]][0]
751
                input_len = len(op.input)
752 753 754
                if (consumer_op.type == MaceOp.BatchNorm.name
                        and (input_len == 2 or (input_len == 3 and op.input[-1] in self._consts))  # noqa
                        and len(self._consumers[op.input[1]]) == 1):
755 756 757 758
                    print("Fold depthwise conv and bn: %s(%s)"
                          % (op.name, op.type))
                    filter = self._consts[op.input[1]]
                    scale = self._consts[consumer_op.input[1]]
759
                    offset = self._consts[consumer_op.input[2]]
760 761 762
                    idx = 0

                    filter_format = self.filter_format()
763
                    if filter_format == DataFormat.HWIO:
764 765 766 767
                        for hw in six.moves.range(filter.dims[0]
                                                  * filter.dims[1]):
                            for i in six.moves.range(filter.dims[2]):
                                for o in six.moves.range(filter.dims[3]):
768
                                    filter.float_data[idx] *= scale.float_data[
769
                                                        i * filter.dims[3] + o]
770
                                    idx += 1
771
                    elif filter_format == DataFormat.OIHW:
772 773 774 775
                        for o in six.moves.range(filter.dims[0]):
                            for i in six.moves.range(filter.dims[1]):
                                for hw in six.moves.range(filter.dims[2]
                                                          * filter.dims[3]):
776 777 778 779 780 781 782
                                    filter.float_data[idx] *= scale.float_data[
                                        i * filter.dims[0] + o]
                                    idx += 1
                    else:
                        mace_check(False, "filter format %s not supported" %
                                   filter_format)

783 784 785 786 787 788 789 790
                    if len(op.input) == 3:
                        conv_bias = self._consts[op.input[2]]
                        for c in six.moves.range(conv_bias.dims[0]):
                            conv_bias.float_data[c] *= scale.float_data[c]
                            conv_bias.float_data[c] += offset.float_data[c]
                        net.tensors.remove(offset)
                    else:
                        op.input.extend([consumer_op.input[2]])
791

792 793
                    # remove bn
                    del consumer_op.input[:]
794
                    net.tensors.remove(scale)
795
                    self.replace_quantize_info(op, consumer_op)
796 797
                    self.safe_remove_node(consumer_op, op)

798 799 800 801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818
                    return True

        return False

    @staticmethod
    def sort_feature_map_shape(shape, data_format):
        """Return shape in NHWC order"""
        batch = shape[0]
        if data_format == DataFormat.NHWC:
            height = shape[1]
            width = shape[2]
            channels = shape[3]
        else:
            height = shape[2]
            width = shape[3]
            channels = shape[1]
        return batch, height, width, channels

    @staticmethod
    def sort_filter_shape(filter_shape, filter_format):
        """Return filter shape in HWIO order"""
819
        if filter_format == DataFormat.HWIO:
820 821 822 823
            filter_height = filter_shape[0]
            filter_width = filter_shape[1]
            in_channels = filter_shape[2]
            out_channels = filter_shape[3]
824
        elif filter_format == DataFormat.OIHW:
825 826 827 828
            filter_height = filter_shape[2]
            filter_width = filter_shape[3]
            in_channels = filter_shape[1]
            out_channels = filter_shape[0]
829
        elif filter_format == DataFormat.HWOI:
830 831 832 833 834 835 836 837 838 839 840
            filter_height = filter_shape[0]
            filter_width = filter_shape[1]
            in_channels = filter_shape[3]
            out_channels = filter_shape[2]
        else:
            mace_check(False, "filter format %s not supported" % filter_format)
        return filter_height, filter_width, in_channels, out_channels

    def transform_add_to_biasadd(self):
        net = self._model
        for op in net.op:
李寅 已提交
841 842 843 844 845
            if (op.type == 'Eltwise'
                    and ConverterUtil.get_arg(op, MaceKeyword.mace_element_type_str).i == EltwiseType.SUM.value  # noqa
                    and len(op.input) == 2
                    and op.input[1] in self._consts
                    and len(self._consts[op.input[1]].dims) == 1):
846 847 848 849 850 851
                print("Transform add to biasadd: %s(%s)" % (op.name, op.type))
                op.type = MaceOp.BiasAdd.name
                return True

        return False

李寅 已提交
852 853 854 855
    def replace_quantize_info(self, op, replace_op):
        if len(replace_op.quantize_info) > 0:
            del op.quantize_info[:]
            op.quantize_info.extend(replace_op.quantize_info)
B
Bin Li 已提交
856 857 858
            for i in range(len(op.quantize_info)):
                self._quantize_activation_info[op.output[i]] = \
                    op.quantize_info[i]
李寅 已提交
859

860 861 862
    def fold_biasadd(self):
        net = self._model
        for op in net.op:
863 864 865 866
            if (((op.type == MaceOp.Conv2D.name
                  or op.type == MaceOp.DepthwiseConv2d.name
                  or op.type == MaceOp.FullyConnected.name)
                 and len(op.input) == 2)
L
liutuo 已提交
867 868 869 870 871 872 873 874 875 876 877
                or (op.type == MaceOp.Deconv2D.name
                    and ((ConverterUtil.get_arg(
                                op,
                                MaceKeyword.mace_framework_type_str).i ==
                          FrameworkType.CAFFE.value
                          and len(op.input) == 2)
                         or (ConverterUtil.get_arg(
                                        op,
                                        MaceKeyword.mace_framework_type_str).i
                             == FrameworkType.TENSORFLOW.value
                             and len(op.input) == 3)))) \
878 879 880 881
                    and len(self._consumers.get(op.output[0], [])) == 1:
                consumer_op = self._consumers[op.output[0]][0]
                if consumer_op.type == MaceOp.BiasAdd.name:
                    print("Fold biasadd: %s(%s)" % (op.name, op.type))
B
Bin Li 已提交
882 883
                    op.name = consumer_op.name
                    op.output[0] = consumer_op.output[0]
884
                    op.input.append(consumer_op.input[1])
李寅 已提交
885
                    self.replace_quantize_info(op, consumer_op)
李寅 已提交
886
                    self.safe_remove_node(consumer_op, op)
887 888 889 890
                    return True

        return False

891
    def flatten_atrous_conv(self):
892
        if self._option.device != DeviceType.GPU.value \
893 894
               and self._option.device != DeviceType.APU.value \
               and self._option.device != DeviceType.HTA.value:
895 896 897 898 899 900 901 902 903 904 905 906
            return

        net = self._model
        for op in net.op:
            if (op.type == MaceOp.SpaceToBatchND.name
                    and len(self._consumers.get(op.output[0], [])) == 1):
                conv_op = self._consumers.get(op.output[0])[0]
                if (conv_op.type == MaceOp.Conv2D.name
                        or conv_op.type == MaceOp.DepthwiseConv2d.name) \
                        and len(self._consumers.get(conv_op.output[0], [])) == 1:  # noqa
                    b2s_op = self._consumers.get(conv_op.output[0])[0]
                    if b2s_op.type == MaceOp.BatchToSpaceND.name:
907
                        six.print_("Flatten atrous convolution")
908 909 910 911 912 913 914 915 916 917 918 919 920 921 922 923 924 925 926 927 928 929 930 931 932 933 934 935 936 937 938 939 940 941 942 943 944 945 946 947 948 949 950 951
                        # Add args.
                        padding_arg_values = ConverterUtil.get_arg(
                            op,
                            MaceKeyword.mace_paddings_str).ints
                        blocks_arg_values = ConverterUtil.get_arg(
                            b2s_op,
                            MaceKeyword.mace_space_batch_block_shape_str).ints
                        dilation_arg = ConverterUtil.get_arg(
                            conv_op,
                            MaceKeyword.mace_dilations_str)
                        if dilation_arg is None:
                            dilation_arg = conv_op.arg.add()
                        dilation_arg.name = MaceKeyword.mace_dilations_str
                        dilation_arg.ints[:] = blocks_arg_values

                        padding_arg = ConverterUtil.get_arg(
                            conv_op,
                            MaceKeyword.mace_padding_str)
                        if padding_arg is None:
                            padding_arg = conv_op.arg.add()
                        padding_arg.name = MaceKeyword.mace_padding_str
                        if len(padding_arg_values) > 0 \
                                and padding_arg_values[0] > 0:
                            padding_arg.i = PaddingMode.SAME.value
                        else:
                            padding_arg.i = PaddingMode.VALID.value

                        strides_arg = ConverterUtil.get_arg(
                            conv_op,
                            MaceKeyword.mace_strides_str)
                        if strides_arg is None:
                            strides_arg = conv_op.arg.add()
                        strides_arg.name = MaceKeyword.mace_strides_str
                        strides_arg.ints[:] = [1, 1]

                        # update output shape
                        conv_op.output_shape[0].dims[:] = \
                            b2s_op.output_shape[0].dims[:]

                        self.safe_remove_node(op, None)
                        self.safe_remove_node(b2s_op, conv_op)
                        return True
        return False

952 953 954 955 956 957 958
    def fold_activation(self):
        net = self._model
        for op in net.op:
            if (op.type == MaceOp.Conv2D.name
                or op.type == MaceOp.Deconv2D.name
                or op.type == MaceOp.DepthwiseConv2d.name
                or op.type == MaceOp.FullyConnected.name
959
                or op.type == MaceOp.BatchNorm.name) \
960 961 962 963 964 965 966 967 968 969 970
                    and len(self._consumers.get(op.output[0], [])) == 1:
                consumer_op = self._consumers[op.output[0]][0]
                if consumer_op.type == MaceOp.Activation.name \
                        and ConverterUtil.get_arg(
                            consumer_op,
                            MaceKeyword.mace_activation_type_str).s != 'PRELU':
                    print("Fold activation: %s(%s)" % (op.name, op.type))
                    op.name = consumer_op.name
                    op.output[0] = consumer_op.output[0]
                    for arg in consumer_op.arg:
                        if arg.name == MaceKeyword.mace_activation_type_str \
Y
yejianwu 已提交
971 972 973
                                or arg.name == \
                                    MaceKeyword.mace_activation_max_limit_str \
                                or arg.name == MaceKeyword.mace_activation_leakyrelu_coefficient_str:  # noqa
974 975
                            op.arg.extend([arg])

李寅 已提交
976
                    self.replace_quantize_info(op, consumer_op)
李寅 已提交
977
                    self.safe_remove_node(consumer_op, op)
978 979 980 981
                    return True

        return False

982 983 984 985 986 987 988 989 990
    def transform_global_conv_to_fc(self):
        """Transform global conv to fc should be placed after transposing
        input/output and filter"""

        if self._option.quantize:
            return

        net = self._model
        for op in net.op:
991 992 993
            if op.type == MaceOp.Conv2D.name \
                    and len(op.input) >= 2 \
                    and op.input[1] in self._consts:
994 995 996 997 998 999 1000 1001 1002 1003 1004 1005 1006 1007 1008 1009 1010 1011 1012 1013 1014 1015
                producer = self._producer[op.input[0]]
                input_shape = producer.output_shape[0].dims
                batch, height, width, channels = self.sort_feature_map_shape(
                    input_shape, ConverterUtil.data_format(producer))
                filter = self._consts[op.input[1]]
                filter_shape = filter.dims
                filter_height, filter_width, in_channels, out_channels = \
                    self.sort_filter_shape(filter_shape, self.filter_format())
                zero_padding = True
                padding_arg = ConverterUtil.get_arg(op,
                                                    MaceKeyword.mace_padding_str)  # noqa
                if padding_arg is not None:
                    if padding_arg.i != PaddingMode.VALID.value:
                        zero_padding = False
                else:
                    padding_value_arg = ConverterUtil.get_arg(op,
                                                              MaceKeyword.mace_padding_values_str)  # noqa
                    if padding_value_arg is not None:
                        if not all(v == 0 for v in padding_value_arg.ints):
                            zero_padding = False

                if height == filter_height and width == filter_width \
1016 1017
                        and zero_padding \
                        and len(self._consumers[op.input[1]]) == 1:
1018 1019 1020 1021 1022 1023 1024 1025 1026 1027 1028 1029 1030
                    print("transform global conv to fc %s(%s)"
                          % (op.name, op.type))
                    op.type = MaceOp.FullyConnected.name

        return False

    def reshape_fc_weight(self):
        net = self._model
        filter_format = self.filter_format()
        for op in net.op:
            if op.type == MaceOp.FullyConnected.name:
                weight = self._consts[op.input[1]]
                if len(weight.dims) == 2:
L
liutuo 已提交
1031
                    print("Reshape fully connected weight shape")
1032 1033 1034 1035
                    input_op = self._producer[op.input[0]]
                    input_shape = list(input_op.output_shape[0].dims)
                    weight.dims[:] = [weight.dims[0]] + input_shape[1:]
                    if len(input_shape) == 2:
1036
                        if filter_format == DataFormat.HWIO:
1037
                            weight.dims[:] = [1, 1] + weight.dims[:]
1038
                        elif filter_format == DataFormat.OIHW:
1039 1040
                            weight.dims[:] = weight.dims[:] + [1, 1]
                        else:
L
liuqi 已提交
1041 1042
                            mace_check(False,
                                       "FC does not support filter format %s" %
1043 1044 1045
                                       filter_format.name)
        return False

1046 1047 1048 1049
    def add_winograd_arg(self):
        if self._wino_arg == 0:
            return False
        net = self._model
李寅 已提交
1050

1051 1052 1053 1054 1055
        for op in net.op:
            if op.type == MaceOp.Conv2D.name:
                winograd_arg = op.arg.add()
                winograd_arg.name = MaceKeyword.mace_wino_arg_str
                winograd_arg.i = self._wino_arg
李寅 已提交
1056

1057 1058
        return False

1059 1060 1061 1062
    def transpose_matmul_weight(self):
        if self._option.device != DeviceType.CPU.value:
            return False
        net = self._model
1063
        transposed_weights = []
1064 1065
        for op in net.op:
            if op.type == MaceOp.MatMul.name:  # noqa
1066 1067 1068
                rhs = op.input[1]
                if rhs in self._consts and len(self._consts[rhs].dims) == 2:
                    arg = ConverterUtil.get_arg(op, MaceKeyword.mace_transpose_b_str)  # noqa
L
liutuo 已提交
1069
                    # six.print_("Transpose matmul weight %s" % rhs)
1070 1071 1072 1073 1074 1075
                    if arg is None:
                        arg = op.arg.add()
                        arg.name = MaceKeyword.mace_transpose_b_str
                        arg.i = 0
                    if arg.i == 0:
                        arg.i = 1
1076 1077 1078 1079 1080 1081 1082 1083 1084 1085
                        if rhs not in transposed_weights:
                            filter = self._consts[rhs]
                            filter_data = np.array(filter.float_data).reshape(
                                filter.dims)
                            filter_data = filter_data.transpose(1, 0)
                            filter.float_data[:] = filter_data.flat
                            filter.dims[:] = filter_data.shape
                            transposed_weights.append(rhs)
                            six.print_('Transpose matmul weight to shape:',
                                       filter.dims)
1086

1087 1088 1089
    def transpose_filters(self):
        net = self._model
        filter_format = self.filter_format()
1090
        transposed_filter = set()
L
liutuo 已提交
1091
        transposed_deconv_filter = set()
1092

李寅 已提交
1093
        if self._option.quantize and \
1094 1095
                (self._option.device == DeviceType.CPU.value or
                 self._option.device == DeviceType.APU.value):
1096
            print("Transpose filters to OHWI")
1097
            if filter_format == DataFormat.HWIO:
1098
                transpose_order = [3, 0, 1, 2]
1099
            elif filter_format == DataFormat.OIHW:
1100 1101
                transpose_order = [0, 2, 3, 1]
            else:
L
liuqi 已提交
1102
                mace_check(False, "Quantize model does not support conv "
1103 1104
                           "filter format: %s" % filter_format.name)

1105
            for op in net.op:
L
liutuo 已提交
1106
                if (op.type == MaceOp.Conv2D.name or
1107 1108 1109
                    op.type == MaceOp.Deconv2D.name or
                    (op.type == MaceOp.DepthwiseConv2d.name and
                     self._option.device == DeviceType.APU.value)) and\
L
liutuo 已提交
1110
                        op.input[1] not in transposed_filter:
1111 1112 1113
                    filter = self._consts[op.input[1]]
                    filter_data = np.array(filter.float_data).reshape(
                        filter.dims)
1114
                    filter_data = filter_data.transpose(transpose_order)
1115 1116
                    filter.float_data[:] = filter_data.flat
                    filter.dims[:] = filter_data.shape
1117
                    transposed_filter.add(op.input[1])
L
liutuo 已提交
1118 1119 1120 1121 1122 1123 1124 1125 1126 1127 1128
            # deconv's filter's output channel and input channel is reversed
            for op in net.op:
                if op.type == MaceOp.Deconv2D.name and \
                        op.input[1] not in transposed_deconv_filter:
                    filter = self._consts[op.input[1]]
                    filter_data = np.array(filter.float_data).reshape(
                        filter.dims)
                    filter_data = filter_data.transpose(3, 1, 2, 0)
                    filter.float_data[:] = filter_data.flat
                    filter.dims[:] = filter_data.shape
                    transposed_deconv_filter.add(op.input[1])
1129

1130
            self.set_filter_format(DataFormat.OHWI)
B
Bin Li 已提交
1131
        elif self._option.quantize and \
B
Bin Li 已提交
1132 1133
                (self._option.device == DeviceType.HEXAGON.value or
                 self._option.device == DeviceType.HTA.value):
B
Bin Li 已提交
1134
            print("Transpose filters to HWIO/HWIM")
1135
            mace_check(filter_format == DataFormat.HWIO,
B
Bin Li 已提交
1136
                       "HEXAGON only support HWIO/HWIM filter format.")
1137 1138
        else:
            # transpose filter to OIHW/MIHW for tensorflow (HWIO/HWIM)
1139
            if filter_format == DataFormat.HWIO:
1140
                for op in net.op:
1141 1142 1143
                    if (op.type == MaceOp.Conv2D.name
                            or op.type == MaceOp.Deconv2D.name
                            or op.type == MaceOp.DepthwiseConv2d.name) \
1144
                            and op.input[1] in self._consts \
1145
                            and op.input[1] not in transposed_filter:
L
liutuo 已提交
1146
                        print("Transpose Conv2D/Deconv2D filters to OIHW/MIHW")
1147 1148 1149 1150 1151 1152
                        filter = self._consts[op.input[1]]
                        filter_data = np.array(filter.float_data).reshape(
                            filter.dims)
                        filter_data = filter_data.transpose(3, 2, 0, 1)
                        filter.float_data[:] = filter_data.flat
                        filter.dims[:] = filter_data.shape
1153
                        transposed_filter.add(op.input[1])
L
liutuo 已提交
1154 1155 1156 1157 1158
                    if (op.type == MaceOp.MatMul.name and
                            (ConverterUtil.get_arg(
                                op,
                                MaceKeyword.mace_winograd_filter_transformed)
                                 is not None)  # noqa
1159
                            and op.input[1] not in transposed_filter):
L
liutuo 已提交
1160
                        print("Transpose Winograd filters to OIHW/MIHW")
1161 1162 1163 1164 1165 1166
                        filter = self._consts[op.input[0]]
                        filter_data = np.array(filter.float_data).reshape(
                            filter.dims)
                        filter_data = filter_data.transpose(3, 2, 0, 1)
                        filter.float_data[:] = filter_data.flat
                        filter.dims[:] = filter_data.shape
1167 1168 1169
                        transposed_filter.add(op.input[0])
                    if op.type == MaceOp.FullyConnected.name \
                            and op.input[1] not in transposed_filter:
1170 1171
                        weight = self._consts[op.input[1]]
                        if len(weight.dims) == 4:
L
liutuo 已提交
1172 1173
                            print("Transpose FullyConnected filters to"
                                  " OIHW/MIHW")
1174 1175 1176 1177 1178
                            weight_data = np.array(weight.float_data).reshape(
                                weight.dims)
                            weight_data = weight_data.transpose(3, 2, 0, 1)
                            weight.float_data[:] = weight_data.flat
                            weight.dims[:] = weight_data.shape
1179
                            transposed_filter.add(op.input[1])
1180

1181
                self.set_filter_format(DataFormat.OIHW)
L
liutuo 已提交
1182
            # deconv's filter's output channel and input channel is reversed
1183
            for op in net.op:
L
liutuo 已提交
1184 1185
                if op.type in [MaceOp.Deconv2D.name,
                               MaceOp.DepthwiseDeconv2d] \
L
liutuo 已提交
1186
                        and op.input[1] not in transposed_deconv_filter:
1187
                    filter = self._consts[op.input[1]]
1188 1189
                    filter_data = np.array(filter.float_data).reshape(
                        filter.dims)
1190
                    filter_data = filter_data.transpose(1, 0, 2, 3)
1191 1192
                    filter.float_data[:] = filter_data.flat
                    filter.dims[:] = filter_data.shape
L
liutuo 已提交
1193
                    transposed_deconv_filter.add(op.input[1])
1194 1195 1196

        return False

李寅 已提交
1197
    def fold_reshape(self):
1198 1199
        net = self._model
        for op in net.op:
1200 1201
            if op.type == MaceOp.Softmax.name:
                # see if possible to fold
1202
                # Reshape(xd->2d) + Softmax(2d) [+ Reshape(xd)] to Softmax(xd)
1203 1204 1205 1206
                should_fold = False
                if op.input[0] in self._producer \
                        and self._producer[op.input[0]].type \
                        == MaceOp.Reshape.name \
1207
                        and len(op.output_shape[0].dims) == 2:
L
liutuo 已提交
1208 1209 1210 1211 1212
                    producer = self._producer[op.input[0]]
                    reshape_input_rank = len(self.get_tensor_shape(
                        producer.input[0]))
                    if reshape_input_rank == 4:
                        should_fold = True
1213 1214 1215 1216 1217

                if should_fold:
                    print(
                        "Fold reshape and softmax: %s(%s)"
                        % (op.name, op.type))
1218
                    producer = self._producer[op.input[0]]
1219 1220 1221
                    op.output_shape[0].dims[:] = self.get_tensor_shape(
                        producer.input[0])

1222 1223 1224
                    if op.output[0] in self._consumers:
                        consumer = self._consumers[op.output[0]][0]
                        # if there is a shape op, remove it too
L
liutuo 已提交
1225 1226 1227 1228 1229 1230 1231
                        if len(consumer.input) > 1:
                            if (consumer.input[1] in self._producer
                                and self._producer[consumer.input[1]].type
                                    == 'Shape'):
                                self.safe_remove_node(
                                    self._producer[consumer.input[1]], None,
                                    remove_input_tensor=True)
1232 1233 1234
                        # remove consumer reshape
                        self.safe_remove_node(consumer, op,
                                              remove_input_tensor=True)
1235 1236 1237
                    # remove producer reshape
                    self.safe_remove_node(producer,
                                          self._producer.get(producer.input[0],
L
liuqi 已提交
1238 1239
                                                             None),
                                          remove_input_tensor=True)
1240

1241
                    return True
1242 1243
        return False

Y
yejianwu 已提交
1244 1245 1246 1247 1248 1249 1250 1251 1252 1253 1254 1255
    def is_after_fc(self, op):
        while op.input[0] in self._producer:
            producer = self._producer[op.input[0]]
            if producer.type in [MaceOp.Activation.name, MaceOp.BiasAdd.name]:
                op = producer
                continue
            elif producer.type == MaceOp.FullyConnected.name:
                return True
            else:
                return False
        return False

李寅 已提交
1256 1257
    def transform_matmul_to_fc(self):
        net = self._model
L
liuqi 已提交
1258
        filter_format = self.filter_format()
李寅 已提交
1259
        for op in net.op:
Y
yejianwu 已提交
1260 1261
            # transform `input(4D) -> reshape(2D) -> matmul` to `fc(2D)`
            # fc output is 2D in transformer, using as 4D in op kernel
L
liuqi 已提交
1262
            # work for TensorFlow
L
liuqi 已提交
1263
            if op.type == MaceOp.Reshape.name and \
L
lichao18 已提交
1264
                    len(op.input) == 2 and \
L
liuqi 已提交
1265 1266
                    op.input[1] in self._consts and \
                    len(op.output_shape[0].dims) == 2 and \
1267
                    filter_format == DataFormat.HWIO and \
李寅 已提交
1268
                    op.input[0] in self._producer:
L
liuqi 已提交
1269 1270 1271 1272 1273 1274 1275 1276 1277 1278 1279 1280 1281 1282 1283 1284 1285
                input_op = self._producer[op.input[0]]
                input_shape = input_op.output_shape[0].dims
                # check input op
                if len(input_shape) == 4 and \
                        np.prod(input_shape[1:]) == op.output_shape[0].dims[1]:
                    is_fc = True
                    consumers = self._consumers[op.output[0]]
                    # check matmul op
                    for matmul_op in consumers:
                        if matmul_op.type != MaceOp.MatMul.name:
                            is_fc = False
                        else:
                            weight = self._consts[matmul_op.input[1]]
                            if len(weight.dims) != 2 or \
                               weight.dims[0] != op.output_shape[0].dims[1]:
                                is_fc = False
                    if is_fc:
1286
                        print('convert reshape and matmul to fc')
L
liuqi 已提交
1287 1288 1289 1290 1291 1292 1293 1294 1295 1296 1297
                        self.safe_remove_node(op, input_op,
                                              remove_input_tensor=True)
                        for matmul_op in consumers:
                            weight = self._consts[matmul_op.input[1]]
                            matmul_op.type = MaceOp.FullyConnected.name
                            weight_data = np.array(weight.float_data).reshape(
                                weight.dims)
                            weight.dims[:] = input_shape[1:] + \
                                [weight_data.shape[1]]
                        return True

Y
yejianwu 已提交
1298 1299
            # transform `fc1(2D) -> matmul` to `fc1(2D) -> fc1(2D)`
            if op.type == MaceOp.MatMul.name and \
1300
                    filter_format == DataFormat.HWIO and \
Y
yejianwu 已提交
1301
                    op.input[1] in self._consts:
Y
yejianwu 已提交
1302 1303 1304 1305 1306 1307 1308 1309 1310 1311 1312 1313
                producer = self._producer[op.input[0]]
                weight = self._consts[op.input[1]]
                if len(weight.dims) == 2 and self.is_after_fc(op) and \
                        len(producer.output_shape[0].dims) == 2 and \
                        weight.dims[0] == producer.output_shape[0].dims[1]:
                    six.print_('convert matmul to fc')
                    op.type = MaceOp.FullyConnected.name
                    weight_data = np.array(weight.float_data).reshape(
                        weight.dims)
                    weight.dims[:] = [1, 1] + list(weight_data.shape)
                    return True

李寅 已提交
1314 1315
        return False

1316 1317 1318
    def update_float_op_data_type(self):
        print("update op with float data type")
        net = self._model
1319
        data_type = self._option.data_type
1320 1321 1322 1323 1324
        net.data_type = data_type

        if self._option.quantize:
            return

1325 1326 1327 1328 1329 1330
        for op in net.op:
            data_type_arg = ConverterUtil.get_arg(
                op, MaceKeyword.mace_op_data_type_str)
            if not data_type_arg:
                data_type_arg = op.arg.add()
                data_type_arg.name = MaceKeyword.mace_op_data_type_str
1331 1332
                data_type_arg.i = data_type
            elif data_type_arg.i != data_type \
L
liuqi 已提交
1333
                    and data_type_arg.i == mace_pb2.DT_FLOAT:
1334
                data_type_arg.i = data_type
1335 1336 1337

        return False

1338
    def sort_dfs(self, op, visited, sorted_nodes):
1339 1340
        if op.name in visited:
            return
1341 1342 1343 1344 1345 1346 1347 1348 1349 1350 1351 1352 1353 1354 1355 1356
        visited.update([op.name])
        if len(op.input) > 0:
            for input_tensor in op.input:
                producer_op = self._producer.get(input_tensor, None)
                if producer_op is None:
                    pass
                elif producer_op.name not in visited:
                    self.sort_dfs(producer_op, visited, sorted_nodes)
        sorted_nodes.append(op)

    def sort_by_execution(self):
        print("Sort by execution")
        net = self._model
        visited = set()
        sorted_nodes = []

B
Bin Li 已提交
1357
        output_nodes = self._option.check_nodes.keys()
B
Bin Li 已提交
1358
        if not self._quantize_activation_info:
B
Bin Li 已提交
1359
            output_nodes.extend(self._option.output_nodes)
B
Bin Li 已提交
1360
        for output_node in output_nodes:
1361 1362 1363
            mace_check(output_node in self._producer,
                       "output_tensor %s not existed in model" % output_node)
            self.sort_dfs(self._producer[output_node], visited, sorted_nodes)
1364 1365 1366

        del net.op[:]
        net.op.extend(sorted_nodes)
李寅 已提交
1367 1368

        print("Final ops:")
B
Bin Li 已提交
1369
        index = 0
李寅 已提交
1370
        for op in net.op:
B
Bin Li 已提交
1371 1372 1373 1374 1375 1376
            if op.type not in [MaceOp.Quantize.name, MaceOp.Dequantize.name]:
                index_str = str(index)
                index += 1
            else:
                index_str = ''
            print("%s (%s, index:%s): %s" % (op.name, op.type, index_str, [
李寅 已提交
1377
                out_shape.dims for out_shape in op.output_shape]))
1378
        return False
李寅 已提交
1379

1380
    def update_data_format(self):
1381
        print("update data format")
1382 1383
        net = self._model
        for op in net.op:
1384
            df_arg = ConverterUtil.get_arg(
1385
                op, MaceKeyword.mace_data_format_str)
1386 1387 1388
            if not df_arg:
                df_arg = op.arg.add()
                df_arg.name = MaceKeyword.mace_data_format_str
1389
            if op.type in MaceFixedDataFormatOps:
1390
                df_arg.i = DataFormat.AUTO.value
1391
            elif op.type in MaceTransposableDataFormatOps:
1392
                input_df = DataFormat.AUTO.value
1393 1394 1395
                for input_tensor in op.input:
                    if input_tensor in self._consts:
                        continue
1396 1397 1398
                    mace_check(
                        input_tensor in self._producer,
                        "Input tensor %s not in producer" % input_tensor)
1399 1400 1401
                    father_op = self._producer[input_tensor]
                    temp_input_df = ConverterUtil.get_arg(
                        father_op, MaceKeyword.mace_data_format_str)
1402
                    if temp_input_df.i != DataFormat.AUTO.value:
1403
                        input_df = temp_input_df.i
1404
                if input_df == DataFormat.AUTO.value:
1405 1406 1407 1408 1409 1410 1411 1412 1413 1414 1415 1416 1417 1418 1419
                    df_arg.i = input_df
                    # add flag to mark the ops may has data format
                    has_data_format_arg = op.arg.add()
                    has_data_format_arg.name = \
                        MaceKeyword.mace_has_data_format_str
                    has_data_format_arg.i = 1
        return False

    def transpose_data_format(self):
        print("Transpose arguments based on data format")
        net = self._model

        src_data_format = ConverterUtil.data_format(net)
        for op in net.op:
            has_data_format = ConverterUtil.data_format(op) == \
1420
                              DataFormat.AUTO
1421 1422 1423 1424 1425 1426 1427 1428 1429 1430 1431 1432 1433 1434 1435 1436 1437 1438 1439 1440 1441 1442 1443 1444 1445 1446
            # transpose args
            if op.type == MaceOp.Pad.name:
                for arg in op.arg:
                    if arg.name == MaceKeyword.mace_paddings_str:
                        mace_check(len(arg.ints) == 8,
                                   "pad dim rank should be 8.")
                        if src_data_format == DataFormat.NCHW and \
                                has_data_format:
                            print("Transpose pad args: %s(%s)"
                                  % (op.name, op.type))
                            self.transpose_shape(arg.ints,
                                                 [0, 1, 4, 5, 6, 7, 2, 3])
            elif op.type == MaceOp.Concat.name or op.type == MaceOp.Split.name:
                for arg in op.arg:
                    if arg.name == MaceKeyword.mace_axis_str:
                        if (src_data_format == DataFormat.NCHW
                                and has_data_format
                                and len(op.output_shape[0].dims) == 4):
                            print("Transpose concat/split args: %s(%s)"
                                  % (op.name, op.type))
                            if arg.i == 1:
                                arg.i = 3
                            elif arg.i == 2:
                                arg.i = 1
                            elif arg.i == 3:
                                arg.i = 2
L
liutuo 已提交
1447 1448 1449 1450 1451 1452 1453 1454 1455
                        if op.input[0] in self._producer:
                            producer = self._producer[op.input[0]]
                            input_shape = producer.output_shape[0].dims
                            if producer.type == MaceOp.FullyConnected.name and\
                                    len(input_shape) == 2:
                                axis_arg = ConverterUtil.get_arg(
                                    op, MaceKeyword.mace_axis_str)
                                if axis_arg.i == 1:
                                    axis_arg.i = 3
1456

Y
yejianwu 已提交
1457 1458 1459 1460 1461 1462 1463 1464 1465 1466 1467 1468
            elif op.type == MaceOp.Squeeze.name:
                for arg in op.arg:
                    if arg.name == MaceKeyword.mace_axis_str:
                        if (src_data_format == DataFormat.NCHW
                                and has_data_format
                                and len(self._producer[op.input[0]].output_shape[0].dims) == 4  # noqa
                                and len(op.output_shape[0].dims) == 2
                                and arg.ints == [2, 3]):
                            print("Transpose squeeze args: %s(%s)"
                                  % (op.name, op.type))
                            arg.ints[:] = [1, 2]

1469 1470 1471 1472 1473 1474 1475 1476 1477 1478 1479 1480 1481 1482 1483 1484 1485 1486 1487 1488 1489 1490 1491 1492 1493 1494 1495 1496 1497 1498 1499
            elif op.type == MaceOp.Reduce.name:
                for arg in op.arg:
                    if arg.name == MaceKeyword.mace_axis_str:
                        if src_data_format == DataFormat.NCHW and \
                                has_data_format:
                            print("Transpose reduce args: %s(%s)"
                                  % (op.name, op.type))
                            reduce_axises = list(arg.ints)
                            new_axises = []
                            for i in range(len(reduce_axises)):
                                idx = reduce_axises[i]
                                if idx == 2 or idx == 3:
                                    new_axises.append(idx - 1)
                                elif idx == 1:
                                    new_axises.append(3)
                                else:
                                    new_axises.append(idx)
                            new_axises.sort()
                            arg.ints[:] = []
                            arg.ints.extend(new_axises)
            elif op.type == MaceOp.Crop.name:
                offset_arg = ConverterUtil.get_arg(op,
                                                   MaceKeyword.mace_offset_str)
                mace_check(offset_arg and
                           src_data_format == DataFormat.NCHW
                           and has_data_format
                           and len(op.output_shape[0].dims) == 4,
                           "MACE only support crop with NCHW format")
                print("Transpose crop args: %s(%s)"
                      % (op.name, op.type))
                self.transpose_shape(offset_arg.ints, [0, 2, 3, 1])
L
luxuhui 已提交
1500 1501 1502 1503 1504 1505 1506
            elif op.type == MaceOp.Reshape.name:
                for arg in op.arg:
                    if arg.name == MaceKeyword.mace_dim_str and \
                            len(arg.ints) == 4 and \
                            src_data_format == DataFormat.NCHW and \
                            has_data_format:
                        self.transpose_shape(arg.ints, [0, 2, 3, 1])
1507 1508 1509 1510 1511 1512 1513 1514 1515 1516

            # transpose op output shape
            if src_data_format == DataFormat.NCHW and \
                    has_data_format:
                print("Transpose output shapes: %s(%s)" % (op.name, op.type))
                for output_shape in op.output_shape:
                    if len(output_shape.dims) == 4:
                        self.transpose_shape(output_shape.dims,
                                             [0, 2, 3, 1])

1517 1518
        return False

李寅 已提交
1519
    def quantize_nodes(self):
李寅 已提交
1520 1521 1522
        if not self._option.quantize:
            return False

李寅 已提交
1523 1524 1525
        print("Add mace quantize and dequantize nodes")

        for op in self._model.op:
1526
            for i in range(len(op.input)):
B
Bin Li 已提交
1527 1528
                if op.input[i] in self.input_name_map:
                    op.input[i] = self.input_name_map[op.input[i]]
1529
            for i in range(len(op.output)):
B
Bin Li 已提交
1530
                if op.output[i] in self.output_name_map:
B
Bin Li 已提交
1531 1532
                    op.name = MaceKeyword.mace_output_node_name \
                              + '_' + op.name
B
Bin Li 已提交
1533
                    new_output_name = self.output_name_map[op.output[i]]
B
Bin Li 已提交
1534 1535
                    self._quantize_activation_info[new_output_name] = \
                        self._quantize_activation_info[op.output[i]]
1536 1537 1538 1539 1540
                    if op.output[i] in self._consumers:
                        for consumer_op in self._consumers[op.output[i]]:
                            self.replace(consumer_op.input,
                                         op.output[i],
                                         new_output_name)
B
Bin Li 已提交
1541
                    op.output[i] = new_output_name
1542

李寅 已提交
1543 1544 1545 1546 1547 1548
            data_type_arg = ConverterUtil.get_arg(
                op, MaceKeyword.mace_op_data_type_str)
            mace_check(data_type_arg, "Data type does not exist for %s(%s)"
                       % (op.name, op.type))
            if data_type_arg.i == mace_pb2.DT_FLOAT:
                data_type_arg.i = mace_pb2.DT_UINT8
1549 1550 1551 1552 1553
            elif data_type_arg.i == mace_pb2.DT_UINT8:
                mace_check(op.type == MaceOp.Quantize.name
                           or op.type == MaceOp.Dequantize.name,
                           "Only Quantization ops support uint8, "
                           "but got %s(%s)" % (op.name, op.type))
李寅 已提交
1554
            else:
1555
                mace_check(op.type == MaceOp.Quantize.name,
李寅 已提交
1556
                           "Quantization only support float ops, "
B
Bin Li 已提交
1557 1558 1559
                           "but get %s(%s, %s)"
                           % (op.name, op.type,
                              mace_pb2.DataType.Name(data_type_arg.i)))
李寅 已提交
1560

B
Bin Li 已提交
1561
        for i, input_node in enumerate(self._option.input_nodes.values()):
B
Bin Li 已提交
1562
            new_input_name = self.input_name_map[input_node.name]
李寅 已提交
1563
            op_def = self._model.op.add()
B
Bin Li 已提交
1564
            op_def.name = self.normalize_op_name(new_input_name)
李寅 已提交
1565
            op_def.type = MaceOp.Quantize.name
1566
            op_def.input.extend([input_node.name])
B
Bin Li 已提交
1567
            op_def.output.extend([new_input_name])
李寅 已提交
1568 1569
            output_shape = op_def.output_shape.add()
            output_shape.dims.extend(input_node.shape)
B
Bin Li 已提交
1570 1571 1572 1573
            quantize_info = self._quantize_activation_info[new_input_name]
            self.copy_quantize_info(op_def, quantize_info)
            self._model.input_info[i].scale = quantize_info.scale
            self._model.input_info[i].zero_point = quantize_info.zero_point
李寅 已提交
1574 1575

            ConverterUtil.add_data_type_arg(op_def, mace_pb2.DT_UINT8)
1576
            ConverterUtil.add_data_format_arg(op_def, input_node.data_format)
B
Bin Li 已提交
1577 1578 1579 1580 1581
            # use actual ranges for model input quantize
            find_range_every_time_arg = op_def.arg.add()
            find_range_every_time_arg.name = \
                MaceKeyword.mace_find_range_every_time
            find_range_every_time_arg.i = 1
李寅 已提交
1582

B
Bin Li 已提交
1583
        output_nodes = self._option.check_nodes.values()
B
Bin Li 已提交
1584
        for i, output_node in enumerate(output_nodes):
李寅 已提交
1585
            op_def = self._model.op.add()
B
Bin Li 已提交
1586
            op_def.name = self.normalize_op_name(output_node.name)
李寅 已提交
1587
            op_def.type = MaceOp.Dequantize.name
B
Bin Li 已提交
1588
            op_def.input.extend([self.output_name_map[output_node.name]])
1589
            op_def.output.extend([output_node.name])
李寅 已提交
1590
            output_shape = op_def.output_shape.add()
B
Bin Li 已提交
1591 1592
            producer_op = self._producer[output_node.name]
            output_shape.dims.extend(producer_op.output_shape[0].dims)
李寅 已提交
1593
            op_def.output_type.extend([mace_pb2.DT_FLOAT])
B
Bin Li 已提交
1594 1595 1596
            quantize_info = producer_op.quantize_info[0]
            self._model.output_info[i].scale = quantize_info.scale
            self._model.output_info[i].zero_point = quantize_info.zero_point
李寅 已提交
1597 1598

            ConverterUtil.add_data_type_arg(op_def, mace_pb2.DT_UINT8)
1599
            ConverterUtil.add_data_format_arg(op_def, output_node.data_format)
李寅 已提交
1600

1601 1602 1603
        quantize_flag_arg = self._model.arg.add()
        quantize_flag_arg.name = MaceKeyword.mace_quantize_flag_arg_str
        quantize_flag_arg.i = 1
李寅 已提交
1604

李寅 已提交
1605
        return False
李寅 已提交
1606 1607 1608 1609 1610

    def quantize_tensor(self, tensor):
        """Assume biasadd has been already folded with convolution and fc"""
        if tensor.data_type == mace_pb2.DT_FLOAT:
            ops = self._consumers.get(tensor.name, None)
L
liutuo 已提交
1611 1612 1613
            check_conv = False
            check_deconv = False
            if ops is not None and len(ops) == 1:
L
liutuo 已提交
1614 1615 1616 1617 1618 1619
                if len(ops[0].input) >= 3:
                    check_conv =\
                        ops[0].type in [MaceOp.Conv2D.name,
                                        MaceOp.DepthwiseConv2d.name,
                                        MaceOp.FullyConnected.name]\
                        and ops[0].input[2] == tensor.name
L
liutuo 已提交
1620
                # in tensorflow deconv's bias is the forth input
L
liutuo 已提交
1621 1622
                if ops[0].type in [MaceOp.Deconv2D.name,
                                   MaceOp.DepthwiseDeconv2d]:
L
liutuo 已提交
1623 1624 1625 1626
                    from_caffe = ConverterUtil.get_arg(
                        ops[0],
                        MaceKeyword.mace_framework_type_str).i ==\
                                 FrameworkType.CAFFE.value
L
liutuo 已提交
1627 1628
                    if from_caffe and len(ops[0].input) >= 3:
                        check_deconv = ops[0].input[2] == tensor.name
L
liutuo 已提交
1629
                    else:
L
liutuo 已提交
1630 1631
                        if len(ops[0].input) >= 4:
                            check_deconv = ops[0].input[3] == tensor.name
L
liutuo 已提交
1632
            if check_conv or check_deconv:
1633 1634
                if self._option.device == DeviceType.CPU.value \
                       or self._option.device == DeviceType.APU.value:
B
Bin Li 已提交
1635 1636 1637 1638 1639 1640 1641 1642 1643 1644
                    conv_op = ops[0]
                    scale_input = self._quantize_activation_info[
                        conv_op.input[0]].scale
                    if conv_op.input[1] not in self._quantized_tensor:
                        self.quantize_tensor(self._consts[conv_op.input[1]])
                    scale_filter = self._consts[conv_op.input[1]].scale
                    scale = scale_input * scale_filter
                    quantized_tensor = \
                        quantize_util.quantize_with_scale_and_zero(
                            tensor.float_data, scale, 0)
B
Bin Li 已提交
1645 1646
                elif self._option.device == DeviceType.HEXAGON.value or \
                        self._option.device == DeviceType.HTA.value:
B
Bin Li 已提交
1647 1648 1649 1650 1651
                    quantized_tensor = \
                        quantize_util.quantize_bias_for_hexagon(
                            tensor.float_data)
                else:
                    mace_check(False, "wrong device.")
李寅 已提交
1652 1653
                tensor.data_type = mace_pb2.DT_INT32
            else:
1654 1655
                non_zero = self._option.device == DeviceType.CPU.value
                quantized_tensor = quantize_util.quantize(tensor.float_data,
B
Bin Li 已提交
1656
                                                          self._option.device,
1657
                                                          non_zero)
李寅 已提交
1658 1659 1660 1661 1662 1663
                tensor.data_type = mace_pb2.DT_UINT8

            del tensor.float_data[:]
            tensor.int32_data.extend(quantized_tensor.data)
            tensor.scale = quantized_tensor.scale
            tensor.zero_point = quantized_tensor.zero
B
Bin Li 已提交
1664 1665
            tensor.minval = quantized_tensor.minval
            tensor.maxval = quantized_tensor.maxval
李寅 已提交
1666
            tensor.quantized = True
李寅 已提交
1667 1668
            self._quantized_tensor.update([tensor.name])

李寅 已提交
1669 1670
        return False

李寅 已提交
1671 1672 1673 1674 1675
    def quantize_weights(self):
        print("Quantize weights")
        net = self._model
        for tensor in net.tensors:
            self.quantize_tensor(tensor)
李寅 已提交
1676 1677 1678

        return False

B
Bin Li 已提交
1679 1680 1681 1682 1683 1684 1685 1686 1687 1688 1689 1690 1691 1692 1693 1694 1695 1696 1697 1698 1699 1700 1701 1702 1703 1704 1705 1706 1707
    def quantize_large_tensor(self, tensor):
        if tensor.data_type == mace_pb2.DT_FLOAT:
            ops = self._consumers.get(tensor.name, None)
            if ops is not None and len(ops) == 1:
                if ops[0].type in [MaceOp.Conv2D.name,
                                   MaceOp.FullyConnected.name]:
                    quantized_tensor = \
                        quantize_util.quantize(tensor.float_data,
                                               self._option.device,
                                               False)
                    tensor.data_type = mace_pb2.DT_UINT8

                    del tensor.float_data[:]
                    tensor.int32_data.extend(quantized_tensor.data)
                    tensor.scale = quantized_tensor.scale
                    tensor.zero_point = quantized_tensor.zero
                    tensor.minval = quantized_tensor.minval
                    tensor.maxval = quantized_tensor.maxval
                    tensor.quantized = True
                    self._quantized_tensor.update([tensor.name])

    def quantize_large_weights(self):
        print("Quantize large weights")
        net = self._model
        for tensor in net.tensors:
            self.quantize_large_tensor(tensor)

        return False

李寅 已提交
1708
    def add_quantize_info(self, op, minval, maxval):
B
Bin Li 已提交
1709
        scale, zero, minval, maxval = \
B
Bin Li 已提交
1710 1711
            quantize_util.adjust_range(minval, maxval, self._option.device,
                                       non_zero=False)
李寅 已提交
1712 1713 1714 1715 1716 1717 1718 1719
        quantize_info = op.quantize_info.add()
        quantize_info.minval = minval
        quantize_info.maxval = maxval
        quantize_info.scale = scale
        quantize_info.zero_point = zero

        return quantize_info

1720 1721 1722 1723 1724 1725 1726
    def copy_quantize_info(self, op, info):
        quantize_info = op.quantize_info.add()
        quantize_info.minval = info.minval
        quantize_info.maxval = info.maxval
        quantize_info.scale = info.scale
        quantize_info.zero_point = info.zero_point

李寅 已提交
1727 1728 1729 1730 1731 1732 1733 1734 1735 1736 1737 1738
    def transform_fake_quantize(self):
        if not self._option.quantize:
            return False

        # Quantize info from fixpoint fine tune
        print("Transform fake quantize")
        range_file = self._option.quantize_range_file
        if range_file:
            return

        net = self._model
        for op in net.op:
1739 1740
            if op.type == 'FakeQuantWithMinMaxVars' or \
                   op.type == 'FakeQuantWithMinMaxArgs':
李寅 已提交
1741 1742 1743 1744 1745 1746
                producer_op = self._producer[op.input[0]]
                minval = ConverterUtil.get_arg(op, 'min').f
                maxval = ConverterUtil.get_arg(op, 'max').f
                quantize_info = \
                    self.add_quantize_info(producer_op, minval, maxval)
                self._quantize_activation_info[op.input[0]] = quantize_info
1747 1748
                # for add -> fakequant pattern
                self._quantize_activation_info[op.output[0]] = quantize_info
1749 1750

                print(op.input[0], op.output[0])
李寅 已提交
1751 1752 1753 1754
                op.type = MaceOp.Identity.name

        return False

B
Bin Li 已提交
1755 1756 1757 1758 1759 1760 1761 1762 1763 1764 1765 1766 1767 1768 1769 1770 1771 1772 1773 1774 1775 1776 1777 1778 1779 1780 1781 1782 1783 1784 1785 1786 1787 1788 1789 1790 1791 1792 1793 1794 1795 1796 1797 1798 1799 1800 1801 1802 1803 1804
    def rearrange_batch_to_space(self):
        if not self._option.quantize:
            return False

        # Put b2s after biasadd and relu
        for conv_op in self._model.op:
            if conv_op.type in [MaceOp.Conv2D.name,
                                MaceOp.DepthwiseConv2d.name] \
                    and self.consumer_count(conv_op.output[0]) == 1:
                b2s_op = self._consumers[conv_op.output[0]][0]
                if b2s_op.type == MaceOp.BatchToSpaceND.name \
                        and self.consumer_count(b2s_op.output[0]) == 1:
                    biasadd_or_act_op = self._consumers[b2s_op.output[0]][0]
                    if biasadd_or_act_op.type == MaceOp.BiasAdd.name:
                        biasadd_op = biasadd_or_act_op
                        if self.consumer_count(biasadd_op.output[0]) == 1 \
                                and self._consumers[biasadd_op.output[0]][0].type == MaceOp.Activation.name:  # noqa
                            act_op = self._consumers[biasadd_op.output[0]][0]
                            biasadd_op.input[0] = conv_op.output[0]
                            b2s_op.input[0] = act_op.output[0]
                            for op in self._consumers[act_op.output[0]]:
                                self.replace(op.input,
                                             act_op.output[0],
                                             b2s_op.output[0])
                        else:
                            biasadd_op.input[0] = conv_op.output[0]
                            b2s_op.input[0] = biasadd_op.output[0]
                            for op in self._consumers[biasadd_op.output[0]]:
                                self.replace(op.input,
                                             biasadd_op.output[0],
                                             b2s_op.output[0])

                        print("Rearrange batch to space: %s(%s)"
                              % (b2s_op.name, b2s_op.type))
                        return True
                    elif biasadd_or_act_op.type == MaceOp.Activation.name:
                        act_op = biasadd_or_act_op
                        act_op.input[0] = conv_op.output[0]
                        b2s_op.input[0] = act_op.output[0]
                        for op in self._consumers[act_op.output[0]]:
                            self.replace(op.input,
                                         act_op.output[0],
                                         b2s_op.output[0])

                        print("Rearrange batch to space: %s(%s)"
                              % (b2s_op.name, b2s_op.type))
                        return True

        return False

李寅 已提交
1805 1806 1807 1808
    def add_quantize_tensor_range(self):
        # Quantize info from range statistics
        range_file = self._option.quantize_range_file
        if range_file:
李寅 已提交
1809
            print("Add quantize tensor range")
李寅 已提交
1810 1811
            with open(range_file) as f:
                for line in f:
1812
                    tensor_name, minmax = line.split("@@")[:2]
李寅 已提交
1813 1814
                    min_val, max_val = [float(i) for i in
                                        minmax.strip().split(",")]
B
Bin Li 已提交
1815
                    scale, zero, min_val, max_val = \
B
Bin Li 已提交
1816 1817 1818
                        quantize_util.adjust_range(min_val, max_val,
                                                   self._option.device,
                                                   non_zero=False)
李寅 已提交
1819 1820 1821 1822 1823 1824 1825 1826 1827 1828 1829 1830 1831 1832
                    activation_info = mace_pb2.QuantizeActivationInfo()
                    activation_info.minval = min_val
                    activation_info.maxval = max_val
                    activation_info.scale = scale
                    activation_info.zero_point = zero
                    self._quantize_activation_info[tensor_name] = activation_info  # noqa

            for op in self._model.op:
                if op.name.find(MaceKeyword.mace_output_node_name) >= 0:
                    continue
                for output in op.output:
                    mace_check(output in self._quantize_activation_info,
                               "%s does not have quantize activation info"
                               % op)
B
Bin Li 已提交
1833 1834
                    op.quantize_info.extend([
                        self._quantize_activation_info[output]])
李寅 已提交
1835

1836 1837
        if not self._option.quantize:
            return False
B
Bin Li 已提交
1838 1839

        print("Add default quantize info for input")
B
Bin Li 已提交
1840
        for i, input_node in enumerate(self._option.input_nodes.values()):
B
Bin Li 已提交
1841 1842 1843 1844 1845 1846 1847
            if input_node.name not in self._quantize_activation_info:
                print("Input range %s: %s" % (input_node.name,
                                              str(input_node.range)))
                new_input_name = self.input_name_map[input_node.name]
                scale, zero, minval, maxval = \
                    quantize_util.adjust_range(input_node.range[0],
                                               input_node.range[1],
B
Bin Li 已提交
1848
                                               self._option.device,
B
Bin Li 已提交
1849
                                               non_zero=False)
1850 1851
                quantize_info = \
                    mace_pb2.QuantizeActivationInfo()
B
Bin Li 已提交
1852 1853 1854 1855 1856 1857
                quantize_info.minval = minval
                quantize_info.maxval = maxval
                quantize_info.scale = scale
                quantize_info.zero_point = zero
                self._quantize_activation_info[new_input_name] = quantize_info

1858
        print("Add default quantize info for ops like Pooling, Softmax")
李寅 已提交
1859 1860
        for op in self._model.op:
            if op.type in [MaceOp.Pooling.name,
B
Bin Li 已提交
1861
                           MaceOp.Reduce.name,
李寅 已提交
1862
                           MaceOp.Squeeze.name,
B
Bin Li 已提交
1863
                           MaceOp.Reshape.name,
李寅 已提交
1864 1865
                           MaceOp.ResizeBilinear.name,
                           MaceOp.BatchToSpaceND.name,
1866 1867 1868
                           MaceOp.SpaceToBatchND.name,
                           MaceOp.SpaceToDepth.name,
                           MaceOp.DepthToSpace.name]:
李寅 已提交
1869 1870
                del op.quantize_info[:]
                producer_op = self._producer[op.input[0]]
B
Bin Li 已提交
1871 1872 1873 1874 1875
                if producer_op.output[0] in self._option.input_nodes:
                    new_input_name = self.input_name_map[producer_op.output[0]]
                    self.copy_quantize_info(
                        op, self._quantize_activation_info[new_input_name])
                else:
1876 1877
                    self.copy_quantize_info(op,
                                            producer_op.quantize_info[0])
1878 1879
                self._quantize_activation_info[op.output[0]] = \
                    op.quantize_info[0]
1880 1881 1882
            elif (op.type == MaceOp.Concat.name
                  and (not op.quantize_info
                       or self._option.change_concat_ranges)):
1883 1884 1885 1886 1887 1888 1889 1890 1891 1892 1893 1894
                if op.quantize_info:
                    maxval = op.quantize_info[0].maxval
                    minval = op.quantize_info[0].minval
                    del op.quantize_info[:]
                else:
                    maxval = float("-inf")
                    minval = float("inf")
                for i in range(len(op.input)):
                    minval = min(minval, self._producer[op.input[i]].quantize_info[0].minval)  # noqa
                    maxval = max(maxval, self._producer[op.input[i]].quantize_info[0].maxval)  # noqa
                quantize_info = \
                    self.add_quantize_info(op, minval, maxval)
李寅 已提交
1895
                self._quantize_activation_info[op.output[0]] = quantize_info
1896 1897 1898 1899 1900 1901 1902
                if self._option.change_concat_ranges:
                    for i in range(len(op.input)):
                        producer_op = self._producer[op.input[i]]
                        del producer_op.quantize_info[:]
                        self.copy_quantize_info(producer_op, quantize_info)
                        self._quantize_activation_info[producer_op.output[0]] \
                            = producer_op.quantize_info[0]
李寅 已提交
1903 1904 1905 1906 1907
            elif op.type == MaceOp.Softmax.name:
                del op.quantize_info[:]
                quantize_info = \
                    self.add_quantize_info(op, 0.0, 1.0)
                self._quantize_activation_info[op.output[0]] = quantize_info
B
Bin Li 已提交
1908 1909 1910 1911 1912 1913 1914
            elif (op.type == MaceOp.Eltwise.name
                  and not op.quantize_info
                  and len(op.input) == 2
                  and len(op.input[0]) not in self._consts
                  and len(op.input[1]) not in self._consts):
                producer_op0 = self._producer[op.input[0]]
                producer_op1 = self._producer[op.input[1]]
B
Bin Li 已提交
1915 1916 1917 1918 1919 1920 1921 1922 1923 1924 1925 1926 1927 1928 1929 1930
                if ConverterUtil.get_arg(
                        op, MaceKeyword.mace_element_type_str).i \
                        == EltwiseType.SUM.value:
                    minval = producer_op0.quantize_info[0].minval \
                        + producer_op1.quantize_info[0].minval
                    maxval = producer_op0.quantize_info[0].maxval \
                        + producer_op1.quantize_info[0].maxval
                elif ConverterUtil.get_arg(
                        op, MaceKeyword.mace_element_type_str).i \
                        == EltwiseType.SUB.value:
                    minval = producer_op0.quantize_info[0].minval \
                        - producer_op1.quantize_info[0].maxval
                    maxval = producer_op0.quantize_info[0].maxval \
                        - producer_op1.quantize_info[0].minval
                else:
                    mace_check(False, "Quantized Elementwise only support:"
B
Bin Li 已提交
1931
                                      " SUM and SUB without ranges now.")
1932 1933
                quantize_info = \
                    self.add_quantize_info(op, minval, maxval)
B
Bin Li 已提交
1934
                self._quantize_activation_info[op.output[0]] = quantize_info
李寅 已提交
1935 1936 1937 1938 1939 1940 1941

        return False

    def check_quantize_info(self):
        if not self._option.quantize:
            return False

B
Bin Li 已提交
1942
        print("Check quantize info")
李寅 已提交
1943 1944 1945 1946 1947 1948 1949
        for op in self._model.op:
            if (op.name.find(MaceKeyword.mace_input_node_name) == -1
                and op.name.find(MaceKeyword.mace_output_node_name) == -1
                and op.type != MaceOp.Quantize.name
                and op.type != MaceOp.Dequantize.name):  # noqa
                mace_check(len(op.output) == len(op.quantize_info),
                           "missing quantize info: %s" % op)
1950 1951 1952 1953 1954
            for i in six.moves.range(len(op.quantize_info)):
                print("Op output %s range: [%f, %f]" % (
                    op.output[i],
                    op.quantize_info[i].minval,
                    op.quantize_info[i].maxval))
1955

Y
yulianfei 已提交
1956 1957 1958 1959 1960 1961 1962 1963 1964 1965 1966 1967 1968 1969 1970 1971 1972 1973 1974 1975 1976 1977 1978 1979 1980 1981 1982 1983 1984 1985 1986 1987 1988 1989 1990 1991 1992 1993 1994 1995 1996 1997 1998 1999
    def fp16_gather_weight(self):
        for op in self._model.op:
            if op.type != MaceOp.Gather.name:
                continue
            if op.input[0] not in self._consts:
                raise KeyError("Not in const tensor: " + str(op.input[0]))

            const_tensor = self._consts[op.input[0]]
            if const_tensor.data_type == mace_pb2.DT_FLOAT16:
                print(str(const_tensor.name) + " is alreay float16")
                continue

            print("FP16 Embedding Lookup Weights: %s" % const_tensor.name)

            op_outputs = [x for x in op.output]
            new_gather_name = op.name + '_fp16'
            new_gather_output_name = new_gather_name + ":0"
            dehalve_name = op.name

            # fp16 weights
            const_tensor.data_type = mace_pb2.DT_FLOAT16

            # change gather
            op.name = new_gather_name
            op.output[:] = [new_gather_output_name]
            # op.output.extend([new_gather_output_name])
            data_type_arg = ConverterUtil.get_arg(op, MaceKeyword.mace_op_data_type_str)  # noqa
            if data_type_arg is None:
                data_type_arg = op.arg.add()
                data_type_arg.name = MaceKeyword.mace_op_data_type_str
            data_type_arg.i = mace_pb2.DT_FLOAT16

            # add dehalve
            dehalve_op = self._model.op.add()
            dehalve_op.name = dehalve_name
            dehalve_op.type = MaceOp.Cast.name
            dehalve_op.input.extend([new_gather_output_name])
            dehalve_op.output.extend(op_outputs)
            dehalve_op.output_shape.extend(op.output_shape)
            dehalve_op.output_type.extend([mace_pb2.DT_FLOAT])
            data_type_arg = dehalve_op.arg.add()
            data_type_arg.name = MaceKeyword.mace_op_data_type_str
            data_type_arg.i = mace_pb2.DT_FLOAT16

2000 2001 2002 2003
    def fp16_matmul_weight(self):
        if self._option.device != DeviceType.CPU.value:
            return

L
liukai6 已提交
2004
        print('Convert matmul weights to fp16 for specific matmul: activation + weights')  # noqa
2005 2006 2007 2008

        for op in self._model.op:
            if op.type != MaceOp.MatMul.name:
                continue
L
liukai6 已提交
2009
            if op.input[0] not in self._consts and op.input[1] not in self._consts:  # noqa
2010
                continue
L
liukai6 已提交
2011
            if op.input[0] in self._consts and op.input[1] in self._consts:
2012 2013 2014 2015 2016 2017 2018 2019 2020 2021 2022 2023 2024 2025 2026 2027 2028 2029 2030 2031 2032 2033 2034 2035 2036 2037 2038 2039 2040 2041 2042 2043 2044 2045 2046 2047 2048 2049 2050 2051
                continue

            # Matmul fp16 Op only support fp32[1,k] x fp16[w,k]T or fp16[w,k] x fp32[k,1] now!  # noqa

            transpose_a_arg = ConverterUtil.get_arg(op, MaceKeyword.mace_transpose_a_str)  # noqa
            transpose_b_arg = ConverterUtil.get_arg(op, MaceKeyword.mace_transpose_b_str)  # noqa
            transpose_a = transpose_a_arg is not None and transpose_a_arg.i == 1  # noqa
            transpose_b = transpose_b_arg is not None and transpose_b_arg.i == 1  # noqa

            left_tensor = op.input[0]
            right_tensor = op.input[1]
            left_shape = self.get_tensor_shape(left_tensor)
            right_shape = self.get_tensor_shape(right_tensor)

            height = left_shape[-1] if transpose_a else left_shape[-2]
            width = right_shape[-2] if transpose_b else right_shape[-1]
            batch = reduce(lambda x, y: x * y, left_shape[: -2], 1)

            if batch != 1:
                continue

            if left_tensor in self._consts:
                if width != 1 or transpose_a:
                    continue
                const_tensor = self._consts[left_tensor]
            else:
                if height != 1 or not transpose_b:
                    continue
                const_tensor = self._consts[right_tensor]

            print('Convert Matmul Weights to fp16: %s' % op.name)

            const_tensor.data_type = mace_pb2.DT_FLOAT16
            data_type_arg = ConverterUtil.get_arg(op, MaceKeyword.mace_op_data_type_str)  # noqa
            if data_type_arg is None:
                data_type_arg = op.arg.add()
                data_type_arg.name = MaceKeyword.mace_op_data_type_str
            data_type_arg.i = mace_pb2.DT_FLOAT16
            op.output_type.extend([mace_pb2.DT_FLOAT])

2052 2053 2054 2055 2056 2057 2058 2059 2060
    def add_opencl_informations(self):
        print("Add OpenCL informations")

        net = self._model

        arg = net.arg.add()
        arg.name = MaceKeyword.mace_opencl_mem_type
        arg.i = mace_pb2.GPU_IMAGE if self._option.cl_mem_type == "image"\
            else mace_pb2.GPU_BUFFER
B
Bin Li 已提交
2061

L
luxuhui 已提交
2062
    def transform_reshape_and_flatten(self):
L
lichao18 已提交
2063 2064
        net = self._model
        for op in net.op:
L
luxuhui 已提交
2065 2066 2067 2068 2069
            if op.type != MaceOp.Reshape.name:
                continue
            dim_arg = ConverterUtil.get_arg(op, MaceKeyword.mace_dim_str)
            shape_tensor = None
            if len(op.input) == 1:
L
lichao18 已提交
2070
                print("Transform Caffe Reshape")
L
lichao18 已提交
2071 2072 2073 2074 2075
                dims = []
                axis_arg = ConverterUtil.get_arg(op, MaceKeyword.mace_axis_str)
                # transform caffe reshape op
                if dim_arg:
                    dims = dim_arg.ints
L
lichao18 已提交
2076 2077 2078 2079
                    shape_tensor = net.tensors.add()
                    shape_tensor.name = op.name + '_shape'
                    shape_tensor.dims.append(len(op.output_shape[0].dims))
                    shape_tensor.data_type = mace_pb2.DT_INT32
L
lichao18 已提交
2080 2081 2082
                # transform caffe flatten op
                elif axis_arg is not None:
                    axis = axis_arg.i
L
lichao18 已提交
2083
                    for i in range(0, axis):
L
lichao18 已提交
2084 2085 2086 2087
                        dims.append(0)
                    dims.append(-1)
                    for i in range(axis + 1, len(op.output_shape[0].dims)):
                        dims.append(0)
L
lichao18 已提交
2088 2089 2090 2091
                    shape_tensor = net.tensors.add()
                    shape_tensor.name = op.name + '_shape'
                    shape_tensor.dims.append(len(dims))
                    shape_tensor.data_type = mace_pb2.DT_INT32
L
lichao18 已提交
2092 2093 2094 2095
                else:
                    mace_check(False, "Only support reshape and flatten")
                shape_tensor.int32_data.extend(dims)
                op.input.append(shape_tensor.name)
L
luxuhui 已提交
2096 2097 2098 2099 2100 2101 2102
            if len(op.input) == 2 and dim_arg is None:
                if shape_tensor is None and op.input[1] in self._consts:
                    shape_tensor = self._consts[op.input[1]]
                if shape_tensor is not None:
                    dim_arg = op.arg.add()
                    dim_arg.name = MaceKeyword.mace_dim_str
                    dim_arg.ints.extend(shape_tensor.int32_data)
Y
yejianwu 已提交
2103 2104 2105 2106 2107

    def fold_fc_reshape(self):
        net = self._model
        for op in net.op:
            # whether to reshape fc output(default 4D)
L
liutuo 已提交
2108 2109
            if op.type == MaceOp.FullyConnected.name and\
                    op.output[0] in self._consumers:
Y
yejianwu 已提交
2110 2111 2112 2113 2114 2115 2116 2117 2118 2119 2120 2121 2122
                consumers = self._consumers[op.output[0]]
                op_output_shape = op.output_shape[0].dims[:]
                for consumer in consumers:
                    if consumer.type == MaceOp.Reshape.name and \
                            consumer.input[1] in self._consts and \
                            self._consts[consumer.input[1]].int32_data[:] == \
                            [op_output_shape[0], 1, 1, op_output_shape[1]]:
                        # work for tensorflow
                        net.tensors.remove(self._consts[consumer.input[1]])
                        del consumer.input[1]
                        self.safe_remove_node(consumer, None)
                        return True
        return False
李寅 已提交
2123 2124 2125 2126 2127 2128 2129 2130 2131 2132 2133 2134 2135 2136 2137 2138 2139 2140 2141 2142 2143 2144 2145 2146 2147 2148 2149 2150 2151 2152 2153 2154 2155 2156 2157 2158 2159 2160 2161 2162 2163 2164

    def transform_channel_shuffle(self):
        net = self._model
        for op in net.op:
            if op.type == MaceOp.Transpose.name and \
                    len(op.output_shape[0].dims) == 5:
                perm = ConverterUtil.get_arg(op,
                                             MaceKeyword.mace_dims_str).ints
                if [0, 1, 2, 4, 3] == list(perm):
                    # Remove the following Reshape op
                    reshape_op = self._consumers.get(op.output[0], None)
                    if (reshape_op and
                            len(reshape_op) == 1 and
                            reshape_op[0].type == MaceOp.Reshape.name and
                            len(reshape_op[0].output_shape[0].dims) == 4):
                        print("Transform channel shuffle")
                        output_shape = reshape_op[0].output_shape[0].dims
                        self.safe_remove_node(reshape_op[0], op,
                                              remove_input_tensor=True)
                    else:
                        return False

                    # Change Transpose op to ChannelShuffle
                    op.type = MaceOp.ChannelShuffle.name
                    del op.arg[:]
                    group_arg = op.arg.add()
                    group_arg.name = MaceKeyword.mace_group_str
                    group_arg.i = op.output_shape[0].dims[4]
                    op.output_shape[0].dims[:] = output_shape

                    # Remove previous Reshape op
                    producer_op = self._producer.get(op.input[0], None)
                    if producer_op:
                        if producer_op.type == MaceOp.Reshape.name:
                            self.safe_remove_node(producer_op, None)
                        elif producer_op.type == MaceOp.Stack.name:
                            print("Change channel shuffle stack to concat")
                            # Change previous Stack op to Concat if any
                            producer_op.type = MaceOp.Concat.name
                            producer_op.output_shape[0].dims[:] = output_shape

                    return True
2165

2166
    def quantize_specific_ops_only(self):
2167 2168 2169 2170
        """
        This transform rule is only used internally, we are not gonna make
        things too complex for users
        """
2171 2172 2173 2174 2175
        to_quantize_ops_output_type = {
            MaceOp.MatMul.name: mace_pb2.DT_INT32,
            MaceOp.Gather.name: mace_pb2.DT_UINT8,
        }

2176
        for op in self._model.op:
2177 2178
            if (op.type not in to_quantize_ops_output_type
                    or len(op.output) > 1
2179 2180 2181 2182 2183 2184
                    or ConverterUtil.get_arg(op,
                                             MaceKeyword.mace_op_data_type_str).i != mace_pb2.DT_FLOAT):  # noqa
                # only support single output
                continue

            quantized_inputs_names = []
L
liyin 已提交
2185

2186
            should_quantize = False
L
liyin 已提交
2187 2188 2189 2190 2191 2192 2193 2194
            has_const = False
            for idx, input_tensor in enumerate(op.input):
                if input_tensor in self._consts:
                    has_const = True
                    break
            if not has_const:
                continue

2195 2196
            for idx, input_tensor in enumerate(op.input):
                if self.get_tensor_data_type(input_tensor) \
2197 2198
                        == mace_pb2.DT_FLOAT:
                    should_quantize = True
2199 2200 2201
                    break
            if not should_quantize:
                continue
2202 2203
            else:
                print("Quantize op %s (%s)" % (op.name, op.type))
2204

L
liyin 已提交
2205 2206
            non_zero = self._option.device == DeviceType.CPU.value \
                and op.type == MaceOp.MatMul.name
2207 2208 2209 2210

            for idx, input_tensor in enumerate(op.input):
                quantized_inputs_names.append(input_tensor)

2211 2212 2213 2214
                if self.get_tensor_data_type(input_tensor) \
                        != mace_pb2.DT_FLOAT:
                    continue

2215 2216 2217
                if input_tensor in self._consts:
                    const_tensor = self._consts[input_tensor]
                    quantized_tensor = quantize_util.quantize(
B
Bin Li 已提交
2218
                        const_tensor.float_data, self._option.device, non_zero)
2219 2220 2221 2222 2223 2224 2225 2226 2227 2228 2229 2230 2231 2232 2233 2234 2235 2236 2237 2238 2239 2240 2241
                    del const_tensor.float_data[:]
                    const_tensor.int32_data.extend(quantized_tensor.data)
                    const_tensor.data_type = mace_pb2.DT_UINT8
                    const_tensor.scale = quantized_tensor.scale
                    const_tensor.zero_point = quantized_tensor.zero
                    const_tensor.minval = quantized_tensor.minval
                    const_tensor.maxval = quantized_tensor.maxval
                    const_tensor.quantized = True
                else:
                    input_shape = self.get_tensor_shape(input_tensor)
                    quantize_op = self._model.op.add()
                    quantize_op.name = self.normalize_op_name(
                        input_tensor) + "_quant"
                    quantize_op.type = MaceOp.Quantize.name
                    quantize_op.input.extend([input_tensor])
                    quantize_output_name = quantize_op.name + '_0'
                    quantize_op.output.extend([quantize_output_name])
                    output_shape = quantize_op.output_shape.add()
                    output_shape.dims.extend(input_shape)
                    quantize_op.output_type.extend([mace_pb2.DT_UINT8])
                    data_type_arg = quantize_op.arg.add()
                    data_type_arg.name = MaceKeyword.mace_op_data_type_str
                    data_type_arg.i = mace_pb2.DT_UINT8
2242 2243 2244
                    ConverterUtil.add_data_format_arg(
                        quantize_op,
                        self.get_tensor_data_format(input_tensor))
2245 2246 2247

                    data_type_arg = quantize_op.arg.add()
                    data_type_arg.name = MaceKeyword.mace_non_zero
L
liyin 已提交
2248
                    data_type_arg.i = 0
2249 2250 2251 2252 2253 2254 2255 2256 2257 2258 2259

                    find_range_arg = quantize_op.arg.add()
                    find_range_arg.name = \
                        MaceKeyword.mace_find_range_every_time
                    find_range_arg.i = 1

                    quantized_inputs_names[-1] = quantize_output_name

            del op.input[:]
            op.input.extend(quantized_inputs_names)

2260 2261
            original_output_name = op.output[0]
            op.output[0] = original_output_name + "_quant"
2262
            op.output_type.extend([to_quantize_ops_output_type[op.type]])
2263 2264 2265 2266 2267 2268 2269 2270 2271 2272 2273
            data_type_arg = ConverterUtil.get_arg(op,
                                                  MaceKeyword.mace_op_data_type_str)  # noqa
            if data_type_arg is None:
                data_type_arg = op.arg.add()
                data_type_arg.name = MaceKeyword.mace_op_data_type_str
            data_type_arg.i = mace_pb2.DT_UINT8

            dequantize_op = self._model.op.add()
            dequantize_op.name = op.name + "_dequant"
            dequantize_op.type = MaceOp.Dequantize.name
            dequantize_op.input.extend([op.output[0]])
2274
            dequantize_op.output.extend([original_output_name])
2275 2276 2277 2278
            dequantize_op.output_shape.extend(op.output_shape)
            dequantize_op.output_type.extend([mace_pb2.DT_FLOAT])
            data_type_arg = dequantize_op.arg.add()
            data_type_arg.name = MaceKeyword.mace_op_data_type_str
2279
            data_type_arg.i = to_quantize_ops_output_type[op.type]
2280 2281 2282
            ConverterUtil.add_data_format_arg(
                dequantize_op,
                self.get_tensor_data_format(original_output_name))
2283 2284 2285 2286 2287 2288 2289 2290 2291 2292
            quantize_flag_arg = ConverterUtil.get_arg(self._model,
                                                      MaceKeyword.mace_quantize_flag_arg_str)  # noqa
            if quantize_flag_arg is None:
                quantize_flag_arg = self._model.arg.add()
                quantize_flag_arg.name = MaceKeyword.mace_quantize_flag_arg_str
                quantize_flag_arg.i = 1

            return True

        return False