def _convert_layers_batchnorm()

in mmdnn/conversion/tensorflow/tensorflow_frozenparser.py [0:0]


    def _convert_layers_batchnorm(self, source_node):
        IR_node = self.IR_graph.node.add()
        TensorflowParser2._copy_and_reop(source_node, IR_node, 'BatchNorm')

        is_transformed = False
        test = self.get_parent(source_node.name, [0])

        if test.type == 'Mul':
            is_transformed = True

        # ssd model is transformed
        if is_transformed:
            # Ax - (Au - b)

            # A
            input_mul_A = self.get_parent(source_node.name, [0, 1])
            tensor_content = input_mul_A.get_attr('value')
            A_content = tensor_util.MakeNdarray(tensor_content)
            self.set_weight(source_node.name, 'A', A_content)

            # b
            input_sub = self.get_parent(source_node.name, [1])
            tensor_content = input_sub.get_attr('value')
            sub_content = tensor_util.MakeNdarray(tensor_content)
            # print(sub_content)
            self.set_weight(source_node.name, 'b', sub_content)

            input_node = self.get_parent(source_node.name, [0])
            IR_node.input.append(input_node.real_name)
            IR_node.attr["_output_shapes"].list.shape.pop()
            IR_node.attr["_output_shapes"].MergeFromString(input_node.layer.attr['_output_shapes'].SerializeToString())

        else:
            # epsilon
            epsilon = self.get_parent(source_node.name, [1])
            IR_node.attr['epsilon'].f = epsilon.layer.attr['value'].tensor.float_val[0]

            # moving variance (var) /read
            moving_variance = self.get_parent(source_node.name, [0])

            if moving_variance.type == 'Identity':
                moving_variance_read = self.src_graph.get_parent(moving_variance.name, [0])
                tensor_content = moving_variance_read.get_attr('value')
                moving_variance_content = tensor_util.MakeNdarray(tensor_content)
                self.set_weight(source_node.name, 'var', moving_variance_content)

            else:
                print(moving_variance.layer)
                assert False

            # gamma (scale)
            Rsqrt = self.get_son(source_node.name, [0], True)
            # print(Rsqrt.out_edges)

            if len(Rsqrt.out_edges) == 2:
                IR_node.attr['scale'].b = False
                output_node = self.get_son(Rsqrt.name, [0, 0], True)
                if output_node.type == 'Sub':
                    output_node = self.get_son(Rsqrt.name, [1, 0], True)
                    Mul = self.get_son(Rsqrt.name, [0], True)
                else:
                    Mul = self.get_son(Rsqrt.name, [1], True)
            else:
                IR_node.attr['scale'].b = True
                son = self.get_son(Rsqrt.name, [0, 0], True)
                gamma_from = self.get_parent(son.name, [1, 1], True)
                gamma = self.check_const(gamma_from)
                gamma_tensor = gamma.get_attr('value')
                scale = tensor_util.MakeNdarray(gamma_tensor)
                self.set_weight(source_node.name, 'scale', scale)
                output_node = self.get_son(source_node.name, [0, 0, 0, 0], True)
                if output_node.type == 'Sub':
                    output_node = self.get_son(source_node.name, [0, 0, 0, 0, 0], True)
                    Mul = self.get_son(Rsqrt.name, [0, 0], True)
                else:
                    Mul = self.get_son(Rsqrt.name, [0, 1], True)

            # beta  (bias)
            beta = self.get_parent(output_node.name, [1, 0, 0], True).get_attr('value')
            bias = tensor_util.MakeNdarray(beta)
            IR_node.attr['bias'].b = True
            self.set_weight(source_node.name, 'bias', bias)

            # moving mean (mean)
            moving_mean = self.get_parent(Mul.name, [0, 0]).get_attr('value')
            mean = tensor_util.MakeNdarray(moving_mean)
            self.set_weight(source_node.name, 'mean', mean)

            # input node
            assert output_node.type == 'Add'
            input_node = self.get_parent(output_node.name, [0, 0])
            IR_node.input.append(input_node.real_name)
            IR_node.attr["_output_shapes"].list.shape.pop()
            IR_node.attr["_output_shapes"].MergeFromString(input_node.layer.attr['_output_shapes'].SerializeToString())
            output_node.real_name = source_node.name