Method fold_instance_norm

tools/python/transform/transformer.py:3306–3413

(self)

        epsilon_arg.f = epsilon

    def fold_instance_norm(self):
        net = self._model
        for op in net.op:
            is_reduce = (op.type == MaceOp.Reduce.name and
                         len(op.input) == 1 and
                         len(op.output) == 1)
            if not is_reduce:
                continue
            reduce_type_arg = ConverterUtil.get_arg(
                op, MaceKeyword.mace_reduce_type_str)
            reduce_type = reduce_type_arg.i
            axis_arg = ConverterUtil.get_arg(op, MaceKeyword.mace_axis_str)
            axis = axis_arg.ints
            keepdims_arg = ConverterUtil.get_arg(
                op, MaceKeyword.mace_keepdims_str)
            keepdims = keepdims_arg.i
            if not (reduce_type == ReduceType.MEAN.value and
                    len(axis) == 2 and
                    axis[0] == 1 and axis[1] == 2 and
                    keepdims == 1 and
                    len(op.output_shape[0].dims) == 4):
                continue
            # Ops that take Mean as input
            mean_consumers = self._consumers.get(op.output[0], [])
            if len(mean_consumers) != 2:
                continue
            sqr_diff_mean_idx = -1
            sqr_diff_mean_op = None
            for idx in range(2):
                if mean_consumers[idx].type == MaceOp.SqrDiffMean.name:
                    sqr_diff_mean_idx = idx
                    sqr_diff_mean_op = mean_consumers[idx]
                    break
            if sqr_diff_mean_idx == -1:
                # SqrDiffMean is not found, it is not InstanceNorm
                continue
            second_mean_consumer_op = mean_consumers[1 - sqr_diff_mean_idx]
            if second_mean_consumer_op.type != MaceOp.Eltwise.name:
                continue
            elt_type = ConverterUtil.get_arg(
                second_mean_consumer_op, MaceKeyword.mace_element_type_str).i
            scale_offset = False
            # second consumer of Mean can only be NEG or PROD, otherwise,
            # it's not InstanceNorm
            if elt_type == EltwiseType.PROD.value:
                scale_offset = True
            elif elt_type == EltwiseType.NEG.value:
                scale_offset = False
            else:
                continue
            # var + epsilon
            sqr_diff_mean_consumers = self._consumers.get(
                sqr_diff_mean_op.output[0], [])
            if len(sqr_diff_mean_consumers) != 1:
                continue
            var_plus_epsilon_op = sqr_diff_mean_consumers[0]
            if var_plus_epsilon_op.type != MaceOp.Eltwise.name:
                continue