Example #1
0
class TestDPUContrib(unittest.TestCase):

    xgraph_partitioner = XGraphPartitioner()
    xgraph_factory = XGraphFactory()
    target_registry = TargetRegistry()
    rt_manager = RtManager()

    @classmethod
    def setUpClass(cls):
        # Import DPU module
        from pyxir.contrib.dpuv1 import dpuv1
        # from pyxir.contrib.dpuv1.dpuv1_target import\
        #     xgraph_dpu_v1_optimizer,\
        #     xgraph_dpu_v1_quantizer,\
        #     xgraph_dpu_v1_compiler,\
        #     xgraph_dpu_v1_build_func

        # pyxir.register_target(
        #     'dpuv1',
        #     xgraph_dpu_v1_optimizer,
        #     xgraph_dpu_v1_quantizer,
        #     xgraph_dpu_v1_compiler,
        #     xgraph_dpu_v1_build_func
        # )

    @classmethod
    def tearDownClass(cls):
        # Unregister dpu for other tests
        TestDPUContrib.target_registry.unregister_target('dpuv1')
        TestDPUContrib.target_registry.unregister_target('DPUCADX8G')

    def test_supported_ops(self):
        dpuv1_ops = TestDPUContrib.target_registry\
            .get_supported_op_check_names('dpuv1')

        assert 'BatchNorm' in dpuv1_ops
        assert 'BiasAdd' in dpuv1_ops
        assert 'Concat' in dpuv1_ops
        assert 'Convolution' in dpuv1_ops
        assert 'Conv2DTranspose' in dpuv1_ops
        assert 'DPU' in dpuv1_ops
        assert 'Eltwise' in dpuv1_ops
        assert 'Pad' in dpuv1_ops
        assert 'Pooling' in dpuv1_ops
        assert 'Mean' in dpuv1_ops
        assert 'pReLU' in dpuv1_ops
        assert 'ReLU' in dpuv1_ops
        assert 'Scale' in dpuv1_ops

    @unittest.skipIf(skip_tf,
                     "Skipping Tensorflow related test because tensorflow is"
                     "not available")
    def test_import_ext_quantizer(self):
        if TestDPUContrib.target_registry.is_target('DPUCADX8G'):
            TestDPUContrib.target_registry.unregister_target('DPUCADX8G')
        if TestDPUContrib.rt_manager.exists_op('cpu-np', 'DPU'):
            TestDPUContrib.rt_manager.unregister_op('cpu-np', 'DPU')
        from pyxir.contrib.target import DPUCADX8G_external_quantizer
Example #2
0
class TestDPUContrib(unittest.TestCase):

    xgraph_partitioner = XGraphPartitioner()
    xgraph_factory = XGraphFactory()
    target_registry = TargetRegistry()

    @classmethod
    def setUpClass(cls):
        # Import DPU module
        from pyxir.contrib.dpuv1 import dpuv1
        # from pyxir.contrib.dpuv1.dpuv1_target import\
        #     xgraph_dpu_v1_optimizer,\
        #     xgraph_dpu_v1_quantizer,\
        #     xgraph_dpu_v1_compiler,\
        #     xgraph_dpu_v1_build_func

        # pyxir.register_target(
        #     'dpuv1',
        #     xgraph_dpu_v1_optimizer,
        #     xgraph_dpu_v1_quantizer,
        #     xgraph_dpu_v1_compiler,
        #     xgraph_dpu_v1_build_func
        # )

    @classmethod
    def tearDownClass(cls):
        # Unregister dpu for other tests
        TestDPUContrib.target_registry.unregister_target('dpuv1')
        TestDPUContrib.target_registry.unregister_target('DPUCADX8G')

    def test_supported_ops(self):
        dpuv1_ops = TestDPUContrib.target_registry\
            .get_supported_op_check_names('dpuv1')

        assert 'BatchNorm' in dpuv1_ops
        assert 'BiasAdd' in dpuv1_ops
        assert 'Concat' in dpuv1_ops
        assert 'Convolution' in dpuv1_ops
        assert 'Conv2DTranspose' in dpuv1_ops
        assert 'DPU' in dpuv1_ops
        assert 'Eltwise' in dpuv1_ops
        assert 'Pad' in dpuv1_ops
        assert 'Pooling' in dpuv1_ops
        assert 'Mean' in dpuv1_ops
        assert 'pReLU' in dpuv1_ops
        assert 'ReLU' in dpuv1_ops
        assert 'Scale' in dpuv1_ops
Example #3
0
class TestLayoutTransformationPass(unittest.TestCase):

    xgraph_partitioner = XGraphPartitioner()
    xgraph_factory = XGraphFactory()

    @classmethod
    def setUpClass(cls):
        def xgraph_build_func(xgraph):
            raise NotImplementedError("")

        def xgraph_optimizer(xgraph):
            raise NotImplementedError("")

        def xgraph_quantizer(xgraph):
            raise NotImplementedError("")

        def xgraph_compiler(xgraph):
            raise NotImplementedError("")

        target_registry = TargetRegistry()
        target_registry.register_target('npu_test', xgraph_optimizer,
                                        xgraph_quantizer, xgraph_compiler,
                                        xgraph_build_func)

        @register_op_support_check('npu_test', 'Convolution')
        def conv_op_support(X, bXs, tXs):
            return True

    @classmethod
    def tearDownClass(cls):

        target_registry = TargetRegistry()
        target_registry.unregister_target('npu_test')

    def test_simple(self):
        net = [
            XLayer(name='in1',
                   type=['Input'],
                   shapes=[1, 1, 4, 4],
                   sizes=[16],
                   bottoms=[],
                   tops=['conv1'],
                   layer=['in1'],
                   targets=[]),
            XLayer(name='conv1',
                   type=['Convolution'],
                   shapes=[1, 2, 3, 3],
                   sizes=[18],
                   bottoms=['in1'],
                   tops=[],
                   layer=['conv1'],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   attrs={
                       'data_layout': 'NCHW',
                       'padding': [[0, 0], [0, 0], [1, 1], [1, 1]]
                   },
                   targets=[])
        ]
        xgraph = TestLayoutTransformationPass.xgraph_factory\
            .build_from_xlayer(net)

        layout_transform_pass = XGraphLayoutTransformationPass('NHWC')
        new_xgraph = layout_transform_pass.execute(xgraph)

        xlayers = new_xgraph.get_layers()
        # print(xlayers)
        assert len(new_xgraph) == 4
        assert xlayers[0].type[0] == 'Input'
        assert xlayers[1].type[0] == 'Transpose'
        assert xlayers[2].type[0] == 'Convolution'
        assert xlayers[3].type[0] == 'Transpose'

        assert xlayers[0].bottoms == []
        assert xlayers[0].tops == ['conv1_bottom_NCHW>NHWC']
        assert xlayers[0].shapes == [1, 1, 4, 4]
        assert xlayers[1].bottoms == ['in1']
        assert xlayers[1].tops == ['conv1']
        assert xlayers[1].shapes == [1, 4, 4, 1]
        assert xlayers[2].bottoms == ['conv1_bottom_NCHW>NHWC']
        assert xlayers[2].tops == ['conv1_top_NHWC>NCHW']
        assert xlayers[2].shapes == [1, 3, 3, 2]
        assert xlayers[3].bottoms == ['conv1']
        assert xlayers[3].tops == []
        assert xlayers[3].shapes == [1, 2, 3, 3]

        # NCHW -> NHWC
        assert xlayers[1].attrs['axes'] == [0, 2, 3, 1]
        # NHWC -> NCHW
        assert xlayers[3].attrs['axes'] == [0, 3, 1, 2]

        assert xlayers[2].attrs['data_layout'] == 'NHWC'

    def test_complex(self):
        net = [
            XLayer(name='in1',
                   type=['Input'],
                   shapes=[1, 1, 4, 4],
                   sizes=[16],
                   bottoms=[],
                   tops=['conv1'],
                   layer=['in1'],
                   targets=[]),
            XLayer(name='conv1',
                   type=['Convolution'],
                   shapes=[1, 2, 3, 3],
                   sizes=[18],
                   bottoms=['in1'],
                   tops=['pool1'],
                   layer=['conv1'],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   attrs={
                       'data_layout': 'NCHW',
                       'padding': [[0, 0], [0, 0], [1, 1], [1, 1]]
                   },
                   targets=[]),
            XLayer(name='pool1',
                   type=['Pooling'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=['conv1'],
                   tops=['concat1'],
                   layer=['pool1'],
                   attrs={
                       'data_layout': 'NCHW',
                       'padding': [[0, 0], [0, 0], [1, 1], [1, 1]]
                   },
                   targets=[]),
            XLayer(name='in2',
                   type=['Input'],
                   shapes=[1, 4, 4, 1],
                   sizes=[16],
                   bottoms=[],
                   tops=['in2_transpose'],
                   layer=['in2'],
                   targets=[]),
            XLayer(name='in2_transpose',
                   type=['Transpose'],
                   shapes=[1, 1, 4, 4],
                   sizes=[16],
                   bottoms=['in2'],
                   tops=['conv2'],
                   layer=['in2_transpose'],
                   attrs={'axes': [0, 3, 1, 2]},
                   targets=[]),
            XLayer(name='conv2',
                   type=['Convolution'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=['in2_transpose'],
                   tops=['concat1'],
                   layer=['conv2'],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   attrs={
                       'data_layout': 'NCHW',
                       'padding': [[0, 0], [0, 0], [1, 1], [1, 1]]
                   },
                   targets=[]),
            XLayer(name='concat1',
                   type=['Concat'],
                   shapes=[1, 4, 2, 2],
                   sizes=[16],
                   bottoms=['pool1', 'conv2'],
                   tops=['concat1_transpose'],
                   layer=['concat1'],
                   attrs={'axis': 1},
                   targets=[]),
            XLayer(name='concat1_transpose',
                   type=['Transpose'],
                   shapes=[1, 2, 2, 4],
                   sizes=[16],
                   bottoms=['concat1'],
                   tops=['dense1'],
                   layer=['concat1_transpose'],
                   attrs={'axes': [0, 2, 3, 1]},
                   targets=[]),
            XLayer(name='dense1',
                   type=['Dense'],
                   shapes=[1, 20],
                   sizes=[],
                   bottoms=['concat1_transpose'],
                   tops=[],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   layer=['dense1'],
                   targets=[])
        ]
        xgraph = TestLayoutTransformationPass.xgraph_factory\
            .build_from_xlayer(net)

        layout_transform_pass = XGraphLayoutTransformationPass('NHWC')
        new_xgraph = layout_transform_pass.execute(xgraph)

        xlayers = new_xgraph.get_layers()
        # print(xlayers)
        # print(len(xlayers))
        assert len(new_xgraph) == 8

        assert xlayers[0].type[0] == 'Input'
        assert xlayers[0].name == 'in1'
        assert xlayers[0].bottoms == []
        assert xlayers[0].tops == ['conv1_bottom_NCHW>NHWC']
        assert xlayers[0].shapes == [1, 1, 4, 4]

        assert xlayers[1].type[0] == 'Transpose'
        assert xlayers[1].name == 'conv1_bottom_NCHW>NHWC'
        assert xlayers[1].bottoms == ['in1']
        assert xlayers[1].tops == ['conv1']
        assert xlayers[1].shapes == [1, 4, 4, 1]
        assert xlayers[1].attrs['axes'] == [0, 2, 3, 1]

        assert xlayers[2].type[0] == 'Convolution'
        assert xlayers[2].name == 'conv1'
        assert xlayers[2].bottoms == ['conv1_bottom_NCHW>NHWC']
        assert xlayers[2].tops == ['pool1']
        assert xlayers[2].shapes == [1, 3, 3, 2]
        assert xlayers[2].attrs['data_layout'] == 'NHWC'

        assert xlayers[3].type[0] == 'Pooling'
        assert xlayers[3].name == 'pool1'
        assert xlayers[3].bottoms == ['conv1']
        assert xlayers[3].tops == ['concat1']
        assert xlayers[3].shapes == [1, 2, 2, 2]
        assert xlayers[3].attrs['data_layout'] == 'NHWC'

        assert xlayers[4].type[0] == 'Input'
        assert xlayers[4].name == 'in2'
        assert xlayers[4].bottoms == []
        assert xlayers[4].tops == ['conv2']
        assert xlayers[4].shapes == [1, 4, 4, 1]

        assert xlayers[5].type[0] == 'Convolution'
        assert xlayers[5].name == 'conv2'
        assert xlayers[5].bottoms == ['in2']
        assert xlayers[5].tops == ['concat1']
        assert xlayers[5].shapes == [1, 2, 2, 2]
        assert xlayers[5].attrs['data_layout'] == 'NHWC'

        assert xlayers[6].type[0] == 'Concat'
        assert xlayers[6].name == 'concat1'
        assert xlayers[6].bottoms == ['pool1', 'conv2']
        assert xlayers[6].tops == ['dense1']
        assert xlayers[6].shapes == [1, 2, 2, 4]
        assert xlayers[6].attrs['axis'] == 3

        assert xlayers[7].type[0] == 'Dense'
        assert xlayers[7].name == 'dense1'
        assert xlayers[7].bottoms == ['concat1']
        assert xlayers[7].tops == []
        assert xlayers[7].shapes == [1, 20]

    def test_target(self):
        net = [
            XLayer(name='in1',
                   type=['Input'],
                   shapes=[1, 1, 4, 4],
                   sizes=[16],
                   bottoms=[],
                   tops=['conv1'],
                   layer=['in1'],
                   targets=[]),
            XLayer(name='conv1',
                   type=['Convolution'],
                   shapes=[1, 2, 3, 3],
                   sizes=[18],
                   bottoms=['in1'],
                   tops=['pool1'],
                   layer=['conv1'],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   attrs={
                       'data_layout': 'NCHW',
                       'padding': [[0, 0], [0, 0], [1, 1], [1, 1]]
                   },
                   targets=[],
                   target='test'),
            XLayer(name='pool1',
                   type=['Pooling'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=['conv1'],
                   tops=['concat1'],
                   layer=['pool1'],
                   attrs={
                       'data_layout': 'NCHW',
                       'padding': [[0, 0], [0, 0], [1, 1], [1, 1]]
                   },
                   targets=[]),
            XLayer(name='in2',
                   type=['Input'],
                   shapes=[1, 4, 4, 1],
                   sizes=[16],
                   bottoms=[],
                   tops=['in2_transpose'],
                   layer=['in2'],
                   targets=[]),
            XLayer(name='in2_transpose',
                   type=['Transpose'],
                   shapes=[1, 1, 4, 4],
                   sizes=[16],
                   bottoms=['in2'],
                   tops=['conv2'],
                   layer=['in2_transpose'],
                   attrs={'axes': [0, 3, 1, 2]},
                   targets=[]),
            XLayer(name='conv2',
                   type=['Convolution'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=['in2_transpose'],
                   tops=['concat1'],
                   layer=['conv2'],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   attrs={
                       'data_layout': 'NCHW',
                       'padding': [[0, 0], [0, 0], [1, 1], [1, 1]]
                   },
                   targets=[],
                   target='test'),
            XLayer(name='concat1',
                   type=['Concat'],
                   shapes=[1, 4, 2, 2],
                   sizes=[16],
                   bottoms=['pool1', 'conv2'],
                   tops=['concat1_transpose'],
                   layer=['concat1'],
                   attrs={'axis': 1},
                   targets=[]),
            XLayer(name='concat1_transpose',
                   type=['Transpose'],
                   shapes=[1, 2, 2, 4],
                   sizes=[16],
                   bottoms=['concat1'],
                   tops=['dense1'],
                   layer=['concat1_transpose'],
                   attrs={'axes': [0, 2, 3, 1]},
                   targets=[]),
            XLayer(name='dense1',
                   type=['Dense'],
                   shapes=[1, 20],
                   sizes=[],
                   bottoms=['concat1_transpose'],
                   tops=[],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   layer=['dense1'],
                   targets=[])
        ]
        xgraph = TestLayoutTransformationPass.xgraph_factory\
            .build_from_xlayer(net)

        layout_transform_pass = XGraphLayoutTransformationPass('NHWC',
                                                               target='test')
        new_xgraph = layout_transform_pass.execute(xgraph)

        xlayers = new_xgraph.get_layers()
        # print(xlayers)
        # print(len(xlayers))
        assert len(new_xgraph) == 10

        assert xlayers[0].type[0] == 'Input'
        assert xlayers[0].name == 'in1'
        assert xlayers[0].bottoms == []
        assert xlayers[0].tops == ['conv1_bottom_NCHW>NHWC']
        assert xlayers[0].shapes == [1, 1, 4, 4]

        assert xlayers[1].type[0] == 'Transpose'
        assert xlayers[1].name == 'conv1_bottom_NCHW>NHWC'
        assert xlayers[1].bottoms == ['in1']
        assert xlayers[1].tops == ['conv1']
        assert xlayers[1].shapes == [1, 4, 4, 1]
        assert xlayers[1].attrs['axes'] == [0, 2, 3, 1]

        assert xlayers[2].type[0] == 'Convolution'
        assert xlayers[2].name == 'conv1'
        assert xlayers[2].bottoms == ['conv1_bottom_NCHW>NHWC']
        assert xlayers[2].tops == ['conv1_top_NHWC>NCHW']
        assert xlayers[2].shapes == [1, 3, 3, 2]
        assert xlayers[2].attrs['data_layout'] == 'NHWC'
        assert xlayers[2].attrs['padding'] == [[0, 0], [1, 1], [1, 1], [0, 0]]

        assert xlayers[3].type[0] == 'Transpose'
        assert xlayers[3].name == 'conv1_top_NHWC>NCHW'
        assert xlayers[3].bottoms == ['conv1']
        assert xlayers[3].tops == ['pool1']
        assert xlayers[3].shapes == [1, 2, 3, 3]
        assert xlayers[3].attrs['axes'] == (0, 3, 1, 2)

        assert xlayers[4].type[0] == 'Pooling'
        assert xlayers[4].name == 'pool1'
        assert xlayers[4].bottoms == ['conv1_top_NHWC>NCHW']
        assert xlayers[4].tops == ['0_split_concat1_transpose']
        assert xlayers[4].shapes == [1, 2, 2, 2]
        assert xlayers[4].attrs['data_layout'] == 'NCHW'
        assert xlayers[4].attrs['padding'] == [[0, 0], [0, 0], [1, 1], [1, 1]]

        assert xlayers[5].type[0] == 'Transpose'
        assert xlayers[5].name == '0_split_concat1_transpose'
        assert xlayers[5].bottoms == ['pool1']
        assert xlayers[5].tops == ['concat1']
        assert xlayers[5].shapes == [1, 2, 2, 2]
        assert xlayers[5].attrs['axes'] == [0, 2, 3, 1]

        assert xlayers[6].type[0] == 'Input'
        assert xlayers[6].name == 'in2'
        assert xlayers[6].bottoms == []
        assert xlayers[6].tops == ['conv2']
        assert xlayers[6].shapes == [1, 4, 4, 1]

        assert xlayers[7].type[0] == 'Convolution'
        assert xlayers[7].name == 'conv2'
        assert xlayers[7].bottoms == ['in2']
        assert xlayers[7].tops == ['concat1']
        assert xlayers[7].shapes == [1, 2, 2, 2]
        assert xlayers[7].attrs['data_layout'] == 'NHWC'

        assert xlayers[8].type[0] == 'Concat'
        assert xlayers[8].name == 'concat1'
        assert xlayers[8].bottoms == ['0_split_concat1_transpose', 'conv2']
        assert xlayers[8].tops == ['dense1']
        assert xlayers[8].shapes == [1, 2, 2, 4]
        assert xlayers[8].attrs['axis'] == 3

        assert xlayers[9].type[0] == 'Dense'
        assert xlayers[9].name == 'dense1'
        assert xlayers[9].bottoms == ['concat1']
        assert xlayers[9].tops == []
        assert xlayers[9].shapes == [1, 20]
Example #4
0
class TfGenerator(object):
    """
    Responsible for generating tensorflow model from xgraph data structure
    """

    runtime_factory = RuntimeFactory()
    xgraph_factory = XGraphFactory()
    xgraph_partitioner = XGraphPartitioner()

    @classmethod
    def generate(cls,
                 xgraph,
                 base_name,
                 subgraphs_only=False,
                 layout='NCHW',
                 batch_size=-1,
                 placeholder=False,
                 out_dir=os.getcwd()):
        # type: (XGraph, str, boolean, str, int) -> Dict[str, str]
        """
        Generate one or multiple tensorflow pb file from an xgraph and
        return dictionary of the base_name/partitions mapping to the pb files
        """
        # layout_transform_pass = XGraphLayoutTransformationPass(layout)
        # xgraph = layout_transform_pass.execute(xgraph, subgraphs_only=False)

        # Import tensorflow only when needed
        import tensorflow as tf

        executors = []
        if not subgraphs_only:
            executors.append((base_name, base_name,
                              TfGenerator.runtime_factory.build_runtime(
                                  xgraph,
                                  batch_size=batch_size,
                                  placeholder=placeholder)))
        else:
            for Xp in \
                    TfGenerator.xgraph_partitioner.get_subgraphs(xgraph):

                sub_xgraph = TfGenerator.xgraph_factory.build_from_xlayer(
                    Xp.subgraph_data)
                executors.append((base_name + '_' + Xp.name, Xp.name,
                                  TfGenerator.runtime_factory.build_runtime(
                                      sub_xgraph,
                                      batch_size=batch_size,
                                      placeholder=placeholder),
                                  sub_xgraph.get_output_names()))

        ret = {}
        for file_name, name, executor, output_names in executors:
            graph_def = executor.tf_graph.as_graph_def()
            # with executor.tf_step_graph.as_default():
            #    graph_def = tf.get_default_graph().as_graph_def()

            with tf.compat.v1.Session(graph=executor.tf_graph) as sess:
                sess.run(tf.compat.v1.global_variables_initializer())
                graph_def = tf.graph_util.convert_variables_to_constants(
                    sess, graph_def, output_names)

            file_path = os.path.join(out_dir, file_name + '.pb')

            with tf.gfile.GFile(file_path, "wb") as f:
                f.write(graph_def.SerializeToString())

            ret[name] = file_path

        return ret

    @classmethod
    def run(self, xgraph, pb_file, inputs):
        # type: (XGraph, str, dict[str, numpy.ndarray])
        """ Run frozen tensorflow graph for corresponding XGraph """

        # Import Tensorflow only when needed
        import tensorflow as tf

        input_graph_def = tf.GraphDef()
        with tf.gfile.GFile(frozen_graph, "rb") as f:
            input_graph_def.ParseFromString(f.read())

        tf.compat.v1.reset_default_graph()
        tf.import_graph_def(input_graph_def, name='')

        input_names = xgraph.get_input_names()
        output_names = xgraph.get_output_names()

        inputs_tf = {}
        for in_name in input_names:
            input_tensor = tf.get_default_graph().get_tensor_by_name(in_name +
                                                                     ':0')

            inputs_tf[input_tensor] = inputs[in_name]

        outputs = [
            tf.get_default_graph().get_tensor_by_name(out_name + ':0')
            for out_name in output_names
        ]

        with tf.compat.v1.Session() as sess:
            res = sess.run(outputs, feed_dict=inputs_tf)

        return res
Example #5
0
class VAICompiler(XGraphBaseCompiler):

    """ Vitis-AI compiler wrapper for DPUCZDX8G """

    xgraph_partitioner = XGraphPartitioner()
    xgraph_factory = XGraphFactory()

    def __init__(self,
                 xgraph,
                 arch,
                 meta,
                 dcf,
                 cpu_arch='arm64',
                 work_dir=os.path.join(os.getcwd(), 'work'),
                 build_dir=os.getcwd(),
                 mode='debug'):

        super(VAICompiler, self).__init__(xgraph)

        if not os.path.isfile(arch):
            raise ValueError("Arch file: {} does not exist".format(arch))

        if cpu_arch != 'arm64':
            raise ValueError("Unsupported CPU architecture: {}. Supported"
                             " architectures are: 'arm64'")

        q_output = self.xgraph.get_quantizer_output()
        self.netcfgs = {q_key: q_output.get_q_file(q_key)
                        for q_key in q_output.keys()}
        assert(len(self.netcfgs) == 1)
        self.arch = arch
        self.meta = meta
        self.dcf = dcf
        self.cpu_arch = cpu_arch
        self.work_dir = work_dir

        if not os.path.exists(self.work_dir):
            os.makedirs(self.work_dir)
        self.build_dir = build_dir if build_dir is not None else work_dir
        if not os.path.exists(self.build_dir):
            os.makedirs(self.build_dir)
        self.mode = mode
        self.c_output = CompilerOutput(name=xgraph.get_name())
        

    def compile(self) -> None:
        """ Start DPUv2 compilation """

        net_name = list(self.netcfgs.keys())[0]
        netcfg = list(self.netcfgs.values())[0]

        # We only handle one partition at the moment
        Xp = VAICompiler.xgraph_partitioner\
            .get_subgraphs(self.xgraph)[0]
        subxg_layers = Xp.subgraph_data
        xgraph = VAICompiler.xgraph_factory.build_from_xlayer(subxg_layers)
        # assert xgraph.get_name() == net_name

        input_names = xgraph.get_input_names()
        input_shapes = [xgraph.get(in_name).shapes[:]
                        for in_name in input_names]
        output_names = list(Xp.attrs['__top_tensors'].keys()) # xgraph.get_output_names()
        output_shapes = [xgraph.get(out_name).shapes[:]
                         for out_name in output_names]

        if len(input_names) > 1:
            raise NotImplementedError("VAICompiler only handles models with"
                                      " one input at the moment but found: {}"
                                      .format(len(input_names)))

        #command = """
        #vai_c_tensorflow \
        #    --frozen_pb {} \
        #    --arch {} \
        #    --output_dir {} \
        #    --net_name {} \
        #    --options "{}"
        #""".format(netcfg, self.arch, self.work_dir, net_name, str(dict()))
        # import pdb; pdb.set_trace()

        command = """
        dnnc-dpuv2 --parser tensorflow\
            --frozen_pb {} \
            --cpu_arch {} \
            --output_dir {} \
            --net_name {} \
            --dcf {}
        """.format(netcfg, self.cpu_arch, self.work_dir, net_name, self.dcf)


        logger.info("Command: {}".format(command))

        process = subprocess.Popen(command,
                                   shell=True,
                                   cwd=FILE_PATH,
                                   stdout=subprocess.PIPE)

        output, error = process.communicate()
        logger.debug("{} {}".format(output, error))

        if output is not None:
            output = output.decode('utf-8')

            logger.info("Output: {}".format(output))
            logger.info("Output names: {}".format(output_names))

            do = DNNCOutput(str(repr(output)))

            dpu_input_nodes = do.get_input_nodes()
            dpu_output_nodes = do.get_output_nodes()
            dpu_output_nodes_on_shapes = do.get_output_nodes_on_shapes()

            in_shapes_log = ["{}*{}*{}".format(ishape[1], ishape[2], ishape[3])
                             for ishape in input_shapes]
            out_shapes_log = ["{}*{}*{}".format(os[1], os[2], os[3])
                              for os in output_shapes]

            in_map = {in_name: in_name + ':0' for in_name, _ in zip(input_names, in_shapes_log)}
            out_map = {}

            for out_name, out_shape_str in zip(output_names, out_shapes_log):
                # DNNC changes naming
                dnnc_out_name = do.get_dnnc_str(out_name)
                if dnnc_out_name in dpu_output_nodes:
                    out_map[out_name] = dpu_output_nodes[dnnc_out_name]
                # out_name: dpu_output_nodes[out_shape_str] + ':0'
                else:
                    assert len(dpu_output_nodes_on_shapes) == len(output_names),\
                        "Can't retrieve right out tensor names from DNNC compiler output"
                    out_map[out_name] = dpu_output_nodes_on_shapes[out_shape_str]

            logger.info("DPU kernel in_map: {}".format(in_map))
            logger.info("DPU kernel out_map: {}".format(out_map))

        if error is not None:
            error = error.decode('utf-8')
            raise ValueError(error)

        logger.info("VAI_C Output: {}".format(output))
        logger.info("VAI_C Error: {}".format(error))

        logger.debug("CROSS COMPILATION")
        command = """
        aarch64-linux-gnu-gcc -fPIC -shared {}/dpu_{}.elf -o {}/libdpumodel{}.so
        """.format(self.work_dir, net_name, self.work_dir, net_name)

        logger.debug("Command: {}".format(command))

        process = subprocess.Popen(command.split(),
                                   cwd=FILE_PATH,
                                   stdout=subprocess.PIPE)
        output, error = process.communicate()

        if output is not None:
            output = output.decode('utf-8')
        if error is not None:
            error = error.decode('utf-8')
            raise ValueError(error)

        logger.debug("Output: {}".format(output))
        logger.debug("Error: {}".format(error))

        lib_file = "{}/libdpumodel{}.so".format(self.work_dir, net_name)
        to_lib_file = "{}/libdpumodel{}.so".format(self.build_dir, net_name)
        shutil.move(lib_file, to_lib_file)

        # meta_file = "{}/meta.json".format(self.work_dir)
        self.meta["vitis_dpu_kernel"] = net_name
        to_meta_file = "{}/meta.json".format(self.build_dir)
        # shutil.move(meta_file, to_meta_file)

        with open(to_meta_file, 'w') as f:
            json.dump(self.meta, f)

        self.c_output.add(net_name, [to_lib_file], in_map, out_map)

        self.xgraph.set_compiler_output(self.c_output)

        return self.xgraph
Example #6
0
class ExternalQuantizer(XGraphBaseSubgraphQuantizer, ABC):

    xgraph_factory = XGraphFactory()
    xgraph_partitioner = XGraphPartitioner()

    def __init__(self,
                 xgraph,
                 inputs_func,
                 work_dir=os.path.join(os.getcwd(), 'work')):

        super(ExternalQuantizer, self).__init__(xgraph, inputs_func, work_dir)

        self.gen = TfGenerator()
        self.partition_graphs = {}
        self.res = {}
        self.q_output = QuantizerOutput(name=xgraph.get_name())

    def _propagate_quant_info(self, xgraph):
        # setup empty vqi and vqo for every layer w/o vai_quant
        for layer in xgraph.get_layers():
            if 'vai_quant' not in layer.attrs:
                layer.attrs['vai_quant'] = ['vai_quant_in', 'vai_quant_out']
                layer.attrs['vai_quant_in'] = ''
                layer.attrs['vai_quant_out'] = ''
        # for every layer
        for layer in xgraph.get_layers():
            # if the layer has non empty vqo, propagate it to the output layers
            if layer.attrs['vai_quant_out'] != '':
                l_vqo = layer.attrs['vai_quant_out']
                # for every output layer
                for t_idx, t_name in enumerate(layer.tops):
                    t_layer = xgraph.get(t_name)
                    # if the input quant is not specified in the output layer
                    if t_layer.attrs['vai_quant_in'] == '':
                        # get quant info from current layer, two by two
                        t_vqi = [l_vqo[2 * t_idx], l_vqo[2 * t_idx + 1]]
                        t_layer.attrs['vai_quant_in'] = t_vqi
            # if the layer has non empty vqi, propagate it to the input layers
            if layer.attrs['vai_quant_in'] != '':
                l_vqi = layer.attrs['vai_quant_in']
                # for every input layer
                for b_idx, b_name in enumerate(layer.bottoms):
                    b_layer = xgraph.get(b_name)
                    if b_layer.attrs['vai_quant_out'] == '':
                        b_vqo = [l_vqi[2 * b_idx], l_vqi[2 * b_idx + 1]]
                        b_layer.attrs['vai_quant_out'] = b_vqo

    def quantize(self):
        # NOTE For Conv2Dtranspose layers we need the specific batch size in tensorflow 1.13
        batch_size = list(self.inputs_func(0).values())[0].shape[0]
        fs = self.gen.generate(
            self.xgraph,
            'graph',
            subgraphs_only=True,
            layout='NHWC',
            batch_size=batch_size)
        assert len(fs) == 1, 'Too many partitions'
        partition_key = list(fs.keys())[0]
        pb_path = list(fs.values())[0]
        self.partition_graphs[partition_key] = pb_path

        q_xgraph = super(ExternalQuantizer, self).quantize()

        self.xgraph.meta_attrs["is_quantized"] = True
        for qkey in self.q_output.keys():
            if 'quant_keys' not in self.xgraph.meta_attrs:
                self.xgraph.meta_attrs['quant_keys'] = [qkey]
            else:
                self.xgraph.meta_attrs['quant_keys'].append(qkey)
            quant_file = self.q_output.get_q_file(qkey)
            quant_info_file = self.q_output.get_q_info(qkey)
            quant_orig_pb = self.q_output.get_orig_pb(qkey)
            self.xgraph.meta_attrs[qkey] = {
                'q_file': quant_file,
                'q_info': quant_info_file,
                'orig_pb': quant_orig_pb}
        return q_xgraph
class RuntimeDecentQSim(BaseRuntime):
    """Runtime for Decent quantizer simulation"""

    xgraph_partitioner = XGraphPartitioner()
    xgraph_factory = XGraphFactory()
    target_registry = TargetRegistry()

    def __init__(self,
                 name,
                 xgraph: XGraph,
                 device: str = 'cpu',
                 batch_size: int = -1,
                 placeholder: bool = False,
                 last_layers: List[str] = None,
                 **kwargs):
        super(RuntimeDecentQSim,
              self).__init__(name, xgraph, device, batch_size, placeholder,
                             last_layers)

        meta_attrs = self.xgraph.meta_attrs

        if 'quant_keys' not in meta_attrs:
            raise ValueError("Trying to simulate unquantized model. Make sure to first"\
                             " quantize the model.")

        qkey = meta_attrs['quant_keys'][0]
        self.q_eval = meta_attrs[qkey]['q_eval']
        self.gpu = 0

        Xps = RuntimeDecentQSim.xgraph_partitioner.get_subgraphs(xgraph)
        assert len(Xps) == 1, "Decent quantizer simulation only supports one partition"\
            " currently"
        self.Xp = Xps[0]
        target = self.Xp.attrs['target']
        opt_xgraph = RuntimeDecentQSim.target_registry.get_target_optimizer(
            target)(self.xgraph, target=target)
        self.rt_xgraph = RuntimeDecentQSim.target_registry.get_target_build_func(
            target
        )(
            copy.deepcopy(opt_xgraph),
            data_layout=
            'NHWC'  # NOTE XGraph's should be built in NHWC data layout, this is
            # important for DPUCADX8G where DPU execution happens in NCHW
            # but quantization simulation in NHWC
        )

    def _init_net(self, network: List[XLayer], params: Dict[str, np.ndarray]):
        # Do nothing
        pass

    def run_input(self, X: XLayer, inputs: Dict[str, Union[np.ndarray, List[np.ndarray]]])\
            -> Dict[str, Union[np.ndarray, List[np.ndarray]]]:
        return None

    def run_transpose(self, X: XLayer, inputs: Dict[str, Union[np.ndarray, List[np.ndarray]]])\
            -> Dict[str, Union[np.ndarray, List[np.ndarray]]]:
        assert len(X.bottoms) == 1
        return np.transpose(inputs[X.bottoms[0]],
                            axes=tuple(X.attrs['axes'][:]))

    def run_dpu(self, X: XLayer, inputs: Dict[str, Union[np.ndarray, List[np.ndarray]]])\
            -> Dict[str, Union[np.ndarray, List[np.ndarray]]]:
        import tensorflow as tf
        tf.compat.v1.reset_default_graph()
        os.environ["CUDA_VISIBLE_DEVICES"] = str(self.gpu)
        input_graph_def = tf.Graph().as_graph_def()
        input_graph_def.ParseFromString(
            tf.io.gfile.GFile(self.q_eval, "rb").read())
        tf.import_graph_def(input_graph_def, name='')

        input_names = X.attrs["input_names"]
        input_map = {
            X.attrs["__bottom_tensors"][in_name][0]: in_name
            for in_name in input_names
        }
        in_tensors = {
            k:
            tf.compat.v1.get_default_graph().get_tensor_by_name(input_map[k] +
                                                                ":0")
            for k in X.bottoms
        }

        feed_dict = {in_tensors[k]: inputs[k] for k in X.bottoms}

        out_names = X.attrs["output_names"]
        out_tensor_names = [X.attrs["output_layers"][o][-1] for o in out_names]

        out_tensors = [
            tf.compat.v1.get_default_graph().get_tensor_by_name(o + "/aquant" +
                                                                ":0")
            for o in out_names
        ]

        with tf.compat.v1.Session() as sess:
            out = sess.run(out_tensors, feed_dict=feed_dict)
            return out if isinstance(out, list) else [out]

    def run_tuple_get_item(self, X: XLayer, inputs: Dict[str, Union[np.ndarray, List[np.ndarray]]])\
            -> Dict[str, Union[np.ndarray, List[np.ndarray]]]:
        assert len(X.bottoms) == 1
        index = X.attrs['index']
        data = inputs[X.bottoms[0]][index]
        if 'transpose' in X.attrs and X.attrs['transpose'] is True:
            return np.transpose(data, axes=tuple(X.attrs['axes'][:]))
        return data

    def run_tuple(self, X: XLayer, inputs: Dict[str, Union[np.ndarray, List[np.ndarray]]])\
            -> Dict[str, Union[np.ndarray, List[np.ndarray]]]:
        return [inputs[b] for b in X.bottoms]

    def run(self,
            inputs: Dict[str, np.ndarray],
            outputs: List[str] = [],
            stop: str = None,
            force_stepwise: bool = False,
            debug: bool = False) -> List[np.ndarray]:
        """Override run method"""
        for X in self.rt_xgraph.get_layers():
            if 'Input' in X.type:
                outs = self.run_input(X, inputs)
            elif 'Transpose' in X.type:
                outs = self.run_transpose(X, inputs)
            elif 'DPU' in X.type:
                outs = self.run_dpu(X, inputs)
            elif 'TupleGetItem' in X.type:
                outs = self.run_tuple_get_item(X, inputs)
            elif 'Tuple' in X.type:
                outs = self.run_tuple(X, inputs)
            else:
                raise NotImplementedError(
                    "Unsupported operation in decentq simulation: {}".format(
                        X.type[0]))
            if outs is not None:
                inputs[X.name] = outs
        return [inputs[o] for o in outputs]
Example #8
0
class TestDPUContrib(unittest.TestCase):

    xgraph_partitioner = XGraphPartitioner()
    xgraph_factory = XGraphFactory()
    target_registry = TargetRegistry()

    @classmethod
    def setUpClass(cls):

        # Import DPU module
        from pyxir.contrib.dpuv2 import dpuv2

    @classmethod
    def tearDownClass(cls):
        # Unregister dpu for other tests
        TestDPUContrib.target_registry.unregister_target('dpuv2-zcu104')
        TestDPUContrib.target_registry.unregister_target('dpuv2-zcu102')
        TestDPUContrib.target_registry.unregister_target('DPUCZDX8G-zcu102')
        TestDPUContrib.target_registry.unregister_target('DPUCZDX8G-zcu104')
        TestDPUContrib.target_registry.unregister_target('dpuv2-ultra96')
        TestDPUContrib.target_registry.unregister_target('DPUCZDX8G-ultra96')

    def test_supported_ops(self):
        ultra96_ops = TestDPUContrib.target_registry\
            .get_supported_op_check_names('dpuv2-ultra96')

        assert 'BatchNorm' in ultra96_ops
        assert 'BiasAdd' in ultra96_ops
        assert 'Concat' in ultra96_ops
        assert 'Convolution' in ultra96_ops
        assert 'Conv2DTranspose' in ultra96_ops
        assert 'DPU' in ultra96_ops
        assert 'Eltwise' in ultra96_ops
        assert 'Pad' in ultra96_ops
        assert 'Pooling' in ultra96_ops
        assert 'Mean' in ultra96_ops
        assert 'pReLU' in ultra96_ops
        assert 'ReLU' in ultra96_ops
        assert 'ReLU6' in ultra96_ops
        assert 'Scale' in ultra96_ops

        zcu102_ops = TestDPUContrib.target_registry\
            .get_supported_op_check_names('dpuv2-zcu102')

        assert 'BatchNorm' in zcu102_ops
        assert 'BiasAdd' in zcu102_ops
        assert 'Concat' in zcu102_ops
        assert 'Convolution' in zcu102_ops
        assert 'Conv2DTranspose' in zcu102_ops
        assert 'DPU' in zcu102_ops
        assert 'Eltwise' in zcu102_ops
        assert 'Pad' in zcu102_ops
        assert 'Pooling' in zcu102_ops
        assert 'Mean' in zcu102_ops
        assert 'pReLU' in zcu102_ops
        assert 'ReLU' in zcu102_ops
        assert 'ReLU6' in zcu102_ops
        assert 'Scale' in zcu102_ops

        zcu104_ops = TestDPUContrib.target_registry\
            .get_supported_op_check_names('dpuv2-zcu104')

        assert 'BatchNorm' in zcu104_ops
        assert 'BiasAdd' in zcu104_ops
        assert 'Concat' in zcu104_ops
        assert 'Convolution' in zcu104_ops
        assert 'Conv2DTranspose' in zcu104_ops
        assert 'DPU' in zcu104_ops
        assert 'Eltwise' in zcu104_ops
        assert 'Pad' in zcu104_ops
        assert 'Pooling' in zcu104_ops
        assert 'Mean' in zcu104_ops
        assert 'pReLU' in zcu104_ops
        assert 'ReLU' in zcu104_ops
        assert 'ReLU6' in zcu104_ops
        assert 'Scale' in zcu104_ops

    def test_small(self):
        net = [
            XLayer(name='in1',
                   type=['Input'],
                   shapes=[1, 1, 4, 4],
                   sizes=[16],
                   bottoms=[],
                   tops=['conv1'],
                   layer=['in1'],
                   targets=[]),
            XLayer(name='in2',
                   type=['Input'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=[],
                   tops=['dense1'],
                   layer=['in2'],
                   targets=[]),
            XLayer(name='conv1',
                   type=['Convolution'],
                   shapes=[1, 2, 3, 3],
                   sizes=[18],
                   bottoms=['in1'],
                   tops=['pool1'],
                   layer=['conv1'],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   attrs={
                       'data_layout': 'NCHW',
                       'padding': [[0, 0], [0, 0], [1, 1], [1, 1]],
                       'kernel_size': [3, 3],
                       'strides': [1, 1],
                       'dilation': [1, 1],
                       'groups': 1,
                       'channels': [2, 2]
                   },
                   targets=[]),
            XLayer(name='pool1',
                   type=['Pooling'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=['conv1'],
                   tops=['dense1'],
                   layer=['pool1'],
                   attrs={
                       'data_layout': 'NCHW',
                       'padding': [[0, 0], [0, 0], [1, 1], [1, 1]],
                       'kernel_size': [3, 3],
                       'strides': [1, 1],
                   },
                   targets=[]),
            XLayer(name='dense1',
                   type=['Dense'],
                   shapes=[1, 20],
                   sizes=[20],
                   bottoms=['pool1', 'in2'],
                   tops=[],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   layer=['dense1'],
                   targets=[])
        ]
        xgraph = TestDPUContrib.xgraph_factory.build_from_xlayer(net)
        p_xgraph = partition(xgraph, ['dpuv2-zcu104'])
        dpu_xgraph = TestDPUContrib.target_registry\
            .get_target_build_func('dpuv2-zcu104')(p_xgraph)

        assert len(dpu_xgraph) == 6
        layers = dpu_xgraph.get_layers()

        assert layers[0].type[0] == 'Input'

        assert layers[1].type[0] == 'Transpose'
        assert layers[1].bottoms == ['in1']
        assert layers[1].tops == ['xp0']

        assert layers[2].type[0] == 'DPU'
        assert layers[2].bottoms == ['conv1_bottom_NCHW>NHWC']
        assert layers[2].tops == ['pool1']
        assert layers[2].shapes == [[1, 2, 2, 2]]
        assert layers[2].attrs['target'] == 'dpuv2-zcu104'
        assert layers[2].attrs['input_names'] == ['xinput0']
        assert layers[2].attrs['output_names'] == ['pool1']
        assert layers[2].attrs['input_layers']['xinput0'] == ['conv1']
        assert layers[2].attrs['output_layers']['pool1'] == ['pool1']
        assert layers[2].attrs['__top_tensors'] ==\
            {'pool1': ['pool1_top_NHWC>NCHW']}
        assert layers[2].attrs['orig_top_tensors'] ==\
            {'pool1': ['dense1']}
        assert layers[2].attrs['__bottom_tensors'] ==\
            {'xinput0': ['conv1_bottom_NCHW>NHWC']}
        assert layers[2].attrs['orig_bottom_tensors'] ==\
            {'xinput0': ['in1']}

        # Merged TupleGetItem and Transpose layer
        assert layers[3].type[0] == 'TupleGetItem'
        assert layers[3].name == 'pool1'
        assert layers[3].shapes == [1, 2, 2, 2]
        assert layers[3].bottoms == ['xp0']
        assert layers[3].tops == ['dense1']
        assert layers[3].attrs['transpose'] is True

        assert layers[4].type[0] == 'Input'
        assert layers[4].name == 'in2'
        assert layers[4].tops == ['dense1']

        assert layers[5].type[0] == 'Dense'
        assert layers[5].name == 'dense1'
        assert layers[5].shapes == [1, 20]
        assert layers[5].bottoms == ['pool1', 'in2']
        assert layers[5].tops == []
Example #9
0
from pyxir.shared.xbuffer import XBuffer
from pyxir.graph.xgraph import XGraph
from pyxir.graph.io.xgraph_io import XGraphIO
from pyxir.io.api import visualize, save, load, get_xgraph_str
from pyxir.runtime import runtime_factory
from pyxir.runtime.base_runtime import BaseRuntime
from pyxir.graph.partitioning.xgraph_partitioner import XGraphPartitioner
from pyxir.graph.optimization.optimizers.basic_optimizer \
    import XGraphBasicOptimizer
from pyxir.graph.transformers.layout_transformation_pass \
    import XGraphLayoutTransformationPass

logger = logging.getLogger("pyxir")
fancy_logger = fancy_logging.getLogger("pyxir")

xgraph_partitioner = XGraphPartitioner()
target_registry = TargetRegistry()


def stringify(name: str):
    # Tensorflow raises invalid scope name errors if name is invalid
    name = re.sub('[^A-Za-z0-9_.\\-/]', '-', name)
    try:
        # Some modules in Tensorflow subgraph contrib have issues with names
        #   that look like ints
        int(name)
        return str(name) + "_"
    except ValueError:
        return str(name)

Example #10
0
class TestQuantSimPass(unittest.TestCase):

    xgraph_partitioner = XGraphPartitioner()
    xgraph_factory = XGraphFactory()
    xf_exec_graph_factory = RuntimeFactory()

    @classmethod
    def setUpClass(cls):
        def xgraph_build_func(xgraph):
            raise NotImplementedError("")

        def xgraph_optimizer(xgraph):
            raise NotImplementedError("")

        def xgraph_quantizer(xgraph):
            raise NotImplementedError("")

        def xgraph_compiler(xgraph):
            raise NotImplementedError("")

        target_registry = TargetRegistry()
        target_registry.register_target('npu_test', xgraph_optimizer,
                                        xgraph_quantizer, xgraph_compiler,
                                        xgraph_build_func)

        @register_op_support_check('npu_test', 'Convolution')
        def conv_op_support(X, bXs, tXs):
            return True

    @classmethod
    def tearDownClass(cls):

        target_registry = TargetRegistry()
        target_registry.unregister_target('npu_test')

    def test_conv(self):
        W = np.reshape(
            np.array([[[1, 2], [3, 0]], [[1, 1], [0, 1]]], dtype=np.float32),
            (2, 1, 2, 2))
        B = np.array([0., 0.], dtype=np.float32)

        net = [
            XLayer(name='in',
                   type=['Input'],
                   shapes=[1, 1, 4, 4],
                   sizes=[16],
                   bottoms=[],
                   tops=['conv2d0'],
                   layer=['in'],
                   targets=[]),
            XLayer(name='conv2d0',
                   type=['Convolution'],
                   shapes=[1, 2, 3, 3],
                   sizes=[18],
                   bottoms=['in'],
                   tops=[],
                   layer=['conv2d0'],
                   data=ConvData(W, B),
                   attrs={
                       'data_layout': 'NCHW',
                       'kernel_layout': 'OIHW',
                       'shape': [1, 2, 3, 3],
                       'padding': [[0, 0], [0, 0], [0, 0], [0, 0]],
                       'strides': [1, 1],
                       'dilation': [1, 1],
                       'groups': 1
                   },
                   targets=[])
        ]
        xgraph = TestQuantSimPass.xgraph_factory.build_from_xlayer(
            net, name='test1')

        quant_sim_pass = XGraphQuantSimPass(fdir=FILE_PATH,
                                            name=xgraph.get_name() + '_qsim')
        qsim_xgraph = quant_sim_pass.execute(xgraph=xgraph,
                                             subgraphs_only=False)

        exec_graph = TestQuantSimPass.xf_exec_graph_factory.build_runtime(
            qsim_xgraph)

        inpts = {
            'in':
            np.reshape(
                np.array([[10, 10, 0, 40], [50, 10, 0, 80], [30, 50, 10, 0],
                          [10, 90, 30, 40]],
                         dtype=np.float32), (1, 1, 4, 4))
        }
        res = exec_graph.run(inpts)
        outpt = res[0]
        # for idx, layer, inpts, outpt, _ in exec_graph.run_stepwise(inpts):
        #     print(layer.name, outpt)

        expected_outpt = np.array([[[[182.28346, 36.45669, 80.20472],
                                     [160.40944, 160.40944, 189.5748],
                                     [160.40944, 342.6929, 102.078735]],
                                    [[29.165354, 7.2913384, 123.95275],
                                     [109.37008, 21.874016, 80.20472],
                                     [167.70079, 87.49606, 51.039368]]]],
                                  dtype=np.float32)

        np.testing.assert_array_almost_equal(outpt, expected_outpt, decimal=4)

    def test_conv_maxpool_subgraph(self):
        W = np.reshape(
            np.array([[[1, 2], [3, 0]], [[1, 1], [0, 1]]], dtype=np.float32),
            (2, 1, 2, 2))
        B = np.array([0., 0.], dtype=np.float32)

        net = [
            XLayer(name='in',
                   type=['Input'],
                   shapes=[1, 1, 4, 4],
                   sizes=[16],
                   bottoms=[],
                   tops=['conv2d0'],
                   layer=['in'],
                   targets=[]),
            XLayer(name='conv2d0',
                   type=['Convolution'],
                   shapes=[1, 2, 3, 3],
                   sizes=[18],
                   bottoms=['in'],
                   tops=[],
                   layer=['conv2d0'],
                   data=ConvData(W, B),
                   attrs={
                       'data_layout': 'NCHW',
                       'kernel_layout': 'OIHW',
                       'shape': [1, 2, 3, 3],
                       'padding': [[0, 0], [0, 0], [0, 0], [0, 0]],
                       'strides': [1, 1],
                       'dilation': [1, 1],
                       'groups': 1
                   },
                   targets=[]),
            XLayer(name='max_pool2d0',
                   type=['Pooling'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=['conv2d0'],
                   tops=[],
                   layer=['max_pool2d0'],
                   attrs={
                       'kernel_size': [2, 2],
                       'insize': [3, 3],
                       'outsize': [2, 2],
                       'data_layout': 'NCHW',
                       'padding': [[0, 0], [0, 0], [0, 0], [0, 0]],
                       'strides': [1, 1],
                       'pool_type': 'Max'
                   },
                   targets=[])
        ]
        xgraph = TestQuantSimPass.xgraph_factory.build_from_xlayer(
            net, name='testtest')

        p_xgraph = TestQuantSimPass.xgraph_partitioner.partition(
            xgraph, ['npu_test'])

        assert (p_xgraph.get_layers()[0].target == 'cpu')
        assert (p_xgraph.get_layers()[1].target == 'npu_test')
        assert (p_xgraph.get_layers()[2].target == 'cpu')

        assert (p_xgraph.get_layers()[0].subgraph is None)
        assert (p_xgraph.get_layers()[1].subgraph == 'xp0')
        assert (p_xgraph.get_layers()[2].subgraph is None)

        quant_sim_pass = XGraphQuantSimPass(fdir=FILE_PATH,
                                            name=xgraph.get_name() + '_qsim')
        qsim_xgraph = quant_sim_pass.execute(xgraph=p_xgraph,
                                             subgraphs_only=True)

        exec_graph = TestQuantSimPass.xf_exec_graph_factory.build_runtime(
            qsim_xgraph)

        inpts = {
            'in':
            np.reshape(
                np.array([[10, 10, 0, 40], [50, 10, 0, 80], [30, 50, 10, 0],
                          [10, 90, 30, 40]],
                         dtype=np.float32), (1, 1, 4, 4))
        }
        res = exec_graph.run(inpts)
        outpt = res[0]

        expected_outpt = np.array(
            [[[[182.28346, 189.5748], [342.6929, 342.6929]],
              [[109.37008, 123.95275], [167.70079, 87.49606]]]],
            dtype=np.float32)

        np.testing.assert_array_almost_equal(outpt, expected_outpt, decimal=4)
Example #11
0
class TestXGraphPartitioner(unittest.TestCase):

    xgraph_partitioner = XGraphPartitioner()
    xgraph_factory = XGraphFactory()

    @classmethod
    def setUpClass(cls):
        def xgraph_build_func(xgraph):
            raise NotImplementedError("")

        def xgraph_optimizer(xgraph):
            raise NotImplementedError("")

        def xgraph_quantizer(xgraph):
            raise NotImplementedError("")

        def xgraph_compiler(xgraph):
            raise NotImplementedError("")

        target_registry = TargetRegistry()
        target_registry.register_target('test', xgraph_optimizer,
                                        xgraph_quantizer, xgraph_compiler,
                                        xgraph_build_func)

        @register_op_support_check('test', 'Convolution')
        def conv_op_support(X, bXs, tXs):
            return True

        @register_op_support_check('test', 'Pooling')
        def pooling_op_support(X, bXs, tXs):
            return True

        @register_op_support_check('test', 'Concat')
        def concat_op_support(X, bXs, tXs):
            return True

    @classmethod
    def tearDownClass(cls):

        target_registry = TargetRegistry()
        target_registry.unregister_target('test')

    def test_basic(self):
        net = [
            XLayer(name='in1',
                   type=['Input'],
                   shapes=[1, 1, 4, 4],
                   sizes=[16],
                   bottoms=[],
                   tops=['conv1'],
                   layer=['in1'],
                   targets=[]),
            XLayer(name='in2',
                   type=['Input'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=[],
                   tops=['add1'],
                   layer=['in2'],
                   targets=[]),
            XLayer(name='conv1',
                   type=['Convolution'],
                   shapes=[1, 2, 3, 3],
                   sizes=[18],
                   bottoms=['in1'],
                   tops=['pool1'],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   layer=['conv1'],
                   targets=[]),
            XLayer(name='pool1',
                   type=['Pooling'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=['conv1'],
                   tops=['add1'],
                   layer=['pool1'],
                   targets=[]),
            XLayer(name='add1',
                   type=['Eltwise'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=['pool1', 'in2'],
                   tops=[],
                   layer=['add1'],
                   targets=[])
        ]
        xgraph = TestXGraphPartitioner.xgraph_factory.build_from_xlayer(net)
        TargetRegistry().annotate_ops(xgraph)
        p_xgraph = TestXGraphPartitioner.xgraph_partitioner.partition(
            xgraph, ['test'])

        assert len(p_xgraph.get_layer_names()) == 5
        assert p_xgraph.get_subgraph_names() == ['xp0']

        p_xlayers = p_xgraph.get_layers()
        assert p_xlayers[0].type[0] in ['Input']
        assert p_xlayers[1].type[0] in ['Convolution']
        assert p_xlayers[2].type[0] in ['Pooling']
        assert p_xlayers[3].type[0] in ['Input']
        assert p_xlayers[4].type[0] in ['Eltwise']

        assert p_xlayers[0].target == 'cpu'
        assert p_xlayers[1].target == 'test'
        assert p_xlayers[2].target == 'test'
        assert p_xlayers[3].target == 'cpu'
        assert p_xlayers[4].target == 'cpu'

        assert p_xlayers[0].subgraph is None
        assert p_xlayers[1].subgraph == 'xp0'
        assert p_xlayers[2].subgraph == 'xp0'
        assert p_xlayers[3].subgraph is None
        assert p_xlayers[4].subgraph is None

        subgraphs = TestXGraphPartitioner.xgraph_partitioner.get_subgraphs(
            p_xgraph)

        assert len(subgraphs) == 1
        xp0 = subgraphs[0]
        assert xp0.name == 'xp0'
        xp0_xgraph = TestXGraphPartitioner.xgraph_factory\
            .build_from_xlayer(xp0.subgraph_data)

        assert xp0.bottoms == ['in1']
        assert xp0.tops == ['add1']
        assert xp0.shapes == [[1, 2, 2, 2]]
        assert xp0.sizes == [8]

        assert len(xp0_xgraph) == 3
        xp0_layers = xp0_xgraph.get_layers()

        assert xp0_layers[0].type[0] == 'Input'
        assert xp0_layers[0].layer[0] == 'conv1'
        assert xp0_layers[1].type[0] == 'Convolution'
        assert xp0_layers[2].type[0] == 'Pooling'

        assert xp0_layers[0].bottoms == []
        assert xp0_layers[0].tops == ['conv1']
        assert xp0_layers[1].bottoms == ['xinput0']
        assert xp0_layers[1].tops == ['pool1']
        assert xp0_layers[2].bottoms == ['conv1']
        assert xp0_layers[2].tops == []

    def test_complete_partition(self):
        net = [
            XLayer(name='in1',
                   type=['Input'],
                   shapes=[1, 1, 4, 4],
                   sizes=[16],
                   bottoms=[],
                   tops=['conv1'],
                   layer=['in1'],
                   targets=[]),
            XLayer(name='conv1',
                   type=['Convolution'],
                   shapes=[1, 2, 3, 3],
                   sizes=[18],
                   bottoms=['in1'],
                   tops=['pool1'],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   layer=['conv1'],
                   targets=[]),
            XLayer(name='pool1',
                   type=['Pooling'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=['conv1'],
                   tops=[],
                   layer=['pool1'],
                   targets=[])
        ]
        xgraph = TestXGraphPartitioner.xgraph_factory.build_from_xlayer(net)
        TargetRegistry().annotate_ops(xgraph)
        p_xgraph = TestXGraphPartitioner.xgraph_partitioner.partition(
            xgraph, ['test'])

        assert len(p_xgraph.get_layer_names()) == 3
        assert p_xgraph.get_subgraph_names() == ['xp0']

        p_xlayers = p_xgraph.get_layers()
        assert p_xlayers[0].type[0] in ['Input']
        assert p_xlayers[1].type[0] in ['Convolution']
        assert p_xlayers[2].type[0] in ['Pooling']

        assert p_xlayers[0].target == 'cpu'
        assert p_xlayers[1].target == 'test'
        assert p_xlayers[2].target == 'test'

        assert p_xlayers[0].subgraph is None
        assert p_xlayers[1].subgraph == 'xp0'
        assert p_xlayers[2].subgraph == 'xp0'

        subgraphs = TestXGraphPartitioner.xgraph_partitioner.get_subgraphs(
            p_xgraph)

        assert len(subgraphs) == 1
        xp0 = subgraphs[0]
        assert xp0.name == 'xp0'
        xp0_xgraph = TestXGraphPartitioner.xgraph_factory\
            .build_from_xlayer(xp0.subgraph_data)

        assert xp0.bottoms == ['in1']
        assert xp0.tops == []
        assert xp0.shapes == [[1, 2, 2, 2]]
        assert xp0.sizes == [8]
        assert xp0.attrs['target'] == 'test'
        assert xp0.attrs['__bottom_tensors'] == {'xinput0': ['in1']}
        assert xp0.attrs['orig_bottom_tensors'] == {'xinput0': ['in1']}
        assert xp0.attrs['__top_tensors'] == {'pool1': []}
        assert xp0.attrs['orig_top_tensors'] == {'pool1': []}

        assert len(xp0_xgraph) == 3
        xp0_layers = xp0_xgraph.get_layers()

        assert xp0_layers[0].type[0] == 'Input'
        assert xp0_layers[0].layer[0] == 'conv1'
        assert xp0_layers[1].type[0] == 'Convolution'
        assert xp0_layers[2].type[0] == 'Pooling'

        assert xp0_layers[0].bottoms == []
        assert xp0_layers[0].tops == ['conv1']
        assert xp0_layers[1].bottoms == ['xinput0']
        assert xp0_layers[1].tops == ['pool1']
        assert xp0_layers[2].bottoms == ['conv1']
        assert xp0_layers[2].tops == []

    def test_two_partitions_through_interruption(self):
        # A layer inside a residual type branch os not supported
        # Here: BatchNorm
        net = [
            XLayer(name='in1',
                   type=['Input'],
                   shapes=[1, 1, 4, 4],
                   sizes=[16],
                   bottoms=[],
                   tops=['conv1'],
                   layer=['in1'],
                   targets=[]),
            XLayer(name='conv1',
                   type=['Convolution'],
                   shapes=[1, 2, 3, 3],
                   sizes=[18],
                   bottoms=['in1'],
                   tops=['pool1', 'bn1'],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   layer=['conv1'],
                   targets=[]),
            XLayer(name='pool1',
                   type=['Pooling'],
                   shapes=[1, 4, 3, 3],
                   sizes=[36],
                   bottoms=['conv1'],
                   tops=['concat1'],
                   layer=['pool1'],
                   targets=[]),
            XLayer(name='bn1',
                   type=['BatchNorm'],
                   shapes=[1, 2, 3, 3],
                   sizes=[18],
                   bottoms=['conv1'],
                   tops=['concat1'],
                   data=BatchData(np.array([1, 1]), np.array([0, 0]),
                                  np.array([1, 1]), np.array([0, 0])),
                   layer=['bn1'],
                   targets=[]),
            XLayer(name='concat1',
                   type=['Concat'],
                   shapes=[1, 6, 3, 3],
                   sizes=[54],
                   bottoms=['pool1', 'bn1'],
                   tops=['conv2'],
                   layer=['concat1'],
                   targets=[]),
            XLayer(name='conv2',
                   type=['Convolution'],
                   shapes=[1, 10, 2, 2],
                   sizes=[40],
                   bottoms=['concat1'],
                   tops=[],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   layer=['conv2'],
                   targets=[])
        ]
        xgraph = TestXGraphPartitioner.xgraph_factory.build_from_xlayer(net)
        TargetRegistry().annotate_ops(xgraph)
        p_xgraph = TestXGraphPartitioner.xgraph_partitioner.partition(
            xgraph, ['test'])

        assert len(p_xgraph.get_layer_names()) == 6
        assert p_xgraph.get_subgraph_names() == ['xp0']

        p_xlayers = p_xgraph.get_layers()
        assert p_xlayers[0].type[0] in ['Input']
        assert p_xlayers[1].type[0] in ['Convolution']
        assert p_xlayers[2].type[0] in ['Pooling']
        assert p_xlayers[3].type[0] in ['BatchNorm']
        assert p_xlayers[4].type[0] in ['Concat']
        assert p_xlayers[5].type[0] in ['Convolution']

        assert p_xlayers[0].target == 'cpu'
        assert p_xlayers[1].target == 'test'
        assert p_xlayers[2].target == 'test'
        assert p_xlayers[3].target == 'cpu'
        assert p_xlayers[4].target == 'cpu'
        assert p_xlayers[5].target == 'cpu'

        assert p_xlayers[0].subgraph is None
        assert p_xlayers[1].subgraph == 'xp0'
        assert p_xlayers[2].subgraph == 'xp0'
        assert p_xlayers[3].subgraph is None
        assert p_xlayers[4].subgraph is None
        assert p_xlayers[5].subgraph is None

        assert p_xlayers[3].name == 'bn1'
        assert p_xlayers[3].bottoms == ['conv1']
        assert p_xlayers[3].tops == ['concat1']

        assert p_xlayers[4].name == 'concat1'
        assert p_xlayers[4].bottoms == ['pool1', 'bn1']
        assert p_xlayers[4].tops == ['conv2']

        subgraphs = TestXGraphPartitioner.xgraph_partitioner.get_subgraphs(
            p_xgraph)

        assert len(subgraphs) == 1
        xp0 = subgraphs[0]
        assert xp0.name == 'xp0'
        xp0_xgraph = TestXGraphPartitioner.xgraph_factory\
            .build_from_xlayer(xp0.subgraph_data)

        assert xp0.bottoms == ['in1']
        assert xp0.tops == ['bn1', 'concat1']
        assert xp0.shapes == [[1, 2, 3, 3], [1, 4, 3, 3]]
        assert xp0.sizes == [18, 36]
        assert xp0.attrs['target'] == 'test'
        assert xp0.attrs['__bottom_tensors'] == {'xinput0': ['in1']}
        assert xp0.attrs['orig_bottom_tensors'] == {'xinput0': ['in1']}
        assert xp0.attrs['__top_tensors'] == \
            {'conv1': ['bn1'], 'pool1': ['concat1']}
        assert xp0.attrs['orig_top_tensors'] == \
            {'conv1': ['bn1'], 'pool1': ['concat1']}

        assert (len(xp0_xgraph) == 3)
        xp0_layers = xp0_xgraph.get_layers()

        assert [X.name for X in xp0_xgraph.get_input_layers()] == ['xinput0']
        # TODO: XGraph only recognizes output layers when they have no top
        #   layers
        assert [X.name for X in xp0_xgraph.get_output_layers()] ==\
            ['pool1']

        assert xp0_layers[0].type[0] == 'Input'
        assert xp0_layers[0].layer[0] == 'conv1'
        assert xp0_layers[1].type[0] == 'Convolution'
        assert xp0_layers[2].type[0] == 'Pooling'

        assert xp0_layers[0].bottoms == []
        assert xp0_layers[0].tops == ['conv1']
        assert xp0_layers[1].bottoms == ['xinput0']
        assert xp0_layers[1].tops == ['pool1']
        assert xp0_layers[2].bottoms == ['conv1']
        assert xp0_layers[2].tops == []

    def test_multiple_partitions(self):
        net = [
            XLayer(name='in1',
                   type=['Input'],
                   shapes=[1, 1, 4, 4],
                   sizes=[16],
                   bottoms=[],
                   tops=['conv1'],
                   layer=['in1'],
                   targets=[]),
            XLayer(name='in2',
                   type=['Input'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=[],
                   tops=['add1'],
                   layer=['in2'],
                   targets=[]),
            XLayer(name='conv1',
                   type=['Convolution'],
                   shapes=[1, 2, 3, 3],
                   sizes=[18],
                   bottoms=['in1'],
                   tops=['pool1'],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   layer=['conv1'],
                   targets=[]),
            XLayer(name='pool1',
                   type=['Pooling'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=['conv1'],
                   tops=['add1'],
                   layer=['pool1'],
                   targets=[]),
            XLayer(name='add1',
                   type=['Eltwise'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=['pool1', 'in2'],
                   tops=[],
                   layer=['add1'],
                   targets=[]),
            XLayer(name='bn1',
                   type=['BatchNorm'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=['add1'],
                   tops=['pool2'],
                   data=BatchData(np.array([1, 1]), np.array([0, 0]),
                                  np.array([1, 1]), np.array([0, 0])),
                   layer=['bn1'],
                   targets=[]),
            XLayer(name='pool2',
                   type=['Pooling'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=['bn1'],
                   tops=[],
                   layer=['pool2'],
                   targets=[])
        ]
        xgraph = TestXGraphPartitioner.xgraph_factory.build_from_xlayer(net)
        TargetRegistry().annotate_ops(xgraph)
        p_xgraph = TestXGraphPartitioner.xgraph_partitioner.partition(
            xgraph, ['test'])

        assert (len(p_xgraph.get_layer_names()) == 7)
        # ! Only xp0 because only one subgraph can exist for now (largest)
        assert (set(p_xgraph.get_subgraph_names()) == set(['xp0']))

        p_xlayers = p_xgraph.get_layers()
        assert (p_xlayers[0].type[0] in ['Input'])
        assert (p_xlayers[1].type[0] in ['Convolution'])
        assert (p_xlayers[2].type[0] in ['Pooling'])
        assert (p_xlayers[3].type[0] in ['Input'])
        assert (p_xlayers[4].type[0] in ['Eltwise'])
        assert (p_xlayers[5].type[0] in ['BatchNorm'])
        assert (p_xlayers[6].type[0] in ['Pooling'])

        assert (p_xlayers[0].target == 'cpu')
        assert (p_xlayers[1].target == 'test')
        assert (p_xlayers[2].target == 'test')
        assert (p_xlayers[3].target == 'cpu')
        assert (p_xlayers[4].target == 'cpu')
        assert (p_xlayers[5].target == 'cpu')
        # ! CPU because only one subgraph can exist for now (largest)
        assert (p_xlayers[6].target == 'cpu')

        assert (p_xlayers[0].subgraph is None)
        assert (p_xlayers[1].subgraph == 'xp0')
        assert (p_xlayers[2].subgraph == 'xp0')
        assert (p_xlayers[3].subgraph is None)
        assert (p_xlayers[4].subgraph is None)
        assert (p_xlayers[5].subgraph is None)
        assert (p_xlayers[6].subgraph is None)

        subgraphs = TestXGraphPartitioner.xgraph_partitioner.get_subgraphs(
            p_xgraph)

        assert (len(subgraphs) == 1)
        xp0 = subgraphs[0]
        assert (xp0.name == 'xp0')
        xp0_xgraph = TestXGraphPartitioner.xgraph_factory\
            .build_from_xlayer(xp0.subgraph_data)

        assert (xp0.bottoms == ['in1'])
        assert (xp0.tops == ['add1'])
        assert (xp0.shapes == [[1, 2, 2, 2]])
        assert (xp0.sizes == [8])

        assert (len(xp0_xgraph) == 3)
        xp0_layers = xp0_xgraph.get_layers()

        assert (xp0_layers[0].type[0] == 'Input')
        assert (xp0_layers[0].layer[0] == 'conv1')
        assert (xp0_layers[1].type[0] == 'Convolution')
        assert (xp0_layers[2].type[0] == 'Pooling')

        assert (xp0_layers[0].bottoms == [])
        assert (xp0_layers[0].tops == ['conv1'])
        assert (xp0_layers[1].bottoms == ['xinput0'])
        assert (xp0_layers[1].tops == ['pool1'])
        assert (xp0_layers[2].bottoms == ['conv1'])
        assert (xp0_layers[2].tops == [])

    def test_multiple_partitions_largest_last(self):
        net = [
            XLayer(name='in1',
                   type=['Input'],
                   shapes=[1, 1, 4, 4],
                   sizes=[16],
                   bottoms=[],
                   tops=['conv1'],
                   layer=['in1'],
                   targets=[]),
            XLayer(name='conv1',
                   type=['Convolution'],
                   shapes=[1, 2, 3, 3],
                   sizes=[18],
                   bottoms=['in1'],
                   tops=['t1'],
                   layer=['conv1'],
                   data=ConvData(weights=np.array([1, 1], dtype=np.float32),
                                 biases=np.array([0, 0], dtype=np.float32)),
                   targets=[]),
            XLayer(name='t1',
                   type=['Transpose'],
                   shapes=[1, 3, 3, 2],
                   sizes=[18],
                   bottoms=['conv1'],
                   tops=['conv2'],
                   layer=['t1'],
                   targets=[],
                   attrs={'axes': [0, 2, 3, 1]}),
            XLayer(name='conv2',
                   type=['Convolution'],
                   shapes=[1, 3, 3, 2],
                   sizes=[18],
                   bottoms=['t1'],
                   tops=['pool1'],
                   layer=['conv2'],
                   data=ConvData(weights=np.array([1, 1], dtype=np.float32),
                                 biases=np.array([0, 0], dtype=np.float32)),
                   targets=[]),
            XLayer(name='pool1',
                   type=['Pooling'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=['conv2'],
                   tops=[],
                   layer=['pool1'],
                   targets=[])
        ]
        xgraph = TestXGraphPartitioner.xgraph_factory.build_from_xlayer(net)
        TargetRegistry().annotate_ops(xgraph)
        p_xgraph = TestXGraphPartitioner.xgraph_partitioner.partition(
            xgraph, ['test'])

        assert len(p_xgraph.get_layer_names()) == 5
        # ! Only xp1 because only one subgraph can exist for now (largest)
        assert set(p_xgraph.get_subgraph_names()) == set(['xp1'])

        p_xlayers = p_xgraph.get_layers()
        assert (p_xlayers[0].type[0] in ['Input'])
        assert (p_xlayers[1].type[0] in ['Convolution'])
        assert (p_xlayers[2].type[0] in ['Transpose'])
        assert (p_xlayers[3].type[0] in ['Convolution'])
        assert (p_xlayers[4].type[0] in ['Pooling'])

        assert (p_xlayers[0].target == 'cpu')
        assert (p_xlayers[1].target == 'cpu')
        assert (p_xlayers[2].target == 'cpu')
        assert (p_xlayers[3].target == 'test')
        assert (p_xlayers[4].target == 'test')

        assert (p_xlayers[0].subgraph is None)
        assert (p_xlayers[1].subgraph is None)
        assert (p_xlayers[2].subgraph is None)
        assert (p_xlayers[3].subgraph == 'xp1')
        assert (p_xlayers[4].subgraph == 'xp1')

        subgraphs = TestXGraphPartitioner.xgraph_partitioner.get_subgraphs(
            p_xgraph)

        assert (len(subgraphs) == 1)
        xp1 = subgraphs[0]
        assert (xp1.name == 'xp1')
        xp1_xgraph = TestXGraphPartitioner.xgraph_factory\
            .build_from_xlayer(xp1.subgraph_data)

        assert (xp1.bottoms == ['t1'])
        assert (xp1.tops == [])
        assert (xp1.shapes == [[1, 2, 2, 2]])
        assert (xp1.sizes == [8])

        assert (len(xp1_xgraph) == 3)
        xp1_layers = xp1_xgraph.get_layers()

        assert (xp1_layers[0].type[0] == 'Input')
        assert (xp1_layers[0].layer[0] == 'conv2')
        assert (xp1_layers[1].type[0] == 'Convolution')
        assert (xp1_layers[2].type[0] == 'Pooling')

        assert (xp1_layers[0].bottoms == [])
        assert (xp1_layers[0].tops == ['conv2'])
        assert (xp1_layers[1].bottoms == ['xinput0'])
        assert (xp1_layers[1].tops == ['pool1'])
        assert (xp1_layers[2].bottoms == ['conv2'])
        assert (xp1_layers[2].tops == [])

    def test_two_partition_inputs(self):
        net = [
            XLayer(name='in1',
                   type=['Input'],
                   shapes=[1, 1, 4, 4],
                   sizes=[16],
                   bottoms=[],
                   tops=['conv1'],
                   layer=['in1'],
                   targets=[]),
            XLayer(name='in2',
                   type=['Input'],
                   shapes=[1, 1, 4, 4],
                   sizes=[16],
                   bottoms=[],
                   tops=['conv2'],
                   layer=['in2'],
                   targets=[]),
            XLayer(name='conv1',
                   type=['Convolution'],
                   shapes=[1, 2, 3, 3],
                   sizes=[18],
                   bottoms=['in1'],
                   tops=['pool1'],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   layer=['conv1'],
                   targets=[]),
            XLayer(name='pool1',
                   type=['Pooling'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=['conv1'],
                   tops=['concat1'],
                   layer=['pool1'],
                   targets=[]),
            XLayer(name='conv2',
                   type=['Convolution'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=['in2'],
                   tops=['concat1'],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   layer=['conv2'],
                   targets=[]),
            XLayer(name='concat1',
                   type=['Concat'],
                   shapes=[1, 4, 2, 2],
                   sizes=[16],
                   bottoms=['pool1', 'conv2'],
                   tops=['dense1'],
                   layer=['concat1'],
                   targets=[]),
            XLayer(name='dense1',
                   type=['Dense'],
                   shapes=[1, 20],
                   sizes=[],
                   bottoms=['concat1'],
                   tops=[],
                   layer=['dense1'],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   targets=[])
        ]
        xgraph = TestXGraphPartitioner.xgraph_factory.build_from_xlayer(net)
        TargetRegistry().annotate_ops(xgraph)
        p_xgraph = TestXGraphPartitioner.xgraph_partitioner.partition(
            xgraph, ['test'])

        assert len(p_xgraph.get_layer_names()) == 7
        assert p_xgraph.get_subgraph_names() == ['xp2']

        p_xlayers = p_xgraph.get_layers()

        assert p_xlayers[0].target == 'cpu'
        assert p_xlayers[1].target == 'test'
        assert p_xlayers[2].target == 'test'
        assert p_xlayers[3].target == 'cpu'
        assert p_xlayers[4].target == 'test'
        assert p_xlayers[5].target == 'test'
        assert p_xlayers[6].target == 'cpu'

        assert p_xlayers[0].subgraph is None
        assert p_xlayers[1].subgraph == 'xp2'
        assert p_xlayers[2].subgraph == 'xp2'
        assert p_xlayers[3].subgraph is None
        assert p_xlayers[4].subgraph == 'xp2'
        assert p_xlayers[5].subgraph == 'xp2'
        assert p_xlayers[6].subgraph is None

        subgraphs = TestXGraphPartitioner.xgraph_partitioner.get_subgraphs(
            p_xgraph)

        assert len(subgraphs) == 1
        xp2 = subgraphs[0]
        assert xp2.name == 'xp2'
        xp2_xgraph = TestXGraphPartitioner.xgraph_factory\
            .build_from_xlayer(xp2.subgraph_data)

        assert xp2.bottoms == ['in1', 'in2']
        assert xp2.tops == ['dense1']
        assert xp2.shapes == [[1, 4, 2, 2]]
        assert xp2.sizes == [16]

        assert len(xp2_xgraph) == 6
        xp2_layers = xp2_xgraph.get_layers()

        assert (xp2_layers[0].type[0] == 'Input')
        assert (xp2_layers[0].layer[0] == 'conv1')
        assert (xp2_layers[1].type[0] == 'Convolution')
        assert (xp2_layers[2].type[0] == 'Pooling')
        assert (xp2_layers[3].type[0] == 'Input')
        assert (xp2_layers[3].layer[0] == 'conv2')
        assert (xp2_layers[4].type[0] == 'Convolution')
        assert (xp2_layers[5].type[0] == 'Concat')

        assert (xp2_layers[0].bottoms == [])
        assert (xp2_layers[0].tops == ['conv1'])
        assert (xp2_layers[1].bottoms == ['xinput0'])
        assert (xp2_layers[1].tops == ['pool1'])
        assert (xp2_layers[2].bottoms == ['conv1'])
        assert (xp2_layers[2].tops == ['concat1'])
        assert (xp2_layers[3].bottoms == [])
        assert (xp2_layers[3].tops == ['conv2'])
        assert (xp2_layers[4].bottoms == ['xinput1'])
        assert (xp2_layers[4].tops == ['concat1'])
        assert (xp2_layers[5].bottoms == ['pool1', 'conv2'])
        assert (xp2_layers[5].tops == [])

    def test_inception_like_block(self):
        net = [
            XLayer(name='in1',
                   type=['Input'],
                   shapes=[1, 1, 4, 4],
                   sizes=[16],
                   bottoms=[],
                   tops=['concat1'],
                   layer=['in1'],
                   targets=[]),
            XLayer(name='in2',
                   type=['Input'],
                   shapes=[1, 1, 4, 4],
                   sizes=[16],
                   bottoms=[],
                   tops=['concat1'],
                   layer=['in2'],
                   targets=[]),
            XLayer(name='concat1',
                   type=['Concat'],
                   shapes=[1, 2, 4, 4],
                   sizes=[32],
                   bottoms=['in1', 'in2'],
                   tops=['conv1', 'conv2'],
                   layer=['concat1'],
                   targets=[]),
            XLayer(name='conv1',
                   type=['Convolution'],
                   shapes=[1, 4, 3, 3],
                   sizes=[],
                   bottoms=['concat1'],
                   tops=['pool1'],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   layer=['conv1'],
                   targets=[]),
            XLayer(name='pool1',
                   type=['Pooling'],
                   shapes=[1, 4, 2, 2],
                   sizes=[],
                   bottoms=['conv1'],
                   tops=['concat2'],
                   layer=['pool1'],
                   targets=[]),
            XLayer(name='conv2',
                   type=['Convolution'],
                   shapes=[1, 4, 2, 2],
                   sizes=[],
                   bottoms=['concat1'],
                   tops=['concat2'],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   layer=['conv2'],
                   targets=[]),
            XLayer(name='concat2',
                   type=['Concat'],
                   shapes=[1, 8, 2, 2],
                   sizes=[32],
                   bottoms=['pool1', 'conv2'],
                   tops=['dense1'],
                   layer=['concat2'],
                   targets=[]),
            XLayer(name='dense1',
                   type=['Dense'],
                   shapes=[1, 20],
                   sizes=[20],
                   bottoms=['concat2'],
                   tops=[],
                   layer=['dense1'],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   targets=[])
        ]
        xgraph = TestXGraphPartitioner.xgraph_factory.build_from_xlayer(net)
        TargetRegistry().annotate_ops(xgraph)

        p_xgraph = TestXGraphPartitioner.xgraph_partitioner.partition(
            xgraph, ['test'])

        assert (len(p_xgraph.get_layer_names()) == 8)
        p_xlayers = p_xgraph.get_layers()

        assert (p_xlayers[0].target == 'cpu')
        assert (p_xlayers[1].target == 'cpu')
        assert (p_xlayers[2].target == 'test')
        assert (p_xlayers[3].target == 'test')
        assert (p_xlayers[4].target == 'test')
        assert (p_xlayers[5].target == 'test')
        assert (p_xlayers[6].target == 'test')
        assert (p_xlayers[7].target == 'cpu')

        assert (p_xlayers[0].subgraph is None)
        assert (p_xlayers[1].subgraph is None)
        assert (p_xlayers[2].subgraph == 'xp0')
        assert (p_xlayers[3].subgraph == 'xp0')
        assert (p_xlayers[4].subgraph == 'xp0')
        assert (p_xlayers[5].subgraph == 'xp0')
        assert (p_xlayers[6].subgraph == 'xp0')
        assert (p_xlayers[7].subgraph is None)

        subgraphs = TestXGraphPartitioner.xgraph_partitioner.get_subgraphs(
            p_xgraph)

        assert (len(subgraphs) == 1)
        xp0 = subgraphs[0]
        assert (xp0.name == 'xp0')
        xp0_xgraph = TestXGraphPartitioner.xgraph_factory\
            .build_from_xlayer(xp0.subgraph_data)

        assert (xp0.bottoms == ['in1', 'in2'])
        assert (xp0.tops == ['dense1'])
        assert (xp0.shapes == [[1, 8, 2, 2]])
        assert (xp0.sizes == [32])

        assert (len(xp0_xgraph) == 7)
        xp0_layers = xp0_xgraph.get_layers()

        assert (xp0_layers[0].type[0] == 'Input')
        assert (xp0_layers[0].layer[0] == 'concat1')
        assert (xp0_layers[1].type[0] == 'Input')
        assert (xp0_layers[1].layer[0] == 'concat1')
        assert (xp0_layers[2].type[0] == 'Concat')
        assert (xp0_layers[3].type[0] == 'Convolution')
        assert (xp0_layers[4].type[0] == 'Pooling')
        assert (xp0_layers[5].type[0] == 'Convolution')
        assert (xp0_layers[6].type[0] == 'Concat')

        assert (xp0_layers[0].bottoms == [])
        assert (xp0_layers[0].tops == ['concat1'])
        assert (xp0_layers[1].bottoms == [])
        assert (xp0_layers[1].tops == ['concat1'])
        assert (xp0_layers[2].bottoms == ['xinput0', 'xinput1'])
        assert (xp0_layers[2].tops == ['conv1', 'conv2'])
        assert (xp0_layers[3].bottoms == ['concat1'])
        assert (xp0_layers[3].tops == ['pool1'])
        assert (xp0_layers[4].bottoms == ['conv1'])
        assert (xp0_layers[4].tops == ['concat2'])
        assert (xp0_layers[5].bottoms == ['concat1'])
        assert (xp0_layers[5].tops == ['concat2'])
        assert (xp0_layers[6].bottoms == ['pool1', 'conv2'])
        assert (xp0_layers[6].tops == [])

    def test_top_tensors_basic(self):
        net = [
            XLayer(name='in1',
                   type=['Input'],
                   shapes=[1, 1, 4, 4],
                   sizes=[16],
                   bottoms=[],
                   tops=['conv1'],
                   layer=['in1'],
                   targets=[]),
            XLayer(name='conv1',
                   type=['Convolution'],
                   shapes=[1, 2, 3, 3],
                   sizes=[18],
                   bottoms=['in1'],
                   tops=['pool1'],
                   layer=['conv1'],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   targets=[]),
            XLayer(name='pool1',
                   type=['Pooling'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=['conv1'],
                   tops=['t1'],
                   layer=['pool1'],
                   targets=[]),
            XLayer(name='t1',
                   type=['Transpose'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=['pool1'],
                   tops=['s1'],
                   layer=['t1'],
                   internal=1,
                   targets=[],
                   attrs={'axes': [0, 2, 3, 1]}),
            XLayer(name='s1',
                   type=['Sqrt'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=['t1'],
                   tops=[],
                   layer=['s1'],
                   internal=0,
                   targets=[])
        ]
        xgraph = TestXGraphPartitioner.xgraph_factory.build_from_xlayer(net)
        TargetRegistry().annotate_ops(xgraph)
        p_xgraph = TestXGraphPartitioner.xgraph_partitioner.partition(
            xgraph, ['test'])

        assert (len(p_xgraph.get_layer_names()) == 5)
        assert (p_xgraph.get_subgraph_names() == ['xp0'])

        p_xlayers = p_xgraph.get_layers()
        assert (p_xlayers[0].type[0] in ['Input'])
        assert (p_xlayers[1].type[0] in ['Convolution'])
        assert (p_xlayers[2].type[0] in ['Pooling'])
        assert (p_xlayers[3].type[0] in ['Transpose'])
        assert (p_xlayers[4].type[0] in ['Sqrt'])

        assert (p_xlayers[0].target == 'cpu')
        assert (p_xlayers[1].target == 'test')
        assert (p_xlayers[2].target == 'test')
        assert (p_xlayers[3].target == 'cpu')
        assert (p_xlayers[4].target == 'cpu')

        assert (p_xlayers[0].subgraph is None)
        assert (p_xlayers[1].subgraph == 'xp0')
        assert (p_xlayers[2].subgraph == 'xp0')
        assert (p_xlayers[3].subgraph is None)
        assert (p_xlayers[4].subgraph is None)

        subgraphs = TestXGraphPartitioner.xgraph_partitioner.get_subgraphs(
            p_xgraph)

        assert len(subgraphs) == 1
        xp0 = subgraphs[0]
        assert xp0.name == 'xp0'
        xp0_xgraph = TestXGraphPartitioner.xgraph_factory\
            .build_from_xlayer(xp0.subgraph_data)

        assert xp0.bottoms == ['in1']
        assert xp0.tops == ['t1']
        assert xp0.shapes == [[1, 2, 2, 2]]
        assert xp0.sizes == [8]
        assert len(xp0_xgraph) == 3

        __bottom_tensors = xp0.attrs['__bottom_tensors']
        orig_bottom_tensors = xp0.attrs['orig_bottom_tensors']

        assert len(__bottom_tensors) == 1
        assert 'xinput0' in __bottom_tensors
        assert __bottom_tensors['xinput0'] == ['in1']

        assert len(orig_bottom_tensors) == 1
        assert 'xinput0' in orig_bottom_tensors
        assert orig_bottom_tensors['xinput0'] == ['in1']

        __top_tensors = xp0.attrs['__top_tensors']
        orig_top_tensors = xp0.attrs['orig_top_tensors']

        assert len(__top_tensors) == 1
        assert 'pool1' in __top_tensors
        assert __top_tensors['pool1'] == ['t1']

        assert len(orig_top_tensors) == 1
        assert 'pool1' in orig_top_tensors
        assert orig_top_tensors['pool1'] == ['s1']

    def test_multi_top_tensors(self):
        net = [
            XLayer(name='in1',
                   type=['Input'],
                   shapes=[1, 1, 4, 4],
                   sizes=[16],
                   bottoms=[],
                   tops=['conv1'],
                   layer=['in1'],
                   targets=[]),
            XLayer(name='conv1',
                   type=['Convolution'],
                   shapes=[1, 2, 3, 3],
                   sizes=[18],
                   bottoms=['in1'],
                   tops=['pool1'],
                   layer=['conv1'],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   targets=[]),
            XLayer(name='pool1',
                   type=['Pooling'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=['conv1'],
                   tops=['t1', 't2'],
                   layer=['pool1'],
                   targets=[]),
            XLayer(name='t1',
                   type=['Transpose'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=['pool1'],
                   tops=['s1'],
                   layer=['t1'],
                   internal=1,
                   targets=[],
                   attrs={'axes': [0, 2, 3, 1]}),
            XLayer(name='t2',
                   type=['Transpose'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=['pool1'],
                   tops=['s2', 's3'],
                   layer=['t2'],
                   internal=1,
                   targets=[],
                   attrs={'axes': [0, 2, 3, 1]}),
            XLayer(name='s1',
                   type=['Sqrt'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=['t1'],
                   tops=[],
                   layer=['s1'],
                   internal=0,
                   targets=[]),
            XLayer(name='s2',
                   type=['Sqrt'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=['t2'],
                   tops=[],
                   layer=['s2'],
                   internal=0,
                   targets=[]),
            XLayer(name='s3',
                   type=['Sqrt'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=['t2'],
                   tops=[],
                   layer=['s3'],
                   internal=0,
                   targets=[])
        ]
        xgraph = TestXGraphPartitioner.xgraph_factory.build_from_xlayer(net)
        TargetRegistry().annotate_ops(xgraph)
        p_xgraph = TestXGraphPartitioner.xgraph_partitioner.partition(
            xgraph, ['test'])

        assert len(p_xgraph.get_layer_names()) == 8
        assert p_xgraph.get_subgraph_names() == ['xp0']

        p_xlayers = p_xgraph.get_layers()
        assert p_xlayers[0].type[0] in ['Input']
        assert p_xlayers[1].type[0] in ['Convolution']
        assert p_xlayers[2].type[0] in ['Pooling']
        assert p_xlayers[3].type[0] in ['Transpose']
        assert p_xlayers[4].type[0] in ['Sqrt']
        assert p_xlayers[5].type[0] in ['Transpose']
        assert p_xlayers[6].type[0] in ['Sqrt']
        assert p_xlayers[7].type[0] in ['Sqrt']

        assert p_xlayers[0].target == 'cpu'
        assert p_xlayers[1].target == 'test'
        assert p_xlayers[2].target == 'test'
        assert p_xlayers[3].target == 'cpu'
        assert p_xlayers[4].target == 'cpu'
        assert p_xlayers[5].target == 'cpu'
        assert p_xlayers[6].target == 'cpu'
        assert p_xlayers[7].target == 'cpu'

        assert p_xlayers[0].subgraph is None
        assert p_xlayers[1].subgraph == 'xp0'
        assert p_xlayers[2].subgraph == 'xp0'
        assert p_xlayers[3].subgraph is None
        assert p_xlayers[4].subgraph is None
        assert p_xlayers[5].subgraph is None
        assert p_xlayers[6].subgraph is None
        assert p_xlayers[7].subgraph is None

        subgraphs = TestXGraphPartitioner.xgraph_partitioner.get_subgraphs(
            p_xgraph)

        assert len(subgraphs) == 1
        xp0 = subgraphs[0]
        assert xp0.name == 'xp0'
        xp0_xgraph = TestXGraphPartitioner.xgraph_factory\
            .build_from_xlayer(xp0.subgraph_data)

        assert xp0.bottoms == ['in1']
        assert xp0.tops == ['t1', 't2']
        assert xp0.shapes == [[1, 2, 2, 2], [1, 2, 2, 2]]
        assert xp0.sizes == [8, 8]
        assert len(xp0_xgraph) == 3

        __bottom_tensors = xp0.attrs['__bottom_tensors']
        orig_bottom_tensors = xp0.attrs['orig_bottom_tensors']

        assert len(__bottom_tensors) == 1
        assert 'xinput0' in __bottom_tensors
        assert __bottom_tensors['xinput0'] == ['in1']

        assert len(orig_bottom_tensors) == 1
        assert 'xinput0' in orig_bottom_tensors
        assert orig_bottom_tensors['xinput0'] == ['in1']

        __top_tensors = xp0.attrs['__top_tensors']
        orig_top_tensors = xp0.attrs['orig_top_tensors']

        assert len(__top_tensors) == 1
        assert 'pool1' in __top_tensors
        assert __top_tensors['pool1'] == ['t1', 't2']

        assert len(orig_top_tensors) == 1
        assert 'pool1' in orig_top_tensors
        assert orig_top_tensors['pool1'] == ['s1', 's2', 's3']
Example #12
0
class DPUCompiler(XGraphBaseCompiler):
    """ TODO """
    xgraph_partitioner = XGraphPartitioner()
    xgraph_factory = XGraphFactory()
    tf_generator = TfGenerator()

    def __init__(self,
                 xgraph,
                 target,
                 arch,
                 work_dir=os.path.join(os.getcwd(), 'work'),
                 build_dir=None,
                 mode='debug'):

        super(DPUCompiler, self).__init__(xgraph)

        if not os.path.isfile(arch):
            raise ValueError("Arch file: {} does not exist".format(arch))

        warnings.warn("This compilation only works with one network"
                      " configuration at the moment!!")

        q_output = self.xgraph.get_quantizer_output()
        self.netcfgs = {
            q_key: q_output.get_orig_pb(q_key)
            for q_key in q_output.keys()
        }
        self.quant_info = {
            q_key: q_output.get_q_info(q_key)
            for q_key in q_output.keys()
        }
        assert (len(self.netcfgs) == 1)
        self.work_dir = work_dir
        self.build_dir = build_dir if build_dir is not None else work_dir
        self.target = target
        self.arch = arch
        self.mode = mode
        self.c_output = CompilerOutput(name=xgraph.get_name())

    def Getopts(self, input_shapes):
        return {
            "maximumasrelu": True,
            "pipelineconvmaxpool": False,
            "bytesperpixels": 1,
            "dsp": 96,
            "memory": 9,
            "ddr": "256",
            "cpulayermustgo": True,
            "forceweightsfullyconnected": True,
            "mixmemorystrategy": True,
            "maximumasrelu": True,
            "pipelineconvmaxpool": True,
            'bridges': ['bytype', 'Concat'],
            "usedeephi": True,
            'placeholdershape': input_shapes
        }

    def compile(self):
        # type: () -> None
        """ """
        layout_transform_pass = \
            XGraphLayoutTransformationPass('NHWC', target=self.target)
        self.xgraph = layout_transform_pass.execute(self.xgraph,
                                                    subgraphs_only=False)

        # netcfg = list(self.netcfgs.values())[0]  # orig pb file
        quant_info_file = list(self.quant_info.values())[0]  # quant info file

        subxg_layers = DPUCompiler.xgraph_partitioner\
            .get_subgraphs(self.xgraph)[0].subgraph_data
        xgraph = DPUCompiler.xgraph_factory.build_from_xlayer(subxg_layers)
        net_name = list(self.netcfgs.keys())[0]
        fs = DPUCompiler.tf_generator.generate(xgraph,
                                               'graph',
                                               subgraphs_only=True,
                                               layout='NHWC',
                                               batch_size=1,
                                               placeholder=True,
                                               out_dir=self.work_dir)
        netcfg = list(fs.values())[0]

        input_names = xgraph.get_input_names()
        input_shapes = [
            xgraph.get(in_name).shapes.tolist()[:] for in_name in input_names
        ]
        output_names = xgraph.get_output_names()
        output_shapes = [
            xgraph.get(out_name).shapes.tolist()[:]
            for out_name in output_names
        ]
        if len(input_names) > 1:
            raise NotImplementedError(
                "DPUCompiler only handles models with"
                " one input at the moment but found: {}".format(
                    len(input_names)))
        opt_input_shapes = {
            in_name: [e if e != -1 else 1 for e in input_shape]
            for in_name, input_shape in zip(input_names, input_shapes)
        }
        opts = self.Getopts(opt_input_shapes)
        if not os.path.isfile(quant_info_file):
            raise ValueError(
                "quant file: {} does not exist".format(quant_info_file))
        opts['quant_cfgfile'] = quant_info_file
        opts = str(opts)
        command = """
            vai_c_tensorflow \
            --frozen_pb {} \
            --arch {} \
            --output_dir {} \
            --net_name {}\
            --options "{}"
        """.format(netcfg, self.arch, self.build_dir, 'compiler', opts)
        logger.info("command: {}".format(command))
        process = subprocess.Popen(command,
                                   shell=True,
                                   cwd=FILE_PATH,
                                   stdout=subprocess.PIPE)
        output, error = process.communicate()
        if output is not None:
            output = output.decode('utf-8')
            if 'SUCCESSFUL COMPILATION' not in output:
                logger.info(output)
                raise ValueError('compiler is failed. Please see the log for'
                                 ' more details')
        if error is not None:
            error = error.decode('utf-8')
            # raise ValueError(error)

        logger.debug("Output: {}".format(output))
        logger.debug("Error: {}".format(error))
        compiler_json_file = self.build_dir + '/compiler.json'
        with open(compiler_json_file) as json_file:
            json_graph = json.load(json_file)
        graph_inputs = json_graph["inputs"]
        graph_outputs = json_graph["outputs"]
        logger.debug("{} {}".format(input_names, graph_inputs))
        logger.debug("{} {}".format(output_names, graph_outputs))

        in_map = {in_name: in_name for in_name in input_names}
        out_node_merged = []
        out_nodes = [
            graph_output['previous_layers'][0]
            for graph_output in graph_outputs
        ]
        for i in range(len(out_nodes)):
            out_node_merged.append([
                layer['merged'][-1] for layer in json_graph['network']
                if layer['name'] == out_nodes[i]
            ][0])
        in_map = {in_name: in_name for in_name in input_names}
        out_map = {out_name: t for out_name, t in zip(output_names, out_nodes)}
        #out_map = {out_name: out_name for out_name in output_names}

        self.c_output.add(net_name, ['dpuv1lib.so'], in_map, out_map)
        self.xgraph.set_compiler_output(self.c_output)

        # TODO
        self.xgraph.meta_attrs['compiled'] = True
        self.xgraph.meta_attrs['compiler_libs'] = ['dpuv1lib.so']
        self.xgraph.meta_attrs['compiler_in_map'] = in_map
        self.xgraph.meta_attrs['compiler_out_map'] = out_map

        return self.xgraph
Example #13
0
class TestBaseSubgraphQuantizer(unittest.TestCase):

    xgraph_factory = XGraphFactory()
    xgraph_partitioner = XGraphPartitioner()

    @classmethod
    def setUpClass(cls):
        def xgraph_build_func(xgraph):
            raise NotImplementedError("")

        def xgraph_optimizer(xgraph):
            raise NotImplementedError("")

        def xgraph_quantizer(xgraph):
            raise NotImplementedError("")

        def xgraph_compiler(xgraph):
            raise NotImplementedError("")

        target_registry = TargetRegistry()
        target_registry.register_target('test', xgraph_optimizer,
                                        xgraph_quantizer, xgraph_compiler,
                                        xgraph_build_func)

        @register_op_support_check('test', 'Pooling')
        def pooling_op_support(X, bXs, tXs):
            return True

        @register_op_support_check('test', 'Concat')
        def concat_op_support(X, bXs, tXs):
            return True

    @classmethod
    def tearDownClass(cls):

        target_registry = TargetRegistry()
        target_registry.unregister_target('test')

    def test_one_subgraph(self):

        W = np.reshape(
            np.array([[[1, 1], [0, 1]], [[3, 4], [-1, 0]]], dtype=np.float32),
            (2, 1, 2, 2))
        B = np.array([1., -1.], dtype=np.float32)

        net = [
            XLayer(name='in1',
                   type=['Input'],
                   shapes=[-1, 1, 4, 4],
                   sizes=[16],
                   bottoms=[],
                   tops=['conv1'],
                   layer=['in1'],
                   targets=[]),
            XLayer(name='conv1',
                   type=['Convolution'],
                   shapes=[-1, 2, 3, 3],
                   sizes=[18],
                   bottoms=['in1'],
                   tops=['pool1'],
                   layer=['conv2d0'],
                   data=ConvData(W, B),
                   attrs={
                       'data_layout': 'NCHW',
                       'kernel_layout': 'OIHW',
                       'shape': [1, 2, 3, 3],
                       'padding': [[0, 0], [0, 0], [0, 0], [0, 0]],
                       'strides': [1, 1],
                       'dilation': [1, 1],
                       'groups': 1
                   },
                   targets=[]),
            XLayer(
                name='pool1',
                type=['Pooling'],
                shapes=[1, 2, 2, 2],
                sizes=[8],
                bottoms=['conv1'],
                tops=['pool2'],
                layer=['pool1'],
                targets=[],
                attrs={
                    'padding': [[0, 0], [0, 0], [0, 0], [0, 0]],
                    'strides': [1, 1],
                    'kernel_size': [2, 2],
                    'insize': [3, 3],
                    # HW
                    'outsize': [2, 2],
                    'data_layout': 'NCHW',
                    'pool_type': 'Max'
                }),
            XLayer(
                name='pool2',
                type=['Pooling'],
                shapes=[1, 2, 1, 1],
                sizes=[2],
                bottoms=['pool1'],
                tops=[],
                layer=['pool2'],
                targets=[],
                attrs={
                    'padding': [[0, 0], [0, 0], [0, 0], [0, 0]],
                    'strides': [1, 1],
                    'kernel_size': [2, 2],
                    'insize': [2, 2],
                    # HW
                    'outsize': [1, 1],
                    'data_layout': 'NCHW',
                    'pool_type': 'Avg'
                })
        ]
        xgraph = TestBaseSubgraphQuantizer.xgraph_factory\
            .build_from_xlayer(net)
        p_xgraph = TestBaseSubgraphQuantizer.xgraph_partitioner.partition(
            xgraph, ['test'])

        assert (len(p_xgraph.get_layer_names()) == 4)
        assert (p_xgraph.get_subgraph_names() == ['xp0'])

        p_xlayers = p_xgraph.get_layers()
        assert (p_xlayers[0].type[0] in ['Input'])
        assert (p_xlayers[1].type[0] in ['Convolution'])
        assert (p_xlayers[2].type[0] in ['Pooling'])
        assert (p_xlayers[3].type[0] in ['Pooling'])

        inputs = np.reshape(
            np.array([[1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1]],
                     dtype=np.float32), (1, 1, 4, 4))

        def inputs_func(iter):
            return {'in1': inputs}

        quantizer = BaseSubgraphQuantizerTest(xgraph=p_xgraph,
                                              inputs_func=inputs_func)
        quantizer.quantize()

        assert ('xp0' in quantizer.test_inputs)
        assert ('xinput0' in quantizer.test_inputs['xp0'])

        expected = np.reshape(
            np.array([[[4, 4, 4], [4, 4, 4], [4, 4, 4]],
                      [[5, 5, 5], [5, 5, 5], [5, 5, 5]]]), (1, 2, 3, 3))
        np.testing.assert_array_equal(quantizer.test_inputs['xp0']['xinput0'],
                                      expected)

    def test_multiple_subgraph(self):

        W = np.reshape(
            np.array([[[1, 1], [0, 1]], [[3, 4], [-1, 0]]], dtype=np.float32),
            (2, 1, 2, 2))
        B = np.array([1., -1.], dtype=np.float32)

        W2 = np.reshape(
            np.array([[[1, 1], [0, 1]], [[3, 4], [-1, 0]]], dtype=np.float32),
            (1, 2, 2, 2))
        B2 = np.array([-1.], dtype=np.float32)

        net = [
            XLayer(name='in1',
                   type=['Input'],
                   shapes=[-1, 1, 4, 4],
                   sizes=[16],
                   bottoms=[],
                   tops=['conv1'],
                   layer=['in1'],
                   targets=[]),
            XLayer(name='conv1',
                   type=['Convolution'],
                   shapes=[-1, 2, 3, 3],
                   sizes=[18],
                   bottoms=['in1'],
                   tops=['pool1'],
                   layer=['conv2d0'],
                   data=ConvData(W, B),
                   attrs={
                       'data_layout': 'NCHW',
                       'kernel_layout': 'OIHW',
                       'shape': [1, 2, 3, 3],
                       'padding': [[0, 0], [0, 0], [0, 0], [0, 0]],
                       'strides': [1, 1],
                       'dilation': [1, 1],
                       'groups': 1
                   },
                   targets=[]),
            XLayer(
                name='pool1',
                type=['Pooling'],
                shapes=[1, 2, 2, 2],
                sizes=[8],
                bottoms=['conv1'],
                tops=['conv2'],
                layer=['pool1'],
                targets=[],
                attrs={
                    'padding': [[0, 0], [0, 0], [0, 0], [0, 0]],
                    'strides': [1, 1],
                    'kernel_size': [2, 2],
                    'insize': [3, 3],
                    # HW
                    'outsize': [2, 2],
                    'data_layout': 'NCHW',
                    'pool_type': 'Max'
                }),
            XLayer(name='conv2',
                   type=['Convolution'],
                   shapes=[-1, 1, 1, 1],
                   sizes=[18],
                   bottoms=['pool1'],
                   tops=['pool2'],
                   layer=['conv2'],
                   data=ConvData(W2, B2),
                   attrs={
                       'data_layout': 'NCHW',
                       'kernel_layout': 'OIHW',
                       'shape': [1, 1, 1, 1],
                       'padding': [[0, 0], [0, 0], [0, 0], [0, 0]],
                       'strides': [1, 1],
                       'dilation': [1, 1],
                       'groups': 1
                   },
                   targets=[]),
            XLayer(
                name='pool2',
                type=['Pooling'],
                shapes=[1, 1, 1, 1],
                sizes=[2],
                bottoms=['conv2'],
                tops=[],
                layer=['pool2'],
                targets=[],
                attrs={
                    'padding': [[0, 0], [0, 0], [0, 0], [0, 0]],
                    'strides': [1, 1],
                    'kernel_size': [1, 1],
                    'insize': [1, 1],
                    # HW
                    'outsize': [1, 1],
                    'data_layout': 'NCHW',
                    'pool_type': 'Avg'
                })
        ]
        xgraph = TestBaseSubgraphQuantizer.xgraph_factory\
            .build_from_xlayer(net)
        p_xgraph = TestBaseSubgraphQuantizer.xgraph_partitioner.partition(
            xgraph, ['test'])

        assert (len(p_xgraph.get_layer_names()) == 5)
        assert (set(p_xgraph.get_subgraph_names()) == set(['xp0']))

        p_xlayers = p_xgraph.get_layers()
        assert (p_xlayers[0].type[0] in ['Input'])
        assert (p_xlayers[1].type[0] in ['Convolution'])
        assert (p_xlayers[2].type[0] in ['Pooling'])
        assert (p_xlayers[3].type[0] in ['Convolution'])
        assert (p_xlayers[4].type[0] in ['Pooling'])

        inputs = np.reshape(
            np.array([[1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1]],
                     dtype=np.float32), (1, 1, 4, 4))

        def inputs_func(iter):
            return {'in1': inputs}

        quantizer = BaseSubgraphQuantizerTest(xgraph=p_xgraph,
                                              inputs_func=inputs_func)
        quantizer.quantize()

        assert 'xp0' in quantizer.test_inputs
        assert 'xinput0' in quantizer.test_inputs['xp0']

        expected = np.reshape(
            np.array([[[[4, 4, 4], [4, 4, 4], [4, 4, 4]],
                       [[5, 5, 5], [5, 5, 5], [5, 5, 5]]]]), (1, 2, 3, 3))
        np.testing.assert_array_equal(quantizer.test_inputs['xp0']['xinput0'],
                                      expected)
Example #14
0
class TestDPUContrib(unittest.TestCase):

    xgraph_partitioner = XGraphPartitioner()
    xgraph_factory = XGraphFactory()
    target_registry = TargetRegistry()
    rt_manager = RtManager()

    @classmethod
    def setUpClass(cls):
        # Import DPU module
        from pyxir.contrib.dpuv1 import dpuv1

        # from pyxir.contrib.dpuv1.dpuv1_target import\
        #     xgraph_dpu_v1_optimizer,\
        #     xgraph_dpu_v1_quantizer,\
        #     xgraph_dpu_v1_compiler,\
        #     xgraph_dpu_v1_build_func

        # pyxir.register_target(
        #     'dpuv1',
        #     xgraph_dpu_v1_optimizer,
        #     xgraph_dpu_v1_quantizer,
        #     xgraph_dpu_v1_compiler,
        #     xgraph_dpu_v1_build_func
        # )

    @classmethod
    def tearDownClass(cls):
        # Unregister dpu for other tests
        TestDPUContrib.target_registry.unregister_target("dpuv1")
        TestDPUContrib.target_registry.unregister_target("DPUCADX8G")

    def test_supported_ops(self):
        dpuv1_ops = TestDPUContrib.target_registry.get_supported_op_check_names(
            "dpuv1")

        assert "BatchNorm" in dpuv1_ops
        assert "BiasAdd" in dpuv1_ops
        assert "Concat" in dpuv1_ops
        assert "Convolution" in dpuv1_ops
        assert "Conv2DTranspose" in dpuv1_ops
        assert "DPU" in dpuv1_ops
        assert "Eltwise" in dpuv1_ops
        assert "Pad" in dpuv1_ops
        assert "Pooling" in dpuv1_ops
        assert "Mean" in dpuv1_ops
        assert "pReLU" in dpuv1_ops
        assert "ReLU" in dpuv1_ops
        assert "Scale" in dpuv1_ops

    @unittest.skipIf(
        skip_tf,
        "Skipping Tensorflow related test because tensorflow is"
        "not available",
    )
    def test_import_ext_quantizer(self):
        if TestDPUContrib.target_registry.is_target("DPUCADX8G"):
            TestDPUContrib.target_registry.unregister_target("DPUCADX8G")
        if TestDPUContrib.rt_manager.exists_op("cpu-np", "DPU"):
            TestDPUContrib.rt_manager.unregister_op("cpu-np", "DPU")
        from pyxir.contrib.target import DPUCADX8G_external_quantizer
Example #15
0
class VAICompiler(XGraphBaseCompiler):

    """ Vitis-AI compiler wrapper for DPUCAHX8H """

    xgraph_partitioner = XGraphPartitioner()
    xgraph_factory = XGraphFactory()

    def __init__(self,
                 xgraph,
                 arch,
                 work_dir=os.path.join(os.getcwd(), 'work'),
                 build_dir=os.getcwd(),
                 mode='debug'):

        super(VAICompiler, self).__init__(xgraph)

        if not os.path.isfile(arch):
            raise ValueError("Arch file: {} does not exist".format(arch))

        self.arch = arch


        q_output = self.xgraph.get_quantizer_output()
        self.netcfgs = {q_key: q_output.get_q_file(q_key)
                        for q_key in q_output.keys()}
        assert(len(self.netcfgs) == 1)
        self.work_dir = work_dir

        if not os.path.exists(self.work_dir):
            os.makedirs(self.work_dir)
        self.build_dir = build_dir if build_dir is not None else work_dir
        if not os.path.exists(self.build_dir):
            os.makedirs(self.build_dir)
        self.mode = mode
        self.c_output = CompilerOutput(name=xgraph.get_name())
        

    def compile(self) -> None:
        """ Start DPUCAHX8H compilation """

        net_name = list(self.netcfgs.keys())[0]
        netcfg = list(self.netcfgs.values())[0]

        # We only handle one partition at the moment
        Xp = VAICompiler.xgraph_partitioner\
            .get_subgraphs(self.xgraph)[0]
        subxg_layers = Xp.subgraph_data
        xgraph = VAICompiler.xgraph_factory.build_from_xlayer(subxg_layers)
        # assert xgraph.get_name() == net_name

        input_names = xgraph.get_input_names()
        input_shapes = [xgraph.get(in_name).shapes[:]
                        for in_name in input_names]
        output_names = list(Xp.attrs['__top_tensors'].keys()) # xgraph.get_output_names()
        output_shapes = [xgraph.get(out_name).shapes[:]
                         for out_name in output_names]

        if len(input_names) > 1:
            raise NotImplementedError("VAICompiler only handles models with"
                                      " one input at the moment but found: {}"
                                      .format(len(input_names)))

        netcfg=netcfg.replace('deploy_model.pb', 'quantize_eval_model.pb')
        command = """
        vai_c_tensorflow \
            --frozen_pb {} \
            --arch {} \
            --output_dir {} \
            --net_name {} \
            --options "{}"
        """.format(netcfg, self.arch, self.build_dir, net_name, str(dict()))

        logger.info("Command: {}".format(command))

        process = subprocess.Popen(command,
                                   shell=True,
                                   cwd=FILE_PATH,
                                   stdout=subprocess.PIPE)

        output, error = process.communicate()
        logger.debug("{} {}".format(output, error))


        if error is not None:
            error = error.decode('utf-8')
            raise ValueError(error)

        in_map = {in_name: in_name for in_name in input_names}
        out_map = {out_name: out_name for out_name in output_names}
        self.c_output.add(net_name, ['libvart-runner.so'], in_map, out_map)
        self.xgraph.set_compiler_output(self.c_output)

        # TODO
        self.xgraph.meta_attrs['compiled'] = True
        self.xgraph.meta_attrs['compiler_libs'] = ['libvart-runner.so']
        self.xgraph.meta_attrs['compiler_in_map'] = in_map
        self.xgraph.meta_attrs['compiler_out_map'] = out_map

        return self.xgraph
Example #16
0
class XGraphBaseSubgraphQuantizer(XGraphBaseQuantizer):
    """
    Base class for quantization of Xfgraphs with partitioned subgraphs

    Attributes
    ----------
    xgraph: XGraph
        the XGraph instance to be quantized
    inputs_func: function
        the input function to be used for retrieving model inputs
    work_dir: str
        the directory to be used for writing files
    quant_iter:
        the number of quantization iterations to be done
    """

    xgraph_partitioner = XGraphPartitioner()
    xgraph_factory = XGraphFactory()

    def __init__(self,
                 xgraph,
                 inputs_func,
                 work_dir=os.path.join(os.getcwd()),
                 quant_iter=1):
        #
        super(XGraphBaseSubgraphQuantizer, self).__init__(xgraph)

        self.subgraph_Xps = XGraphBaseSubgraphQuantizer.xgraph_partitioner\
            .get_subgraphs(self.xgraph)

        # Maps external (graph) to internal (subgraph) inputs for each subgraph
        self.subgraph_input_map = {}
        self.subgraph_inputs = {}
        self.subgraph_input_names = []
        for Xp in self.subgraph_Xps:
            sub_xgraph = XGraphBaseSubgraphQuantizer.xgraph_factory.\
                build_from_xlayer(Xp.subgraph_data, name=Xp.name)

            self.subgraph_input_map[Xp.name] = {}

            input_names = sub_xgraph.get_input_names()

            for b, in_name in zip(Xp.bottoms, input_names):
                self.subgraph_input_names.append(b)
                self.subgraph_inputs[b] = None
                self.subgraph_input_map[Xp.name][b] = in_name

        # Setup executable graph
        self.runtime = pyxir.build(self.xgraph,
                                   target='cpu',
                                   last_layers=self.subgraph_input_names)

        self.inputs_func = inputs_func
        self.work_dir = work_dir
        os.makedirs(self.work_dir, exist_ok=True)
        self.quant_iter = quant_iter

    def quantize_subgraph(self, xgraph, inputs, input_names, output_names):
        # type: (XGraph, Dict[str, numpy.ndarray]) -> None
        """ Quantize a subgraph with given calibration inputs """
        raise NotImplementedError("")

    def quantize(self):
        # type: () -> None
        """ Start quantization of the partitioned xgraph """

        input_names = self.runtime.get_input_names()
        input_shapes = self.runtime.get_input_shapes()
        assert len(input_names) == len(input_shapes)
        if len(input_names) != 1:
            raise NotImplementedError(
                "Invalid number of inputs to model: {},{}, Vitis-AI"
                " quantization only supports models with one input at the"
                " moment".format(len(input_names), input_names))
        input_name = input_names[0]
        input_shape = input_shapes[0]

        logger.debug("START Compute subgraph inputs for quantization")
        for it in range(self.quant_iter):
            inputs = self.inputs_func(it)

            subgraph_inpts = self.runtime.run(
                inputs, outputs=self.subgraph_input_names)

            for in_name, inpt in \
                    zip(self.subgraph_input_names, subgraph_inpts):
                self.subgraph_inputs[in_name] = inpt

            # for idx, layer, inpts, outpt, _ in \
            #         self.runtime.run_stepwise(inputs):

            #     si = 0
            #     if layer.name in self.subgraph_inputs:
            #         self.subgraph_inputs[layer.name] = outpt
            #         si += 1

            #     if len(self.subgraph_inputs) == si:
            #         break

        logger.debug("START Subgraph quantization")
        for Xp in self.subgraph_Xps:
            # Create sub XGraph from subgraph layer subgraph_data
            sub_xgraph = XGraphBaseSubgraphQuantizer.xgraph_factory.\
                build_from_xlayer(Xp.subgraph_data, name=Xp.name)

            input_names = list(Xp.attrs['__bottom_tensors'].keys())
            output_names = list(Xp.attrs['__top_tensors'].keys())

            original_input_names = \
                list(self.subgraph_input_map[Xp.name].keys())
            inputs = {
                self.subgraph_input_map[Xp.name][in_name]:
                self.subgraph_inputs[in_name]
                for in_name in original_input_names
            }
            self.quantize_subgraph(sub_xgraph, inputs, input_names,
                                   output_names)

        logger.debug("STOP Subgraph quantization")

        return self.xgraph
Example #17
0
class TestXGraphPartitioner(unittest.TestCase):

    xgraph_partitioner = XGraphPartitioner()
    xgraph_factory = XGraphFactory()

    @classmethod
    def setUpClass(cls):
        def xgraph_build_func(xgraph):
            raise NotImplementedError("")

        def xgraph_optimizer(xgraph):
            raise NotImplementedError("")

        def xgraph_quantizer(xgraph):
            raise NotImplementedError("")

        def xgraph_compiler(xgraph):
            raise NotImplementedError("")

        target_registry = TargetRegistry()
        target_registry.register_target(
            "test",
            xgraph_optimizer,
            xgraph_quantizer,
            xgraph_compiler,
            xgraph_build_func,
        )

        @register_op_support_check("test", "Convolution")
        def conv_op_support(X, bXs, tXs):
            return True

        @register_op_support_check("test", "Pooling")
        def pooling_op_support(X, bXs, tXs):
            return True

        @register_op_support_check("test", "Concat")
        def concat_op_support(X, bXs, tXs):
            return False

        @register_op_support_check("test", "Eltwise")
        def eltwise_op_support(X, bXs, tXs):
            return True

        @register_op_support_check("test", "ReLU")
        def relu_op_support(X, bXs, tXs):
            return True

    @classmethod
    def tearDownClass(cls):

        target_registry = TargetRegistry()
        target_registry.unregister_target("test")

    def test_basic(self):
        x1 = px.ops.input("in1", shape=[1, 1, 4, 4])
        x2 = px.ops.input("in2", shape=[1, 2, 2, 2])
        w1 = px.ops.constant("weight", np.ones((2, 1, 2, 2), dtype=np.float32))
        conv = px.ops.conv2d(
            op_name="conv1",
            input_layer=x1,
            weights_layer=w1,
            kernel_size=[2, 2],
            strides=[1, 1],
            padding_hw=[0, 0, 0, 0],
            dilation=[1, 1],
            groups=1,
            channels=2,
            data_layout="NCHW",
            kernel_layout="OIHW",
        )
        pool = px.ops.pool2d(
            op_name="pool1",
            input_layer=conv,
            pool_type="Avg",
            pool_size=[2, 2],
        )
        add = px.ops.eltwise("add1", pool, x2)
        net = [x1, x2, conv, pool, add]
        xgraph = TestXGraphPartitioner.xgraph_factory.build_from_xlayer(net)
        p_xgraph = px.partition(xgraph, ["test"])

        assert len(p_xgraph.get_layer_names()) == 5
        assert p_xgraph.get_subgraph_names() == ["xp0"]

        p_xlayers = p_xgraph.get_layers()
        assert p_xlayers[0].type[0] in ["Input"]
        assert p_xlayers[1].type[0] in ["Convolution"]
        assert p_xlayers[2].type[0] in ["Pooling"]
        assert p_xlayers[3].type[0] in ["Input"]
        assert p_xlayers[4].type[0] in ["Eltwise"]

        assert p_xlayers[0].target == "cpu"
        assert p_xlayers[1].target == "test"
        assert p_xlayers[2].target == "test"
        assert p_xlayers[3].target == "cpu"
        assert p_xlayers[4].target == "test"

        assert p_xlayers[0].subgraph is None
        assert p_xlayers[1].subgraph == "xp0"
        assert p_xlayers[2].subgraph == "xp0"
        assert p_xlayers[3].subgraph is None
        assert p_xlayers[4].subgraph == "xp0"

        subgraphs = TestXGraphPartitioner.xgraph_partitioner.get_subgraphs(
            p_xgraph)

        assert len(subgraphs) == 1
        xp0 = subgraphs[0]
        assert xp0.name == "xp0"
        xp0_xgraph = TestXGraphPartitioner.xgraph_factory.build_from_xlayer(
            xp0.subgraph_data)

        assert xp0.bottoms == ["in1", "in2"]
        assert xp0.tops == []
        assert xp0.shapes == [[-1, 2, 2, 2]]
        assert xp0.sizes == [8]

        assert len(xp0_xgraph) == 5
        xp0_layers = xp0_xgraph.get_layers()

        assert xp0_layers[0].type[0] == "Input"
        assert xp0_layers[0].layer[0] == "conv1"
        assert xp0_layers[1].type[0] == "Convolution"
        assert xp0_layers[2].type[0] == "Pooling"
        assert xp0_layers[3].type[0] == "Input"
        assert xp0_layers[4].type[0] == "Eltwise"

        assert xp0_layers[0].bottoms == []
        assert xp0_layers[0].tops == ["conv1"]
        assert xp0_layers[1].bottoms == ["xinput0"]
        assert xp0_layers[1].tops == ["pool1"]
        assert xp0_layers[2].bottoms == ["conv1"]
        assert xp0_layers[2].tops == ["add1"]

    def test_interrupt_partition_in_add_branch(self):
        x = px.ops.input("in1", shape=[1, 28, 28, 2028])
        w1 = px.ops.constant("weight",
                             np.ones((2048, 2048, 1, 1), dtype=np.float32))
        conv1 = px.ops.conv2d(
            op_name="conv1",
            input_layer=x,
            weights_layer=w1,
            kernel_size=[1, 1],
            strides=[1, 1],
            padding_hw=[0, 0, 0, 0],
            dilation=[1, 1],
            groups=1,
            channels=2048,
            data_layout="NHWC",
            kernel_layout="OIHW",
        )
        r1 = px.ops.relu("r1", [conv1])
        w2 = px.ops.constant("weight",
                             np.ones((512, 2048, 1, 1), dtype=np.float32))
        conv2 = px.ops.conv2d(
            op_name="conv2",
            input_layer=r1,
            weights_layer=w2,
            kernel_size=[1, 1],
            strides=[1, 1],
            padding_hw=[0, 0, 0, 0],
            dilation=[1, 1],
            groups=1,
            channels=512,
            data_layout="NHWC",
            kernel_layout="OIHW",
        )
        sigm = px.ops.sigmoid("sigm", [conv2])  # Unsupported layer
        w3 = px.ops.constant("weight",
                             np.ones((2048, 512, 1, 1), dtype=np.float32))
        conv3 = px.ops.conv2d(
            op_name="conv3",
            input_layer=sigm,
            weights_layer=w3,
            kernel_size=[1, 1],
            strides=[1, 1],
            padding_hw=[0, 0, 0, 0],
            dilation=[1, 1],
            groups=1,
            channels=2048,
            data_layout="NHWC",
            kernel_layout="OIHW",
        )  # Although this layer is supported, it should not be in the partition
        add = px.ops.eltwise(
            "add", r1, conv3
        )  # Although this layer is supported, it should not be in the partition
        net = [x, conv1, r1, conv2, sigm, conv3, add]
        xgraph = TestXGraphPartitioner.xgraph_factory.build_from_xlayer(net)
        p_xgraph = px.partition(xgraph, ["test"])

        assert len(p_xgraph.get_layer_names()) == 7
        assert p_xgraph.get_subgraph_names() == ["xp0"]

        p_xlayers = p_xgraph.get_layers()

        assert p_xgraph.get("in1").target == "cpu"
        assert p_xgraph.get("conv1").target == "test"
        assert p_xgraph.get("r1").target == "test"
        assert p_xgraph.get("conv2").target == "test"
        assert p_xgraph.get("sigm").target == "cpu"
        assert p_xgraph.get("conv3").target == "cpu"
        assert p_xgraph.get("add").target == "cpu"

        subgraphs = TestXGraphPartitioner.xgraph_partitioner.get_subgraphs(
            p_xgraph)

        assert len(subgraphs) == 1
        xp0 = subgraphs[0]
        assert xp0.name == "xp0"
        xp0_xgraph = TestXGraphPartitioner.xgraph_factory.build_from_xlayer(
            xp0.subgraph_data)

        assert xp0.bottoms == ["in1"]
        assert xp0.tops == ["add", "sigm"]
        assert xp0.shapes == [[-1, 28, 28, 2048], [-1, 28, 28, 512]]
        assert xp0.sizes == [28 * 28 * 2048, 28 * 28 * 512]

        assert len(xp0_xgraph) == 4
        xp0_layers = xp0_xgraph.get_layers()

        assert xp0_layers[0].type[0] == "Input"
        assert xp0_layers[0].layer[0] == "conv1"
        assert xp0_layers[1].type[0] == "Convolution"
        assert xp0_layers[2].type[0] == "ReLU"
        assert xp0_layers[3].type[0] == "Convolution"

    def test_complete_partition(self):
        x = px.ops.input("in1", shape=[1, 1, 4, 4])
        w1 = px.ops.constant("weight", np.ones((2, 1, 2, 2), dtype=np.float32))
        conv = px.ops.conv2d(
            op_name="conv1",
            input_layer=x,
            weights_layer=w1,
            kernel_size=[2, 2],
        )
        pool = px.ops.pool2d(
            op_name="pool1",
            input_layer=conv,
            pool_type="Avg",
            pool_size=[2, 2],
        )
        net = [x, conv, pool]
        xgraph = TestXGraphPartitioner.xgraph_factory.build_from_xlayer(net)
        p_xgraph = px.partition(xgraph, ["test"])

        assert len(p_xgraph.get_layer_names()) == 3
        assert p_xgraph.get_subgraph_names() == ["xp0"]

        p_xlayers = p_xgraph.get_layers()
        assert p_xlayers[0].type[0] in ["Input"]
        assert p_xlayers[1].type[0] in ["Convolution"]
        assert p_xlayers[2].type[0] in ["Pooling"]

        assert p_xlayers[0].target == "cpu"
        assert p_xlayers[1].target == "test"
        assert p_xlayers[2].target == "test"

        assert p_xlayers[0].subgraph is None
        assert p_xlayers[1].subgraph == "xp0"
        assert p_xlayers[2].subgraph == "xp0"

        subgraphs = TestXGraphPartitioner.xgraph_partitioner.get_subgraphs(
            p_xgraph)

        assert len(subgraphs) == 1
        xp0 = subgraphs[0]
        assert xp0.name == "xp0"
        xp0_xgraph = TestXGraphPartitioner.xgraph_factory.build_from_xlayer(
            xp0.subgraph_data)

        assert xp0.bottoms == ["in1"]
        assert xp0.tops == []
        assert xp0.shapes == [[-1, 2, 2, 2]]
        assert xp0.sizes == [8]
        assert xp0.attrs["target"] == "test"
        assert xp0.attrs["__bottom_tensors"] == {"xinput0": ["in1"]}
        assert xp0.attrs["orig_bottom_tensors"] == {"xinput0": ["in1"]}
        assert xp0.attrs["__top_tensors"] == {"pool1": []}
        assert xp0.attrs["orig_top_tensors"] == {"pool1": []}

        assert len(xp0_xgraph) == 3
        xp0_layers = xp0_xgraph.get_layers()

        assert xp0_layers[0].type[0] == "Input"
        assert xp0_layers[0].layer[0] == "conv1"
        assert xp0_layers[1].type[0] == "Convolution"
        assert xp0_layers[2].type[0] == "Pooling"

        assert xp0_layers[0].bottoms == []
        assert xp0_layers[0].tops == ["conv1"]
        assert xp0_layers[1].bottoms == ["xinput0"]
        assert xp0_layers[1].tops == ["pool1"]
        assert xp0_layers[2].bottoms == ["conv1"]
        assert xp0_layers[2].tops == []

    def test_two_partitions_through_interruption(self):
        # A layer inside a residual type branch os not supported
        # Here: BatchNorm
        x1 = px.ops.input("in1", shape=[1, 1, 4, 4])
        w1 = px.ops.constant("weight", np.ones((2, 1, 2, 2), dtype=np.float32))
        conv1 = px.ops.conv2d(
            op_name="conv1",
            input_layer=x1,
            weights_layer=w1,
            kernel_size=[2, 2],
        )  # 1, 2, 3, 3
        pool = px.ops.pool2d(
            op_name="pool1",
            input_layer=conv1,
            pool_type="Avg",
            pool_size=[2, 2],
            padding=[1, 1, 0, 0],
        )  # 1, 2, 3, 3
        bn_mean = px.ops.constant("mean", np.ones((2, ), dtype=np.float32))
        bn_var = px.ops.constant("var", np.ones((2, ), dtype=np.float32))
        bn_gamma = px.ops.constant("gamma", np.ones((2, ), dtype=np.float32))
        bn_beta = px.ops.constant("beta", np.ones((2, ), dtype=np.float32))
        bn = px.ops.batch_norm(
            op_name="bn1",
            input_layer=conv1,
            mean_layer=bn_mean,
            variance_layer=bn_var,
            gamma_layer=bn_gamma,
            beta_layer=bn_beta,
            axis=1,
        )  # 1, 2, 3, 3
        concat = px.ops.concat("concat1", [pool, bn], axis=1)  # 1, 4, 3, 3
        w2 = px.ops.constant("weight2", np.ones((6, 2, 2, 2),
                                                dtype=np.float32))
        conv2 = px.ops.conv2d(
            op_name="conv2",
            input_layer=concat,
            weights_layer=w2,
            kernel_size=[2, 2],
            padding_hw=[1, 1, 0, 0],
        )  # 1, 6, 3, 3
        net = [x1, conv1, pool, bn, concat, conv2]
        xgraph = TestXGraphPartitioner.xgraph_factory.build_from_xlayer(net)
        p_xgraph = px.partition(xgraph, ["test"])

        assert len(p_xgraph.get_layer_names()) == 6
        assert p_xgraph.get_subgraph_names() == ["xp0"]

        p_xlayers = p_xgraph.get_layers()
        assert p_xlayers[0].type[0] in ["Input"]
        assert p_xlayers[1].type[0] in ["Convolution"]
        assert p_xlayers[2].type[0] in ["Pooling"]
        assert p_xlayers[3].type[0] in ["BatchNorm"]
        assert p_xlayers[4].type[0] in ["Concat"]
        assert p_xlayers[5].type[0] in ["Convolution"]

        assert p_xlayers[0].target == "cpu"
        assert p_xlayers[1].target == "test"
        assert p_xlayers[2].target == "test"
        assert p_xlayers[3].target == "cpu"
        assert p_xlayers[4].target == "cpu"
        assert p_xlayers[5].target == "cpu"

        assert p_xlayers[0].subgraph is None
        assert p_xlayers[1].subgraph == "xp0"
        assert p_xlayers[2].subgraph == "xp0"
        assert p_xlayers[3].subgraph is None
        assert p_xlayers[4].subgraph is None
        assert p_xlayers[5].subgraph is None

        assert p_xlayers[3].name == "bn1"
        assert p_xlayers[3].bottoms == ["conv1"]
        assert p_xlayers[3].tops == ["concat1"]

        assert p_xlayers[4].name == "concat1"
        assert p_xlayers[4].bottoms == ["pool1", "bn1"]
        assert p_xlayers[4].tops == ["conv2"]

        subgraphs = TestXGraphPartitioner.xgraph_partitioner.get_subgraphs(
            p_xgraph)

        assert len(subgraphs) == 1
        xp0 = subgraphs[0]
        assert xp0.name == "xp0"
        xp0_xgraph = TestXGraphPartitioner.xgraph_factory.build_from_xlayer(
            xp0.subgraph_data)

        assert xp0.bottoms == ["in1"]
        assert xp0.tops == ["bn1", "concat1"]
        assert xp0.shapes == [[-1, 2, 3, 3], [-1, 2, 3, 3]]
        assert xp0.sizes == [18, 18]
        assert xp0.attrs["target"] == "test"
        assert xp0.attrs["__bottom_tensors"] == {"xinput0": ["in1"]}
        assert xp0.attrs["orig_bottom_tensors"] == {"xinput0": ["in1"]}
        assert xp0.attrs["__top_tensors"] == {
            "conv1": ["bn1"],
            "pool1": ["concat1"]
        }
        assert xp0.attrs["orig_top_tensors"] == {
            "conv1": ["bn1"],
            "pool1": ["concat1"]
        }

        assert len(xp0_xgraph) == 3
        xp0_layers = xp0_xgraph.get_layers()

        assert [X.name for X in xp0_xgraph.get_input_layers()] == ["xinput0"]
        # TODO: XGraph only recognizes output layers when they have no top
        #   layers
        assert [X.name for X in xp0_xgraph.get_output_layers()] == ["pool1"]

        assert xp0_layers[0].type[0] == "Input"
        assert xp0_layers[0].layer[0] == "conv1"
        assert xp0_layers[1].type[0] == "Convolution"
        assert xp0_layers[2].type[0] == "Pooling"

        assert xp0_layers[0].bottoms == []
        assert xp0_layers[0].tops == ["conv1"]
        assert xp0_layers[1].bottoms == ["xinput0"]
        assert xp0_layers[1].tops == ["pool1"]
        assert xp0_layers[2].bottoms == ["conv1"]
        assert xp0_layers[2].tops == []

    def test_multiple_partitions(self):
        x1 = px.ops.input("in1", shape=[1, 1, 4, 4])
        x2 = px.ops.input("in2", shape=[1, 2, 2, 2])
        w1 = px.ops.constant("weight", np.ones((2, 1, 2, 2), dtype=np.float32))
        conv1 = px.ops.conv2d(
            op_name="conv1",
            input_layer=x1,
            weights_layer=w1,
            kernel_size=[2, 2],
        )  # 1, 2, 3, 3
        pool = px.ops.pool2d(
            op_name="pool1",
            input_layer=conv1,
            pool_type="Avg",
            pool_size=[2, 2],
        )  # 1, 2, 2, 2
        add = px.ops.eltwise("add1", pool, x2)  # 1, 2, 2, 2
        bn_mean = px.ops.constant("mean", np.ones((2, ), dtype=np.float32))
        bn_var = px.ops.constant("var", np.ones((2, ), dtype=np.float32))
        bn_gamma = px.ops.constant("gamma", np.ones((2, ), dtype=np.float32))
        bn_beta = px.ops.constant("beta", np.ones((2, ), dtype=np.float32))
        bn = px.ops.batch_norm(
            op_name="bn1",
            input_layer=add,
            mean_layer=bn_mean,
            variance_layer=bn_var,
            gamma_layer=bn_gamma,
            beta_layer=bn_beta,
            axis=1,
        )  # 1, 2, 3, 3
        pool2 = px.ops.pool2d(
            op_name="pool2",
            input_layer=bn,
            pool_type="Avg",
            pool_size=[2, 2],
            padding=[0, 0, 1, 1],
        )  # 1, 2, 2, 2
        net = [x1, x2, conv1, pool, add, bn, pool2]
        xgraph = TestXGraphPartitioner.xgraph_factory.build_from_xlayer(net)
        p_xgraph = px.partition(xgraph, ["test"])

        assert len(p_xgraph.get_layer_names()) == 7
        # ! Only xp0 because only one subgraph can exist for now (largest)
        assert set(p_xgraph.get_subgraph_names()) == set(["xp0"])

        p_xlayers = p_xgraph.get_layers()
        assert p_xlayers[0].type[0] in ["Input"]
        assert p_xlayers[1].type[0] in ["Convolution"]
        assert p_xlayers[2].type[0] in ["Pooling"]
        assert p_xlayers[3].type[0] in ["Input"]
        assert p_xlayers[4].type[0] in ["Eltwise"]
        assert p_xlayers[5].type[0] in ["BatchNorm"]
        assert p_xlayers[6].type[0] in ["Pooling"]

        assert p_xlayers[0].target == "cpu"
        assert p_xlayers[1].target == "test"
        assert p_xlayers[2].target == "test"
        assert p_xlayers[3].target == "cpu"
        assert p_xlayers[4].target == "test"
        assert p_xlayers[5].target == "cpu"
        # ! CPU because only one subgraph can exist for now (largest)
        assert p_xlayers[6].target == "cpu"

        assert p_xlayers[0].subgraph is None
        assert p_xlayers[1].subgraph == "xp0"
        assert p_xlayers[2].subgraph == "xp0"
        assert p_xlayers[3].subgraph is None
        assert p_xlayers[4].subgraph == "xp0"
        assert p_xlayers[5].subgraph is None
        assert p_xlayers[6].subgraph is None

        subgraphs = TestXGraphPartitioner.xgraph_partitioner.get_subgraphs(
            p_xgraph)

        assert len(subgraphs) == 1
        xp0 = subgraphs[0]
        assert xp0.name == "xp0"
        xp0_xgraph = TestXGraphPartitioner.xgraph_factory.build_from_xlayer(
            xp0.subgraph_data)

        assert xp0.bottoms == ["in1", "in2"]
        assert xp0.tops == ["bn1"]
        assert xp0.shapes == [[-1, 2, 2, 2]]
        assert xp0.sizes == [8]

        assert len(xp0_xgraph) == 5
        xp0_layers = xp0_xgraph.get_layers()

        assert xp0_layers[0].type[0] == "Input"
        assert xp0_layers[0].layer[0] == "conv1"
        assert xp0_layers[1].type[0] == "Convolution"
        assert xp0_layers[2].type[0] == "Pooling"
        assert xp0_layers[3].type[0] == "Input"
        assert xp0_layers[4].type[0] == "Eltwise"

        assert xp0_layers[0].bottoms == []
        assert xp0_layers[0].tops == ["conv1"]
        assert xp0_layers[1].bottoms == ["xinput0"]
        assert xp0_layers[1].tops == ["pool1"]
        assert xp0_layers[2].bottoms == ["conv1"]
        assert xp0_layers[2].tops == ["add1"]

    def test_multiple_partitions_largest_last(self):
        x1 = px.ops.input("in1", shape=[1, 1, 4, 4])
        w1 = px.ops.constant("weight1", np.ones((2, 1, 2, 2),
                                                dtype=np.float32))
        conv1 = px.ops.conv2d(
            op_name="conv1",
            input_layer=x1,
            weights_layer=w1,
            kernel_size=[2, 2],
        )  # 1, 2, 3, 3
        t1 = px.ops.transpose("t1", conv1, axes=[0, 2, 3, 1])  # 1, 3, 3, 2
        w2 = px.ops.constant("weight2", np.ones((4, 2, 2, 2),
                                                dtype=np.float32))
        conv2 = px.ops.conv2d(
            op_name="conv2",
            input_layer=t1,
            weights_layer=w2,
            kernel_size=[2, 2],
            padding_hw=[1, 0, 1, 0],
            data_layout="NHWC",
        )  # 1, 3, 3, 4
        pool = px.ops.pool2d(
            op_name="pool1",
            input_layer=conv2,
            pool_type="Avg",
            pool_size=[2, 2],
            layout="NHWC",
        )  # 1, 2, 2, 4
        net = [x1, conv1, t1, conv2, pool]
        xgraph = TestXGraphPartitioner.xgraph_factory.build_from_xlayer(net)
        p_xgraph = px.partition(xgraph, ["test"])

        assert len(p_xgraph.get_layer_names()) == 5
        # ! Only xp1 because only one subgraph can exist for now (largest)
        assert set(p_xgraph.get_subgraph_names()) == set(["xp1"])

        p_xlayers = p_xgraph.get_layers()
        assert p_xlayers[0].type[0] in ["Input"]
        assert p_xlayers[1].type[0] in ["Convolution"]
        assert p_xlayers[2].type[0] in ["Transpose"]
        assert p_xlayers[3].type[0] in ["Convolution"]
        assert p_xlayers[4].type[0] in ["Pooling"]

        assert p_xlayers[0].target == "cpu"
        assert p_xlayers[1].target == "cpu"
        assert p_xlayers[2].target == "cpu"
        assert p_xlayers[3].target == "test"
        assert p_xlayers[4].target == "test"

        assert p_xlayers[0].subgraph is None
        assert p_xlayers[1].subgraph is None
        assert p_xlayers[2].subgraph is None
        assert p_xlayers[3].subgraph == "xp1"
        assert p_xlayers[4].subgraph == "xp1"

        subgraphs = TestXGraphPartitioner.xgraph_partitioner.get_subgraphs(
            p_xgraph)

        assert len(subgraphs) == 1
        xp1 = subgraphs[0]
        assert xp1.name == "xp1"
        xp1_xgraph = TestXGraphPartitioner.xgraph_factory.build_from_xlayer(
            xp1.subgraph_data)

        assert xp1.bottoms == ["t1"]
        assert xp1.tops == []
        assert xp1.shapes == [[-1, 2, 2, 4]]
        assert xp1.sizes == [16]

        assert len(xp1_xgraph) == 3
        xp1_layers = xp1_xgraph.get_layers()

        assert xp1_layers[0].type[0] == "Input"
        assert xp1_layers[0].layer[0] == "conv2"
        assert xp1_layers[1].type[0] == "Convolution"
        assert xp1_layers[2].type[0] == "Pooling"

        assert xp1_layers[0].bottoms == []
        assert xp1_layers[0].tops == ["conv2"]
        assert xp1_layers[1].bottoms == ["xinput0"]
        assert xp1_layers[1].tops == ["pool1"]
        assert xp1_layers[2].bottoms == ["conv2"]
        assert xp1_layers[2].tops == []

    def test_two_partition_inputs(self):
        x1 = px.ops.input("in1", shape=[1, 1, 4, 4])
        x2 = px.ops.input("in2", shape=[1, 1, 4, 4])
        w1 = px.ops.constant("weight", np.ones((2, 1, 2, 2), dtype=np.float32))
        conv1 = px.ops.conv2d(
            op_name="conv1",
            input_layer=x1,
            weights_layer=w1,
            kernel_size=[2, 2],
        )  # 1, 2, 3, 3
        pool = px.ops.pool2d(
            op_name="pool1",
            input_layer=conv1,
            pool_type="Avg",
            pool_size=[3, 3],
        )  # 1, 2, 1, 1
        w2 = px.ops.constant("weight2", np.ones((2, 1, 4, 4),
                                                dtype=np.float32))
        conv2 = px.ops.conv2d(
            op_name="conv2",
            input_layer=x2,
            weights_layer=w2,
            kernel_size=[4, 4],
            strides=[1, 1],
        )  # 1, 2, 1, 1
        add = px.ops.eltwise("add1", pool, conv2)
        reshape = px.ops.reshape("reshape1", add, [-1, 2])
        wd = px.ops.constant("weight_dense", np.ones((20, 2),
                                                     dtype=np.float32))
        dense = px.ops.dense(
            op_name="dense1",
            input_layer=reshape,
            weights_layer=wd,
            units=20,
        )
        net = [x1, x2, conv1, pool, conv2, add, reshape, dense]
        xgraph = TestXGraphPartitioner.xgraph_factory.build_from_xlayer(net)
        p_xgraph = px.partition(xgraph, ["test"])

        assert len(p_xgraph.get_layer_names()) == 8
        assert p_xgraph.get_subgraph_names() == ["xp2"]

        p_xlayers = p_xgraph.get_layers()

        assert p_xlayers[0].target == "cpu"
        assert p_xlayers[1].target == "test"
        assert p_xlayers[2].target == "test"
        assert p_xlayers[3].target == "cpu"
        assert p_xlayers[4].target == "test"
        assert p_xlayers[5].target == "test"
        assert p_xlayers[6].target == "cpu"
        assert p_xlayers[7].target == "cpu"

        assert p_xlayers[0].subgraph is None
        assert p_xlayers[1].subgraph == "xp2"
        assert p_xlayers[2].subgraph == "xp2"
        assert p_xlayers[3].subgraph is None
        assert p_xlayers[4].subgraph == "xp2"
        assert p_xlayers[5].subgraph == "xp2"
        assert p_xlayers[6].subgraph is None
        assert p_xlayers[7].subgraph is None

        subgraphs = TestXGraphPartitioner.xgraph_partitioner.get_subgraphs(
            p_xgraph)

        assert len(subgraphs) == 1
        xp2 = subgraphs[0]
        assert xp2.name == "xp2"
        xp2_xgraph = TestXGraphPartitioner.xgraph_factory.build_from_xlayer(
            xp2.subgraph_data)

        assert xp2.bottoms == ["in1", "in2"]
        assert xp2.tops == ["reshape1"]
        assert xp2.shapes == [[-1, 2, 1, 1]]
        assert xp2.sizes == [2]

        assert len(xp2_xgraph) == 6
        xp2_layers = xp2_xgraph.get_layers()

        assert xp2_layers[0].type[0] == "Input"
        assert xp2_layers[0].layer[0] == "conv1"
        assert xp2_layers[1].type[0] == "Convolution"
        assert xp2_layers[2].type[0] == "Pooling"
        assert xp2_layers[3].type[0] == "Input"
        assert xp2_layers[3].layer[0] == "conv2"
        assert xp2_layers[4].type[0] == "Convolution"
        assert xp2_layers[5].type[0] == "Eltwise"

        assert xp2_layers[0].bottoms == []
        assert xp2_layers[0].tops == ["conv1"]
        assert xp2_layers[1].bottoms == ["xinput0"]
        assert xp2_layers[1].tops == ["pool1"]
        assert xp2_layers[2].bottoms == ["conv1"]
        assert xp2_layers[2].tops == ["add1"]
        assert xp2_layers[3].bottoms == []
        assert xp2_layers[3].tops == ["conv2"]
        assert xp2_layers[4].bottoms == ["xinput1"]
        assert xp2_layers[4].tops == ["add1"]
        assert xp2_layers[5].bottoms == ["pool1", "conv2"]
        assert xp2_layers[5].tops == []

    def test_inception_like_block(self):
        x1 = px.ops.input("in1", shape=[1, 1, 4, 4])
        x2 = px.ops.input("in2", shape=[1, 1, 4, 4])
        concat1 = px.ops.concat("concat1", [x1, x2], axis=1)  # 1, 2, 4, 4
        w1 = px.ops.constant("weight", np.ones((4, 2, 2, 2), dtype=np.float32))
        conv1 = px.ops.conv2d(
            op_name="conv1",
            input_layer=concat1,
            weights_layer=w1,
            kernel_size=[2, 2],
        )  # 1, 4, 3, 3
        pool1 = px.ops.pool2d(
            op_name="pool1",
            input_layer=conv1,
            pool_type="Avg",
            pool_size=[2, 2],
        )  # 1, 4, 2, 2
        w2 = px.ops.constant("weight2", np.ones((4, 2, 2, 2),
                                                dtype=np.float32))
        conv2 = px.ops.conv2d(
            op_name="conv2",
            input_layer=concat1,
            weights_layer=w2,
            kernel_size=[2, 2],
            strides=[2, 2],
        )  # 1, 4, 2, 2
        concat2 = px.ops.concat("concat2", [pool1, conv2],
                                axis=1)  # 1, 8, 2, 2
        net = [x1, x2, concat1, conv1, pool1, conv2, concat2]
        xgraph = TestXGraphPartitioner.xgraph_factory.build_from_xlayer(net)
        p_xgraph = px.partition(xgraph, ["test"])

        assert len(p_xgraph.get_layer_names()) == 7
        p_xlayers = p_xgraph.get_layers()

        assert p_xlayers[0].target == "cpu"
        assert p_xlayers[1].target == "cpu"
        assert p_xlayers[2].target == "cpu"
        assert p_xlayers[3].target == "test"
        assert p_xlayers[4].target == "test"
        assert p_xlayers[5].target == "cpu"
        assert p_xlayers[6].target == "cpu"

        assert p_xlayers[0].subgraph is None
        assert p_xlayers[1].subgraph is None
        assert p_xlayers[2].subgraph is None
        assert p_xlayers[3].subgraph == "xp0"
        assert p_xlayers[4].subgraph == "xp0"
        assert p_xlayers[5].subgraph is None
        assert p_xlayers[6].subgraph is None

        subgraphs = TestXGraphPartitioner.xgraph_partitioner.get_subgraphs(
            p_xgraph)

        assert len(subgraphs) == 1
        xp0 = subgraphs[0]
        assert xp0.name == "xp0"
        xp0_xgraph = TestXGraphPartitioner.xgraph_factory.build_from_xlayer(
            xp0.subgraph_data)

        assert xp0.bottoms == ["concat1"]
        assert xp0.tops == ["concat2"]
        assert xp0.shapes == [[-1, 4, 2, 2]]
        assert xp0.sizes == [16]

        assert len(xp0_xgraph) == 3
        xp0_layers = xp0_xgraph.get_layers()

        assert xp0_layers[0].type[0] == "Input"
        assert xp0_layers[1].type[0] == "Convolution"
        assert xp0_layers[2].type[0] == "Pooling"

    def test_top_tensors_basic(self):
        x1 = px.ops.input("in1", shape=[1, 1, 4, 4])
        w1 = px.ops.constant("weight", np.ones((2, 1, 2, 2), dtype=np.float32))
        conv1 = px.ops.conv2d(
            op_name="conv1",
            input_layer=x1,
            weights_layer=w1,
            kernel_size=[2, 2],
        )  # 1, 2, 3, 3
        pool1 = px.ops.pool2d(
            op_name="pool1",
            input_layer=conv1,
            pool_type="Avg",
            pool_size=[2, 2],
        )  # 1, 2, 2, 2

        # ! internal=1
        t1 = px.ops.transpose("t1", pool1, axes=[0, 2, 3, 1], internal=1)
        s1 = px.ops.sqrt("s1", [t1])
        net = [x1, conv1, pool1, t1, s1]
        xgraph = TestXGraphPartitioner.xgraph_factory.build_from_xlayer(net)
        p_xgraph = px.partition(xgraph, ["test"])

        assert len(p_xgraph.get_layer_names()) == 5
        assert p_xgraph.get_subgraph_names() == ["xp0"]

        p_xlayers = p_xgraph.get_layers()
        assert p_xlayers[0].type[0] in ["Input"]
        assert p_xlayers[1].type[0] in ["Convolution"]
        assert p_xlayers[2].type[0] in ["Pooling"]
        assert p_xlayers[3].type[0] in ["Transpose"]
        assert p_xlayers[4].type[0] in ["Sqrt"]

        assert p_xlayers[0].target == "cpu"
        assert p_xlayers[1].target == "test"
        assert p_xlayers[2].target == "test"
        assert p_xlayers[3].target == "cpu"
        assert p_xlayers[4].target == "cpu"

        assert p_xlayers[0].subgraph is None
        assert p_xlayers[1].subgraph == "xp0"
        assert p_xlayers[2].subgraph == "xp0"
        assert p_xlayers[3].subgraph is None
        assert p_xlayers[4].subgraph is None

        subgraphs = TestXGraphPartitioner.xgraph_partitioner.get_subgraphs(
            p_xgraph)

        assert len(subgraphs) == 1
        xp0 = subgraphs[0]
        assert xp0.name == "xp0"
        xp0_xgraph = TestXGraphPartitioner.xgraph_factory.build_from_xlayer(
            xp0.subgraph_data)

        assert xp0.bottoms == ["in1"]
        assert xp0.tops == ["t1"]
        assert xp0.shapes == [[-1, 2, 2, 2]]
        assert xp0.sizes == [8]
        assert len(xp0_xgraph) == 3

        __bottom_tensors = xp0.attrs["__bottom_tensors"]
        orig_bottom_tensors = xp0.attrs["orig_bottom_tensors"]

        assert len(__bottom_tensors) == 1
        assert "xinput0" in __bottom_tensors
        assert __bottom_tensors["xinput0"] == ["in1"]

        assert len(orig_bottom_tensors) == 1
        assert "xinput0" in orig_bottom_tensors
        assert orig_bottom_tensors["xinput0"] == ["in1"]

        __top_tensors = xp0.attrs["__top_tensors"]
        orig_top_tensors = xp0.attrs["orig_top_tensors"]

        assert len(__top_tensors) == 1
        assert "pool1" in __top_tensors
        assert __top_tensors["pool1"] == ["t1"]

        assert len(orig_top_tensors) == 1
        assert "pool1" in orig_top_tensors
        assert orig_top_tensors["pool1"] == ["s1"]

    def test_multi_top_tensors(self):
        x1 = px.ops.input("in1", shape=[1, 1, 4, 4])
        w1 = px.ops.constant("weight", np.ones((2, 1, 2, 2), dtype=np.float32))
        conv1 = px.ops.conv2d(
            op_name="conv1",
            input_layer=x1,
            weights_layer=w1,
            kernel_size=[2, 2],
        )  # 1, 2, 3, 3
        pool1 = px.ops.pool2d(
            op_name="pool1",
            input_layer=conv1,
            pool_type="Avg",
            pool_size=[2, 2],
        )  # 1, 2, 2, 2
        t1 = px.ops.transpose("t1", pool1, axes=[0, 2, 3, 1], internal=1)
        t2 = px.ops.transpose("t2", pool1, axes=[0, 2, 3, 1], internal=1)
        s1 = px.ops.sqrt("s1", [t1])
        s2 = px.ops.sqrt("s2", [t2])
        s3 = px.ops.sqrt("s3", [t2])
        net = [x1, conv1, pool1, t1, t2, s1, s2, s3]

        xgraph = TestXGraphPartitioner.xgraph_factory.build_from_xlayer(net)
        p_xgraph = px.partition(xgraph, ["test"])

        assert len(p_xgraph.get_layer_names()) == 8
        assert p_xgraph.get_subgraph_names() == ["xp0"]

        p_xlayers = p_xgraph.get_layers()
        assert p_xlayers[0].type[0] in ["Input"]
        assert p_xlayers[1].type[0] in ["Convolution"]
        assert p_xlayers[2].type[0] in ["Pooling"]
        assert p_xlayers[3].type[0] in ["Transpose"]
        assert p_xlayers[4].type[0] in ["Sqrt"]
        assert p_xlayers[5].type[0] in ["Transpose"]
        assert p_xlayers[6].type[0] in ["Sqrt"]
        assert p_xlayers[7].type[0] in ["Sqrt"]

        assert p_xlayers[0].target == "cpu"
        assert p_xlayers[1].target == "test"
        assert p_xlayers[2].target == "test"
        assert p_xlayers[3].target == "cpu"
        assert p_xlayers[4].target == "cpu"
        assert p_xlayers[5].target == "cpu"
        assert p_xlayers[6].target == "cpu"
        assert p_xlayers[7].target == "cpu"

        assert p_xlayers[0].subgraph is None
        assert p_xlayers[1].subgraph == "xp0"
        assert p_xlayers[2].subgraph == "xp0"
        assert p_xlayers[3].subgraph is None
        assert p_xlayers[4].subgraph is None
        assert p_xlayers[5].subgraph is None
        assert p_xlayers[6].subgraph is None
        assert p_xlayers[7].subgraph is None

        subgraphs = TestXGraphPartitioner.xgraph_partitioner.get_subgraphs(
            p_xgraph)

        assert len(subgraphs) == 1
        xp0 = subgraphs[0]
        assert xp0.name == "xp0"
        xp0_xgraph = TestXGraphPartitioner.xgraph_factory.build_from_xlayer(
            xp0.subgraph_data)

        assert xp0.bottoms == ["in1"]
        assert xp0.tops == ["t1", "t2"]
        assert xp0.shapes == [[-1, 2, 2, 2], [-1, 2, 2, 2]]
        assert xp0.sizes == [8, 8]
        assert len(xp0_xgraph) == 3

        __bottom_tensors = xp0.attrs["__bottom_tensors"]
        orig_bottom_tensors = xp0.attrs["orig_bottom_tensors"]

        assert len(__bottom_tensors) == 1
        assert "xinput0" in __bottom_tensors
        assert __bottom_tensors["xinput0"] == ["in1"]

        assert len(orig_bottom_tensors) == 1
        assert "xinput0" in orig_bottom_tensors
        assert orig_bottom_tensors["xinput0"] == ["in1"]

        __top_tensors = xp0.attrs["__top_tensors"]
        orig_top_tensors = xp0.attrs["orig_top_tensors"]

        assert len(__top_tensors) == 1
        assert "pool1" in __top_tensors
        assert __top_tensors["pool1"] == ["t1", "t2"]

        assert len(orig_top_tensors) == 1
        assert "pool1" in orig_top_tensors
        assert orig_top_tensors["pool1"] == ["s1", "s2", "s3"]
Example #18
0
class TestDPUCZDX8G(unittest.TestCase):

    xgraph_partitioner = XGraphPartitioner()
    xgraph_factory = XGraphFactory()
    target_registry = TargetRegistry()
    rt_manager = RtManager()

    @classmethod
    def setUpClass(cls):

        # Import DPU module
        from pyxir.contrib.dpuv2 import dpuv2

    @classmethod
    def tearDownClass(cls):
        # Unregister dpu for other tests
        TestDPUCZDX8G.target_registry.unregister_target("dpuv2-zcu104")
        TestDPUCZDX8G.target_registry.unregister_target("dpuv2-zcu102")
        TestDPUCZDX8G.target_registry.unregister_target("DPUCZDX8G-zcu102")
        TestDPUCZDX8G.target_registry.unregister_target("DPUCZDX8G-zcu104")
        TestDPUCZDX8G.target_registry.unregister_target("dpuv2-ultra96")
        TestDPUCZDX8G.target_registry.unregister_target("DPUCZDX8G-ultra96")
        TestDPUCZDX8G.target_registry.unregister_target("dpuv2-som")
        TestDPUCZDX8G.target_registry.unregister_target("DPUCZDX8G-som")

    # @unittest.skipIf(skip_tf, "Skipping Tensorflow related test because tensorflow is"
    #                 "not available")
    # def test_import_ext_quantizer(self):
    #     if TestDPUCZDX8G.target_registry.is_target('DPUCZDX8G-ultra96'):
    #         TestDPUCZDX8G.target_registry.unregister_target('DPUCZDX8G-ultra96')
    #         TestDPUCZDX8G.target_registry.unregister_target('DPUCZDX8G-zcu104')
    #         TestDPUCZDX8G.target_registry.unregister_target('DPUCZDX8G-zcu102')
    #         TestDPUCZDX8G.target_registry.unregister_target('DPUCZDX8G-som')
    #     if TestDPUCZDX8G.rt_manager.exists_op('cpu-np', 'DPU'):
    #         TestDPUCZDX8G.rt_manager.unregister_op('cpu-np', 'DPU')
    #     from pyxir.contrib.target import DPUCZDX8G_external_quantizer

    @unittest.skipIf(not is_dpuczdx8g_vart_flow_enabled(),
                     "DPUCZDX8G VART test")
    def test_compile_conv2d_pool2d(self):
        xcompiler_conv2d_pool2d_nhwc_oihw_test(
            (1, 4, 4, 1),
            (2, 1, 2, 2),
            [0, 0],
            [1, 1],
            [1, 1],
            "Max",
            [2, 2],
            [0, 0],
            targets=[
                "DPUCZDX8G-zcu104",
                "DPUCZDX8G-zcu102",
                "DPUCZDX8G-ultra96",
                "DPUCZDX8G-som",
            ],
        )
        # Strided
        xcompiler_conv2d_pool2d_nhwc_oihw_test(
            (1, 4, 4, 1),
            (2, 1, 2, 2),
            [0, 0],
            [2, 2],
            [1, 1],
            "Max",
            [2, 2],
            [0, 0],
            targets=[
                "DPUCZDX8G-zcu104",
                "DPUCZDX8G-zcu102",
                "DPUCZDX8G-ultra96",
                "DPUCZDX8G-som",
            ],
        )
        xcompiler_conv2d_pool2d_nhwc_oihw_test(
            (1, 4, 4, 1),
            (2, 1, 2, 2),
            [0, 0],
            [3, 3],
            [1, 1],
            "Avg",
            [2, 2],
            [1, 1],
            targets=[
                "DPUCZDX8G-zcu104",
                "DPUCZDX8G-zcu102",
                "DPUCZDX8G-ultra96",
                "DPUCZDX8G-som",
            ],
        )
        # Padded
        xcompiler_conv2d_pool2d_nhwc_oihw_test(
            (1, 4, 4, 1),
            (2, 1, 2, 2),
            [1, 1],
            [1, 1],
            [1, 1],
            "Max",
            [4, 4],
            [0, 0],
            targets=[
                "DPUCZDX8G-zcu104",
                "DPUCZDX8G-zcu102",
                "DPUCZDX8G-ultra96",
                "DPUCZDX8G-som",
            ],
        )
        xcompiler_conv2d_pool2d_nhwc_oihw_test(
            (1, 8, 8, 1),
            (2, 1, 3, 3),
            [2, 2],
            [1, 1],
            [1, 1],
            "Avg",
            [4, 4],
            [0, 0],
            targets=[
                "DPUCZDX8G-zcu104",
                "DPUCZDX8G-zcu102",
                "DPUCZDX8G-ultra96",
                "DPUCZDX8G-som",
            ],
        )
        # Dilated
        xcompiler_conv2d_pool2d_nhwc_oihw_test(
            (1, 4, 4, 1),
            (2, 1, 2, 2),
            [1, 1],
            [1, 1],
            [2, 2],
            "Max",
            [2, 2],
            [0, 0],
            targets=[
                "DPUCZDX8G-zcu104",
                "DPUCZDX8G-zcu102",
                "DPUCZDX8G-ultra96",
                "DPUCZDX8G-som",
            ],
        )
        xcompiler_conv2d_pool2d_nhwc_oihw_test(
            (1, 10, 10, 1),
            (2, 1, 2, 2),
            [1, 1],
            [1, 1],
            [4, 4],
            "Max",
            [2, 2],
            [0, 0],
            targets=[
                "DPUCZDX8G-zcu104",
                "DPUCZDX8G-zcu102",
                "DPUCZDX8G-ultra96",
                "DPUCZDX8G-som",
            ],
        )
        xcompiler_conv2d_pool2d_nhwc_oihw_test(
            (1, 28, 28, 512),
            (512, 512, 3, 3),
            [2, 2, 2, 2],
            [1, 1],
            [2, 2],
            "Max",
            [2, 2],
            [0, 0],
            targets=[
                "DPUCZDX8G-zcu104",
                "DPUCZDX8G-zcu102",
                "DPUCZDX8G-ultra96",
                "DPUCZDX8G-som",
            ],
        )

    @unittest.skipIf(is_dpuczdx8g_vart_flow_enabled(),
                     "DPUCZDX8G DNNC/DNNDK test")
    def test_compile_conv2d_pool2d_dnnc(self):
        conv2d_pool2d_nhwc_oihw_test(
            (1, 4, 4, 1),
            (2, 1, 2, 2),
            [0, 0],
            [1, 1],
            [1, 1],
            "Max",
            [2, 2],
            [0, 0],
            targets=[
                "DPUCZDX8G-zcu104",
                "DPUCZDX8G-zcu102",
                "DPUCZDX8G-ultra96",
                "DPUCZDX8G-som",
            ],
        )
        # Strided
        conv2d_pool2d_nhwc_oihw_test(
            (1, 4, 4, 1),
            (2, 1, 2, 2),
            [0, 0],
            [2, 2],
            [1, 1],
            "Max",
            [2, 2],
            [0, 0],
            targets=[
                "DPUCZDX8G-zcu104",
                "DPUCZDX8G-zcu102",
                "DPUCZDX8G-ultra96",
                "DPUCZDX8G-som",
            ],
        )
        conv2d_pool2d_nhwc_oihw_test(
            (1, 4, 4, 1),
            (2, 1, 2, 2),
            [0, 0],
            [3, 3],
            [1, 1],
            "Avg",
            [2, 2],
            [1, 1],
            targets=[
                "DPUCZDX8G-zcu104",
                "DPUCZDX8G-zcu102",
                "DPUCZDX8G-ultra96",
                "DPUCZDX8G-som",
            ],
        )
        # Padded
        conv2d_pool2d_nhwc_oihw_test(
            (1, 4, 4, 1),
            (2, 1, 2, 2),
            [1, 1],
            [1, 1],
            [1, 1],
            "Max",
            [4, 4],
            [0, 0],
            targets=[
                "DPUCZDX8G-zcu104",
                "DPUCZDX8G-zcu102",
                "DPUCZDX8G-ultra96",
                "DPUCZDX8G-som",
            ],
        )
        conv2d_pool2d_nhwc_oihw_test(
            (1, 8, 8, 1),
            (2, 1, 3, 3),
            [2, 2],
            [1, 1],
            [1, 1],
            "Avg",
            [4, 4],
            [0, 0],
            targets=[
                "DPUCZDX8G-zcu104",
                "DPUCZDX8G-zcu102",
                "DPUCZDX8G-ultra96",
                "DPUCZDX8G-som",
            ],
        )
        # Dilated
        conv2d_pool2d_nhwc_oihw_test(
            (1, 4, 4, 1),
            (2, 1, 2, 2),
            [1, 1],
            [1, 1],
            [2, 2],
            "Max",
            [2, 2],
            [0, 0],
            targets=[
                "DPUCZDX8G-zcu104",
                "DPUCZDX8G-zcu102",
                "DPUCZDX8G-ultra96",
                "DPUCZDX8G-som",
            ],
        )
        conv2d_pool2d_nhwc_oihw_test(
            (1, 10, 10, 1),
            (2, 1, 2, 2),
            [1, 1],
            [1, 1],
            [4, 4],
            "Max",
            [2, 2],
            [0, 0],
            targets=[
                "DPUCZDX8G-zcu104",
                "DPUCZDX8G-zcu102",
                "DPUCZDX8G-ultra96",
                "DPUCZDX8G-som",
            ],
        )
        conv2d_pool2d_nhwc_oihw_test(
            (1, 28, 28, 512),
            (512, 512, 3, 3),
            [2, 2, 2, 2],
            [1, 1],
            [2, 2],
            "Max",
            [2, 2],
            [0, 0],
            targets=[
                "DPUCZDX8G-zcu104",
                "DPUCZDX8G-zcu102",
                "DPUCZDX8G-ultra96",
                "DPUCZDX8G-som",
            ],
        )

    @unittest.skipIf(not is_dpuczdx8g_vart_flow_enabled(),
                     "DPUCZDX8G VART test")
    def test_compile_depthwise_conv2d_pool2d(self):
        xcompiler_conv2d_pool2d_nhwc_oihw_test(
            (1, 3, 3, 8),
            (8, 1, 3, 3),
            [1, 1, 1, 1],
            [1, 1],
            [1, 1],
            "Max",
            [2, 2],
            [0, 0],
            conv_groups=8,
            targets=["DPUCZDX8G-zcu104", "DPUCZDX8G-zcu102", "DPUCZDX8G-som"],
            expected_nb_subgraphs=3,
        )

    @unittest.skipIf(not is_dpuczdx8g_vart_flow_enabled(),
                     "DPUCZDX8G VART test")
    def test_compile_scale_conv2d(self):
        # Standalone scale/batchnorm unsupported in DPUCAHX8H compiler
        xcompiler_scale_conv2d_nhwc_oihw_test(
            (1, 299, 299, 3),
            (64, 3, 7, 7),
            [3, 3],
            [2, 2],
            [1, 1],
            target="DPUCZDX8G-zcu104",
            expected_nb_subgraphs=3,
        )

    @unittest.skipIf(not is_dpuczdx8g_vart_flow_enabled(),
                     "DPUCZDX8G VART test")
    def test_compile_resnetv1_block(self):
        xcompiler_resnetv1_block_test(
            in_shape=(1, 112, 112, 64),
            pool_size=[3, 3],
            pool_strides=[2, 2],
            w1_shape=(256, 64, 1, 1),
            w2_shape=(64, 64, 1, 1),
            w3_shape=(64, 64, 3, 3),
            w4_shape=(256, 64, 1, 1),
            c3_padding=[1, 1, 1, 1],
            target="DPUCZDX8G-zcu104",
        )

    @unittest.skipIf(is_dpuczdx8g_vart_flow_enabled(),
                     "DPUCZDX8G DNNC/DNNDK test")
    def test_compile_conv2d_leaky_relu_dnnc(self):
        conv2d_leaky_relu_nhwc_oihw_test(
            (1, 4, 4, 1),
            (2, 1, 2, 2),
            [0, 0],
            [1, 1],
            [1, 1],
            targets=[
                "DPUCZDX8G-zcu104",
                "DPUCZDX8G-zcu102",
                "DPUCZDX8G-ultra96",
                "DPUCZDX8G-som",
            ],
        )

    @unittest.skipIf(not is_dpuczdx8g_vart_flow_enabled(),
                     "DPUCZDX8G VART test")
    def test_compile_conv2d_leaky_relu(self):
        xcompiler_conv2d_leaky_relu_nhwc_oihw_test(
            (1, 4, 4, 1),
            (2, 1, 2, 2),
            [0, 0],
            [1, 1],
            [1, 1],
            targets=[
                "DPUCZDX8G-zcu104",
                "DPUCZDX8G-zcu102",
                "DPUCZDX8G-ultra96",
                "DPUCZDX8G-som",
            ],
        )

    @unittest.skipIf(is_dpuczdx8g_vart_flow_enabled(),
                     "DPUCZDX8G DNNC/DNNDK test")
    def test_dnnc_out_names(self):
        multi_output_conv2d_naming_test(["nn.conv-134", "nn.relu-22"])
        multi_output_conv2d_naming_test(["nn..con.v--1.", "n.n.-relu-2-"])
        conv2d_pool2d_naming_test(["conv1", "nn.conv-1"],
                                  ["nn.pool1", "nn.pool-1"])

    def test_supported_ops(self):
        ultra96_ops = TestDPUCZDX8G.target_registry.get_supported_op_check_names(
            "dpuv2-ultra96")

        assert "BatchNorm" in ultra96_ops
        assert "BiasAdd" in ultra96_ops
        assert "Concat" in ultra96_ops
        assert "Convolution" in ultra96_ops
        assert "Conv2DTranspose" in ultra96_ops
        assert "DPU" in ultra96_ops
        assert "Eltwise" in ultra96_ops
        assert "Pad" in ultra96_ops
        assert "Pooling" in ultra96_ops
        assert "Mean" in ultra96_ops
        assert "pReLU" in ultra96_ops
        assert "ReLU" in ultra96_ops
        assert "ReLU6" in ultra96_ops
        assert "Scale" in ultra96_ops

        zcu102_ops = TestDPUCZDX8G.target_registry.get_supported_op_check_names(
            "dpuv2-zcu102")

        assert "BatchNorm" in zcu102_ops
        assert "BiasAdd" in zcu102_ops
        assert "Concat" in zcu102_ops
        assert "Convolution" in zcu102_ops
        assert "Conv2DTranspose" in zcu102_ops
        assert "DPU" in zcu102_ops
        assert "Eltwise" in zcu102_ops
        assert "Pad" in zcu102_ops
        assert "Pooling" in zcu102_ops
        assert "Mean" in zcu102_ops
        assert "pReLU" in zcu102_ops
        assert "ReLU" in zcu102_ops
        assert "ReLU6" in zcu102_ops
        assert "Scale" in zcu102_ops

        zcu104_ops = TestDPUCZDX8G.target_registry.get_supported_op_check_names(
            "dpuv2-zcu104")

        assert "BatchNorm" in zcu104_ops
        assert "BiasAdd" in zcu104_ops
        assert "Concat" in zcu104_ops
        assert "Convolution" in zcu104_ops
        assert "Conv2DTranspose" in zcu104_ops
        assert "DPU" in zcu104_ops
        assert "Eltwise" in zcu104_ops
        assert "Pad" in zcu104_ops
        assert "Pooling" in zcu104_ops
        assert "Mean" in zcu104_ops
        assert "pReLU" in zcu104_ops
        assert "ReLU" in zcu104_ops
        assert "ReLU6" in zcu104_ops
        assert "Scale" in zcu104_ops

        som_ops = TestDPUCZDX8G.target_registry.get_supported_op_check_names(
            "dpuv2-som")

        assert "BatchNorm" in som_ops
        assert "BiasAdd" in som_ops
        assert "Concat" in som_ops
        assert "Convolution" in som_ops
        assert "Conv2DTranspose" in som_ops
        assert "DPU" in som_ops
        assert "Eltwise" in som_ops
        assert "Pad" in som_ops
        assert "Pooling" in som_ops
        assert "Mean" in som_ops
        assert "pReLU" in som_ops
        assert "ReLU" in som_ops
        assert "ReLU6" in som_ops
        assert "Scale" in som_ops
        assert "Upsampling2D" in som_ops

    def test_small(self):
        net = [
            XLayer(
                name="in1",
                type=["Input"],
                shapes=[1, 1, 4, 4],
                sizes=[16],
                bottoms=[],
                tops=["conv1"],
                layer=["in1"],
                targets=[],
            ),
            XLayer(
                name="in2",
                type=["Input"],
                shapes=[1, 2, 2, 2],
                sizes=[8],
                bottoms=[],
                tops=["dense1"],
                layer=["in2"],
                targets=[],
            ),
            XLayer(
                name="conv1",
                type=["Convolution"],
                shapes=[1, 2, 3, 3],
                sizes=[18],
                bottoms=["in1"],
                tops=["pool1"],
                layer=["conv1"],
                data=ConvData(np.array([1, 1]), np.array([0, 0])),
                attrs={
                    "data_layout": "NCHW",
                    "padding": [[0, 0], [0, 0], [1, 1], [1, 1]],
                    "kernel_size": [3, 3],
                    "strides": [1, 1],
                    "dilation": [1, 1],
                    "groups": 1,
                    "channels": [2, 2],
                },
                targets=[],
            ),
            XLayer(
                name="pool1",
                type=["Pooling"],
                shapes=[1, 2, 2, 2],
                sizes=[8],
                bottoms=["conv1"],
                tops=["dense1"],
                layer=["pool1"],
                attrs={
                    "data_layout": "NCHW",
                    "padding": [[0, 0], [0, 0], [1, 1], [1, 1]],
                    "kernel_size": [3, 3],
                    "strides": [1, 1],
                },
                targets=[],
            ),
            XLayer(
                name="dense1",
                type=["Dense"],
                shapes=[1, 20],
                sizes=[20],
                bottoms=["pool1", "in2"],
                tops=[],
                data=ConvData(np.array([1, 1]), np.array([0, 0])),
                layer=["dense1"],
                targets=[],
            ),
        ]
        xgraph = TestDPUCZDX8G.xgraph_factory.build_from_xlayer(net)
        p_xgraph = partition(xgraph, ["dpuv2-zcu104"])
        dpu_xgraph = TestDPUCZDX8G.target_registry.get_target_build_func(
            "dpuv2-zcu104")(p_xgraph)

        assert len(dpu_xgraph) == 6
        layers = dpu_xgraph.get_layers()

        assert layers[0].type[0] == "Input"

        assert layers[1].type[0] == "Transpose"
        assert layers[1].bottoms == ["in1"]
        assert layers[1].tops == ["xp0"]

        assert layers[2].type[0] == "DPU"
        assert layers[2].bottoms == ["conv1_bottom_NCHW-NHWC"]
        assert layers[2].tops == ["pool1"]
        assert layers[2].shapes == [[1, 2, 2, 2]]
        assert layers[2].attrs["target"] == "dpuv2-zcu104"
        assert layers[2].attrs["input_names"] == ["xinput0"]
        assert layers[2].attrs["output_names"] == ["pool1"]
        assert layers[2].attrs["input_layers"]["xinput0"] == ["conv1"]
        assert layers[2].attrs["output_layers"]["pool1"] == ["pool1"]
        assert layers[2].attrs["__top_tensors"] == {
            "pool1": ["pool1_top_NHWC-NCHW"]
        }
        assert layers[2].attrs["orig_top_tensors"] == {"pool1": ["dense1"]}
        assert layers[2].attrs["__bottom_tensors"] == {
            "xinput0": ["conv1_bottom_NCHW-NHWC"]
        }
        assert layers[2].attrs["orig_bottom_tensors"] == {"xinput0": ["in1"]}

        # Merged TupleGetItem and Transpose layer
        assert layers[3].type[0] == "TupleGetItem"
        assert layers[3].name == "pool1"
        assert layers[3].shapes == [1, 2, 2, 2]
        assert layers[3].bottoms == ["xp0"]
        assert layers[3].tops == ["dense1"]
        assert layers[3].attrs["transpose"] is True

        assert layers[4].type[0] == "Input"
        assert layers[4].name == "in2"
        assert layers[4].tops == ["dense1"]

        assert layers[5].type[0] == "Dense"
        assert layers[5].name == "dense1"
        assert layers[5].shapes == [1, 20]
        assert layers[5].bottoms == ["pool1", "in2"]
        assert layers[5].tops == []
Example #19
0
class TestSubgraphBuildFunc(unittest.TestCase):

    xgraph_partitioner = XGraphPartitioner()
    xgraph_factory = XGraphFactory()
    target_registry = TargetRegistry()

    @classmethod
    def setUpClass(cls):
        def xgraph_build_func_simple(xgraph):
            return subgraph.xgraph_build_func(
                xgraph=xgraph,
                target='test_simple',
                xtype='TEST_SIMPLE',
                layout='NCHW',
            )

        def xgraph_build_func(xgraph):
            return subgraph.xgraph_build_func(xgraph=xgraph,
                                              target='test',
                                              xtype='TEST',
                                              layout='NHWC')

        def xgraph_optimizer(xgraph):
            raise NotImplementedError("")

        def xgraph_quantizer(xgraph):
            raise NotImplementedError("")

        def xgraph_compiler(xgraph):
            raise NotImplementedError("")

        TestSubgraphBuildFunc.target_registry.register_target(
            'test', xgraph_optimizer, xgraph_quantizer, xgraph_compiler,
            xgraph_build_func)
        TestSubgraphBuildFunc.target_registry.register_target(
            'test_simple', xgraph_optimizer, xgraph_quantizer, xgraph_compiler,
            xgraph_build_func_simple)

        @register_op_support_check('test', 'Convolution')
        def conv_op_support(X, bXs, tXs):
            return True

        @register_op_support_check('test', 'Pooling')
        def pooling_op_support(X, bXs, tXs):
            return True

        @register_op_support_check('test', 'Concat')
        def concat_op_support(X, bXs, tXs):
            return True

        @register_op_support_check('test_simple', 'Convolution')
        def conv_op_support(X, bXs, tXs):
            return True

        @register_op_support_check('test_simple', 'Pooling')
        def pooling_op_support(X, bXs, tXs):
            return True

        @register_op_support_check('test_simple', 'Concat')
        def concat_op_support(X, bXs, tXs):
            return True

    @classmethod
    def tearDownClass(cls):

        TestSubgraphBuildFunc.target_registry.unregister_target('test')
        TestSubgraphBuildFunc.target_registry.unregister_target('test_simple')

    def test_basic(self):
        net = [
            XLayer(name='in1',
                   type=['Input'],
                   shapes=[1, 1, 4, 4],
                   sizes=[16],
                   bottoms=[],
                   tops=['conv1'],
                   layer=['in1'],
                   targets=[]),
            XLayer(name='in2',
                   type=['Input'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=[],
                   tops=['add1'],
                   layer=['in2'],
                   targets=[]),
            XLayer(name='conv1',
                   type=['Convolution'],
                   shapes=[1, 2, 3, 3],
                   sizes=[18],
                   bottoms=['in1'],
                   tops=['pool1'],
                   layer=['conv1'],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   attrs={
                       'data_layout': 'NCHW',
                       'padding': [[0, 0], [0, 0], [1, 1], [1, 1]]
                   },
                   targets=[]),
            XLayer(name='pool1',
                   type=['Pooling'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=['conv1'],
                   tops=['add1'],
                   layer=['pool1'],
                   attrs={
                       'data_layout': 'NCHW',
                       'padding': [[0, 0], [0, 0], [1, 1], [1, 1]]
                   },
                   targets=[]),
            XLayer(name='add1',
                   type=['Eltwise'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=['pool1', 'in2'],
                   tops=[],
                   layer=['add1'],
                   targets=[])
        ]
        xgraph = TestSubgraphBuildFunc.xgraph_factory.build_from_xlayer(net)
        p_xgraph = partition(xgraph, ['test_simple'])
        dpu_xgraph = TestSubgraphBuildFunc.target_registry\
            .get_target_build_func('test_simple')(p_xgraph)

        layers = dpu_xgraph.get_layers()
        # print(layers)
        assert len(dpu_xgraph) == 5

        assert layers[0].type[0] == 'Input'
        assert layers[1].type[0] == 'TEST_SIMPLE'
        assert layers[2].type[0] == 'TupleGetItem'
        assert layers[3].type[0] == 'Input'
        assert layers[4].type[0] == 'Eltwise'

        assert layers[0].bottoms == []
        assert layers[0].tops == ['xp0']

        assert layers[1].bottoms == ['in1']
        assert layers[1].tops == ['pool1']
        assert layers[1].attrs['target'] == 'test_simple'
        assert layers[1].attrs['input_names'] == ['xinput0']
        assert layers[1].attrs['output_names'] == ['pool1']
        assert layers[1].attrs['input_layers']['xinput0'] == ['conv1']
        assert layers[1].attrs['output_layers']['pool1'] == ['pool1']
        assert layers[1].attrs['__bottom_tensors'] == {'xinput0': ['in1']}
        assert layers[1].attrs['__top_tensors'] == {'pool1': ['add1']}

        assert layers[2].bottoms == ['xp0']
        assert layers[2].tops == ['add1']

        assert layers[3].bottoms == []
        assert layers[3].tops == ['add1']

        assert layers[4].bottoms == ['pool1', 'in2']
        assert layers[4].tops == []

    def test_two_partitions_interrupt(self):
        net = [
            XLayer(name='in1',
                   type=['Input'],
                   shapes=[1, 1, 4, 4],
                   sizes=[16],
                   bottoms=[],
                   tops=['conv1'],
                   layer=['in1'],
                   targets=[]),
            XLayer(name='conv1',
                   type=['Convolution'],
                   shapes=[1, 2, 3, 3],
                   sizes=[18],
                   bottoms=['in1'],
                   tops=['pool1', 'bn1'],
                   layer=['conv1'],
                   targets=[],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   attrs={
                       'data_layout': 'NCHW',
                       'padding': [[0, 0], [0, 0], [1, 1], [1, 1]]
                   }),
            XLayer(name='pool1',
                   type=['Pooling'],
                   shapes=[1, 4, 3, 3],
                   sizes=[36],
                   bottoms=['conv1'],
                   tops=['concat1'],
                   layer=['pool1'],
                   targets=[],
                   attrs={
                       'data_layout': 'NCHW',
                       'padding': [[0, 0], [0, 0], [1, 1], [1, 1]]
                   }),
            XLayer(name='bn1',
                   type=['BatchNorm'],
                   shapes=[1, 2, 3, 3],
                   sizes=[18],
                   bottoms=['conv1'],
                   tops=['concat1'],
                   layer=['bn1'],
                   data=BatchData(np.array([1, 1]), np.array([0, 0]),
                                  np.array([1, 1]), np.array([0, 0])),
                   targets=[]),
            XLayer(name='concat1',
                   type=['Concat'],
                   shapes=[1, 6, 3, 3],
                   sizes=[54],
                   bottoms=['pool1', 'bn1'],
                   tops=['conv2'],
                   layer=['concat1'],
                   targets=[]),
            XLayer(name='conv2',
                   type=['Convolution'],
                   shapes=[1, 10, 2, 2],
                   sizes=[40],
                   bottoms=['concat1'],
                   tops=[],
                   layer=['conv2'],
                   targets=[],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   attrs={
                       'data_layout': 'NCHW',
                       'padding': [[0, 0], [0, 0], [1, 1], [1, 1]]
                   })
        ]
        xgraph = TestSubgraphBuildFunc.xgraph_factory\
            .build_from_xlayer(net)

        p_xgraph = partition(xgraph, ['test_simple'])

        dpu_xgraph = TestSubgraphBuildFunc.target_registry\
            .get_target_build_func('test_simple')(p_xgraph)

        layers = dpu_xgraph.get_layers()
        assert len(dpu_xgraph) == 7

        assert layers[0].type[0] == 'Input'
        assert layers[0].bottoms == []
        assert layers[0].tops == ['xp0']

        assert layers[1].type[0] == 'TEST_SIMPLE'
        assert layers[1].shapes == [[1, 2, 3, 3], [1, 4, 3, 3]]
        assert layers[1].bottoms == ['in1']
        assert layers[1].tops == ['conv1', 'pool1']
        assert layers[1].attrs['input_names'] == ['xinput0']
        assert set(layers[1].attrs['output_names']) == set(['pool1', 'conv1'])
        assert layers[1].attrs['target'] == 'test_simple'
        assert layers[1].attrs['__bottom_tensors'] == {'xinput0': ['in1']}
        assert layers[1].attrs['orig_bottom_tensors'] == {'xinput0': ['in1']}
        assert layers[1].attrs['__top_tensors'] == \
            {'conv1': ['bn1'], 'pool1': ['concat1']}
        assert layers[1].attrs['orig_top_tensors'] == \
            {'conv1': ['bn1'], 'pool1': ['concat1']}

        assert layers[2].type[0] == 'TupleGetItem'
        assert layers[2].name == 'pool1'
        assert layers[2].bottoms == ['xp0']
        assert layers[2].shapes == [1, 4, 3, 3]
        assert layers[2].tops == ['concat1']
        assert layers[2].attrs['index'] == 1

        assert layers[3].type[0] == 'TupleGetItem'
        assert layers[3].name == 'conv1'
        assert layers[3].bottoms == ['xp0']
        assert layers[3].shapes == [1, 2, 3, 3]
        assert layers[3].tops == ['bn1']
        assert layers[3].attrs['index'] == 0

        assert layers[4].type[0] == 'BatchNorm'
        assert layers[4].name == 'bn1'
        assert layers[4].bottoms == ['conv1']
        assert layers[4].shapes == [1, 2, 3, 3]
        assert layers[4].tops == ['concat1']

        assert layers[5].type[0] == 'Concat'
        assert layers[5].name == 'concat1'
        assert layers[5].bottoms == ['pool1', 'bn1']
        assert layers[5].shapes == [1, 6, 3, 3]
        assert layers[5].tops == ['conv2']

        assert layers[6].type[0] == 'Convolution'
        assert layers[6].name == 'conv2'
        assert layers[6].bottoms == ['concat1']
        assert layers[6].shapes == [1, 10, 2, 2]
        assert layers[6].tops == []

    def test_basic_diff_layout(self):
        net = [
            XLayer(name='in1',
                   type=['Input'],
                   shapes=[1, 1, 4, 4],
                   sizes=[16],
                   bottoms=[],
                   tops=['conv1'],
                   layer=['in1'],
                   targets=[]),
            XLayer(name='in2',
                   type=['Input'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=[],
                   tops=['add1'],
                   layer=['in2'],
                   targets=[]),
            XLayer(name='conv1',
                   type=['Convolution'],
                   shapes=[1, 2, 3, 3],
                   sizes=[18],
                   bottoms=['in1'],
                   tops=['pool1'],
                   layer=['conv1'],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   attrs={
                       'data_layout': 'NCHW',
                       'padding': [[0, 0], [0, 0], [1, 1], [1, 1]]
                   },
                   targets=[]),
            XLayer(name='pool1',
                   type=['Pooling'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=['conv1'],
                   tops=['add1'],
                   layer=['pool1'],
                   attrs={
                       'data_layout': 'NCHW',
                       'padding': [[0, 0], [0, 0], [1, 1], [1, 1]]
                   },
                   targets=[]),
            XLayer(name='add1',
                   type=['Eltwise'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=['pool1', 'in2'],
                   tops=[],
                   layer=['add1'],
                   targets=[])
        ]
        xgraph = TestSubgraphBuildFunc.xgraph_factory.build_from_xlayer(net)
        p_xgraph = partition(xgraph, ['test'])
        dpu_xgraph = TestSubgraphBuildFunc.target_registry\
            .get_target_build_func('test')(p_xgraph)

        layers = dpu_xgraph.get_layers()
        # print(layers)
        assert (len(dpu_xgraph) == 6)

        assert (layers[0].type[0] == 'Input')
        assert (layers[0].name == 'in1')
        assert (layers[0].bottoms == [])
        assert (layers[0].tops == ['conv1_bottom_NCHW-NHWC'])

        assert (layers[1].type[0] == 'Transpose')
        assert (layers[1].name == 'conv1_bottom_NCHW-NHWC')
        assert (layers[1].bottoms == ['in1'])
        assert (layers[1].tops == ['xp0'])

        assert (layers[2].type[0] == 'TEST')
        assert (layers[2].bottoms == ['conv1_bottom_NCHW-NHWC'])
        assert (layers[2].tops == ['pool1'])
        assert (layers[2].attrs['target'] == 'test')
        assert (layers[2].attrs['input_names'] == ['xinput0'])
        assert (layers[2].attrs['output_names'] == ['pool1'])
        assert (layers[2].attrs['input_layers']['xinput0'] == ['conv1'])
        assert (layers[2].attrs['output_layers']['pool1'] == ['pool1'])
        assert (layers[2].attrs['__bottom_tensors'] == {
            'xinput0': ['conv1_bottom_NCHW-NHWC']
        })
        assert (layers[2].attrs['orig_bottom_tensors'] == {'xinput0': ['in1']})
        assert (layers[2].attrs['__top_tensors'] == {
            'pool1': ['pool1_top_NHWC-NCHW']
        })
        assert (layers[2].attrs['orig_top_tensors'] == {'pool1': ['add1']})

        assert (layers[3].type[0] == 'TupleGetItem')
        assert (layers[3].bottoms == ['xp0'])
        assert (layers[3].tops == ['add1'])
        assert layers[3].attrs['transpose'] is True
        assert layers[3].attrs['axes'] == [0, 3, 1, 2]

        # assert(layers[4].type[0] == 'Transpose')
        # assert(layers[4].name == 'pool1_top_NHWC-NCHW')
        # assert(layers[4].bottoms == ['pool1'])
        # assert(layers[4].tops == ['add1'])
        # assert layers[4].attrs['axes'] == [0, 3, 1, 2]

        assert layers[4].type[0] == 'Input'
        assert layers[4].name == 'in2'
        assert layers[4].bottoms == []
        assert layers[4].tops == ['add1']

        assert layers[5].type[0] == 'Eltwise'
        assert layers[5].name == 'add1'
        assert layers[5].bottoms == ['pool1', 'in2']
        assert layers[5].tops == []

    def test_two_partition_inputs(self):
        net = [
            XLayer(name='in1',
                   type=['Input'],
                   shapes=[1, 1, 4, 4],
                   sizes=[16],
                   bottoms=[],
                   tops=['conv1'],
                   layer=['in1'],
                   targets=[]),
            XLayer(name='in2',
                   type=['Input'],
                   shapes=[1, 1, 4, 4],
                   sizes=[16],
                   bottoms=[],
                   tops=['conv2'],
                   layer=['in2'],
                   targets=[]),
            XLayer(name='conv1',
                   type=['Convolution'],
                   shapes=[1, 2, 3, 3],
                   sizes=[18],
                   bottoms=['in1'],
                   tops=['pool1'],
                   layer=['conv1'],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   attrs={
                       'data_layout': 'NCHW',
                       'padding': [[0, 0], [0, 0], [1, 1], [1, 1]]
                   },
                   targets=[]),
            XLayer(name='pool1',
                   type=['Pooling'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=['conv1'],
                   tops=['concat1'],
                   layer=['pool1'],
                   attrs={
                       'data_layout': 'NCHW',
                       'padding': [[0, 0], [0, 0], [1, 1], [1, 1]]
                   },
                   targets=[]),
            XLayer(name='conv2',
                   type=['Convolution'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=['in2'],
                   tops=['concat1'],
                   layer=['conv2'],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   attrs={
                       'data_layout': 'NCHW',
                       'padding': [[0, 0], [0, 0], [1, 1], [1, 1]]
                   },
                   targets=[]),
            XLayer(name='concat1',
                   type=['Concat'],
                   shapes=[1, 4, 2, 2],
                   sizes=[16],
                   bottoms=['pool1', 'conv2'],
                   tops=['dense1'],
                   layer=['concat1'],
                   attrs={'axis': 1},
                   targets=[]),
            XLayer(name='dense1',
                   type=['Dense'],
                   shapes=[1, 20],
                   sizes=[],
                   bottoms=['concat1'],
                   tops=[],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   layer=['dense1'],
                   targets=[])
        ]
        xgraph = TestSubgraphBuildFunc.xgraph_factory.build_from_xlayer(net)
        p_xgraph = partition(xgraph, ['test'])
        dpu_xgraph = TestSubgraphBuildFunc.target_registry\
            .get_target_build_func('test')(p_xgraph)

        layers = dpu_xgraph.get_layers()
        assert len(dpu_xgraph) == 7

        assert layers[0].type[0] == 'Input'
        assert layers[0].name == 'in1'
        assert layers[0].bottoms == []
        assert layers[0].tops == ['conv1_bottom_NCHW-NHWC']
        assert layers[0].target == 'cpu'
        assert layers[0].subgraph is None

        assert layers[1].type[0] == 'Transpose'
        assert layers[1].name == 'conv1_bottom_NCHW-NHWC'
        assert layers[1].bottoms == ['in1']
        assert layers[1].tops == ['xp2']
        assert layers[1].target == 'cpu'
        assert layers[1].subgraph is None

        assert layers[2].type[0] == 'Input'
        assert layers[2].name == 'in2'
        assert layers[2].bottoms == []
        assert layers[2].tops == ['conv2_bottom_NCHW-NHWC']
        assert layers[2].target == 'cpu'
        assert layers[2].subgraph is None

        assert layers[3].type[0] == 'Transpose'
        assert layers[3].name == 'conv2_bottom_NCHW-NHWC'
        assert layers[3].bottoms == ['in2']
        assert layers[3].tops == ['xp2']
        assert layers[3].target == 'cpu'
        assert layers[3].subgraph is None

        assert layers[4].type[0] == 'TEST'
        assert layers[4].name == 'xp2'
        assert layers[4].bottoms == [
            'conv1_bottom_NCHW-NHWC', 'conv2_bottom_NCHW-NHWC'
        ]
        assert layers[4].tops == ['concat1']
        assert layers[4].attrs['target'] == 'test'
        assert layers[4].attrs['input_names'] == ['xinput0', 'xinput1']
        assert layers[4].attrs['output_names'] == ['concat1']
        assert layers[4].attrs['input_layers']['xinput0'] == ['conv1']
        assert layers[4].attrs['input_layers']['xinput1'] == ['conv2']
        assert layers[4].attrs['output_layers']['concat1'] == ['concat1']
        assert (layers[4].attrs['__bottom_tensors'] == {
            'xinput0': ['conv1_bottom_NCHW-NHWC'],
            'xinput1': ['conv2_bottom_NCHW-NHWC']
        })
        assert (layers[4].attrs['orig_bottom_tensors'] == {
            'xinput0': ['in1'],
            'xinput1': ['in2']
        })
        assert layers[4].attrs['__top_tensors'] == {
            'concat1': ['merge_pool1_top_NHWC-NCHW_conv2_top_NHWC-NCHW']
        }
        assert layers[4].attrs['orig_top_tensors'] == {'concat1': ['dense1']}
        assert layers[4].target == 'cpu'
        assert layers[4].subgraph is None

        assert layers[5].type[0] == 'TupleGetItem'
        assert layers[5].name == 'concat1'
        assert layers[5].bottoms == ['xp2']
        assert layers[5].tops == ['dense1']
        assert layers[5].target == 'cpu'
        assert layers[5].subgraph is None
        assert layers[5].attrs['transpose'] is True

        # assert layers[6].type[0] == 'Transpose'
        # assert layers[6].name ==\
        #     'merge_pool1_top_NHWC-NCHW_conv2_top_NHWC-NCHW'
        # assert layers[6].bottoms == ['concat1']
        # assert layers[6].tops == ['dense1']
        # assert layers[6].target == 'cpu'
        # assert layers[6].subgraph is None

        assert layers[6].type[0] == 'Dense'
        assert layers[6].name == 'dense1'
        assert layers[6].bottoms == ['concat1']
        assert layers[6].tops == []
        assert layers[6].target == 'cpu'
        assert layers[6].subgraph is None

    def test_two_partition_diff_layout(self):
        net = [
            XLayer(name='in1',
                   type=['Input'],
                   shapes=[1, 1, 4, 4],
                   sizes=[16],
                   bottoms=[],
                   tops=['conv1'],
                   layer=['in1'],
                   targets=[]),
            XLayer(name='in2',
                   type=['Input'],
                   shapes=[1, 4, 4, 1],
                   sizes=[16],
                   bottoms=[],
                   tops=['in2_transpose'],
                   layer=['in2'],
                   targets=[]),
            XLayer(name='in2_transpose',
                   type=['Transpose'],
                   shapes=[1, 1, 4, 4],
                   sizes=[16],
                   bottoms=['in2'],
                   tops=['conv2'],
                   layer=['in2'],
                   attrs={'axes': [0, 3, 1, 2]},
                   targets=[]),
            XLayer(name='conv1',
                   type=['Convolution'],
                   shapes=[1, 2, 3, 3],
                   sizes=[18],
                   bottoms=['in1'],
                   tops=['pool1'],
                   layer=['conv1'],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   attrs={
                       'data_layout': 'NCHW',
                       'padding': [[0, 0], [0, 0], [1, 1], [1, 1]]
                   },
                   targets=[]),
            XLayer(name='pool1',
                   type=['Pooling'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=['conv1'],
                   tops=['concat1'],
                   layer=['pool1'],
                   attrs={
                       'data_layout': 'NCHW',
                       'padding': [[0, 0], [0, 0], [1, 1], [1, 1]]
                   },
                   targets=[]),
            XLayer(name='conv2',
                   type=['Convolution'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=['in2_transpose'],
                   tops=['concat1'],
                   layer=['conv2'],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   attrs={
                       'data_layout': 'NCHW',
                       'padding': [[0, 0], [0, 0], [1, 1], [1, 1]]
                   },
                   targets=[]),
            XLayer(name='concat1',
                   type=['Concat'],
                   shapes=[1, 4, 2, 2],
                   sizes=[16],
                   bottoms=['pool1', 'conv2'],
                   tops=['concat1_transpose'],
                   layer=['concat1'],
                   attrs={'axis': 1},
                   targets=[]),
            XLayer(name='concat1_transpose',
                   type=['Transpose'],
                   shapes=[1, 2, 2, 4],
                   sizes=[16],
                   bottoms=['concat1'],
                   tops=['dense1'],
                   layer=['concat1'],
                   attrs={'axes': [0, 2, 3, 1]},
                   targets=[]),
            XLayer(name='dense1',
                   type=['Dense'],
                   shapes=[1, 20],
                   sizes=[],
                   bottoms=['concat1_transpose'],
                   tops=[],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   layer=['dense1'],
                   targets=[])
        ]
        xgraph = TestSubgraphBuildFunc.xgraph_factory.build_from_xlayer(net)
        p_xgraph = partition(xgraph, ['test'])
        p_xlayers = p_xgraph.get_layers()

        dpu_xgraph = TestSubgraphBuildFunc.target_registry\
            .get_target_build_func('test')(p_xgraph)

        layers = dpu_xgraph.get_layers()
        assert len(dpu_xgraph) == 6

        assert layers[0].type[0] == 'Input'
        assert layers[0].name == 'in1'
        assert layers[0].shapes == [1, 1, 4, 4]
        assert layers[0].bottoms == []
        assert layers[0].tops == ['conv1_bottom_NCHW-NHWC']
        assert layers[0].target == 'cpu'
        assert layers[0].subgraph is None

        assert layers[1].type[0] == 'Transpose'
        assert layers[1].name == 'conv1_bottom_NCHW-NHWC'
        assert layers[1].shapes == [1, 4, 4, 1]
        assert layers[1].bottoms == ['in1']
        assert layers[1].tops == ['xp2']
        assert layers[1].target == 'cpu'
        assert layers[1].subgraph is None

        assert layers[2].type[0] == 'Input'
        assert layers[2].name == 'in2'
        assert layers[2].shapes == [1, 4, 4, 1]
        assert layers[2].bottoms == []
        assert layers[2].tops == ['xp2']
        assert layers[2].target == 'cpu'
        assert layers[2].subgraph is None

        assert layers[3].type[0] == 'TEST'
        assert layers[3].name == 'xp2'
        assert layers[3].shapes == [[1, 2, 2, 4]]
        assert layers[3].bottoms == ['conv1_bottom_NCHW-NHWC', 'in2']
        assert layers[3].tops == ['concat1']
        assert layers[3].target == 'cpu'
        assert layers[3].subgraph is None
        assert layers[3].attrs['target'] == 'test'
        assert layers[3].attrs['input_names'] == ['xinput0', 'xinput1']
        assert layers[3].attrs['output_names'] == ['concat1']
        assert layers[3].attrs['input_layers']['xinput0'] == ['conv1']
        assert layers[3].attrs['input_layers']['xinput1'] == ['conv2']
        assert layers[3].attrs['output_layers']['concat1'] == ['concat1']
        assert (layers[3].attrs['__bottom_tensors'] == {
            'xinput0': ['conv1_bottom_NCHW-NHWC'],
            'xinput1': ['in2']
        })
        assert (layers[3].attrs['orig_bottom_tensors'] == {
            'xinput0': ['in1'],
            'xinput1': ['in2']
        })
        assert layers[3].attrs['__top_tensors'] == {'concat1': ['dense1']}
        assert layers[3].attrs['orig_top_tensors'] == {'concat1': ['dense1']}

        assert layers[4].type[0] == 'TupleGetItem'
        assert layers[4].name == 'concat1'
        assert layers[4].shapes == [1, 2, 2, 4]
        assert layers[4].bottoms == ['xp2']
        assert layers[4].tops == ['dense1']
        assert layers[4].target == 'cpu'
        assert layers[4].subgraph is None

        assert layers[5].type[0] == 'Dense'
        assert layers[5].name == 'dense1'
        assert layers[5].shapes == [1, 20]
        assert layers[5].bottoms == ['concat1']
        assert layers[5].tops == []
        assert layers[5].target == 'cpu'
        assert layers[5].subgraph is None

    def test_two_partition_inputs_complex(self):
        net = [
            XLayer(name='in1',
                   type=['Input'],
                   shapes=[1, 1, 4, 4],
                   sizes=[16],
                   bottoms=[],
                   tops=['conv1'],
                   layer=['in1'],
                   targets=[]),
            XLayer(name='conv1',
                   type=['Convolution'],
                   shapes=[1, 2, 3, 3],
                   sizes=[18],
                   bottoms=['in1'],
                   tops=['pool1'],
                   layer=['conv1'],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   attrs={
                       'data_layout': 'NCHW',
                       'padding': [[0, 0], [0, 0], [1, 1], [1, 1]]
                   },
                   targets=[]),
            XLayer(name='pool1',
                   type=['Pooling'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=['conv1'],
                   tops=['concat1'],
                   layer=['pool1'],
                   attrs={
                       'data_layout': 'NCHW',
                       'padding': [[0, 0], [0, 0], [1, 1], [1, 1]]
                   },
                   targets=[]),
            XLayer(name='in2',
                   type=['Input'],
                   shapes=[1, 1, 4, 4],
                   sizes=[16],
                   bottoms=[],
                   tops=['conv2'],
                   layer=['in2'],
                   targets=[]),
            XLayer(name='conv2',
                   type=['Convolution'],
                   shapes=[1, 2, 2, 2],
                   sizes=[8],
                   bottoms=['in2'],
                   tops=['concat1'],
                   layer=['conv2'],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   attrs={
                       'data_layout': 'NCHW',
                       'padding': [[0, 0], [0, 0], [1, 1], [1, 1]]
                   },
                   targets=[]),
            XLayer(name='concat1',
                   type=['Concat'],
                   shapes=[1, 4, 2, 2],
                   sizes=[16],
                   bottoms=['pool1', 'conv2'],
                   tops=['dense1'],
                   layer=['concat1'],
                   attrs={'axis': 1},
                   targets=[]),
            XLayer(name='dense1',
                   type=['Dense'],
                   shapes=[1, 20],
                   sizes=[],
                   bottoms=['concat1'],
                   tops=[],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   layer=['dense1'],
                   targets=[])
        ]
        xgraph = TestSubgraphBuildFunc.xgraph_factory.build_from_xlayer(net)
        p_xgraph = partition(xgraph, ['test'])
        dpu_xgraph = TestSubgraphBuildFunc.target_registry\
            .get_target_build_func('test')(p_xgraph)

        layers = dpu_xgraph.get_layers()
        assert len(dpu_xgraph) == 7

        assert layers[0].type[0] == 'Input'
        assert layers[0].name == 'in1'
        assert layers[0].shapes == [1, 1, 4, 4]
        assert layers[0].bottoms == []
        assert layers[0].tops == ['conv1_bottom_NCHW-NHWC']
        assert layers[0].target == 'cpu'
        assert layers[0].subgraph is None

        assert layers[1].type[0] == 'Transpose'
        assert layers[1].name == 'conv1_bottom_NCHW-NHWC'
        assert layers[1].shapes == [1, 4, 4, 1]
        assert layers[1].bottoms == ['in1']
        assert layers[1].tops == ['xp2']
        assert layers[1].target == 'cpu'
        assert layers[1].subgraph is None

        assert layers[2].type[0] == 'Input'
        assert layers[2].name == 'in2'
        assert layers[2].shapes == [1, 1, 4, 4]
        assert layers[2].bottoms == []
        assert layers[2].tops == ['conv2_bottom_NCHW-NHWC']
        assert layers[2].target == 'cpu'
        assert layers[2].subgraph is None

        assert layers[3].type[0] == 'Transpose'
        assert layers[3].name == 'conv2_bottom_NCHW-NHWC'
        assert layers[3].shapes == [1, 4, 4, 1]
        assert layers[3].bottoms == ['in2']
        assert layers[3].tops == ['xp2']
        assert layers[3].target == 'cpu'
        assert layers[3].subgraph is None

        assert layers[4].type[0] == 'TEST'
        assert layers[4].name == 'xp2'
        assert layers[4].shapes == [[1, 2, 2, 4]]
        assert layers[4].bottoms == [
            'conv1_bottom_NCHW-NHWC', 'conv2_bottom_NCHW-NHWC'
        ]
        assert layers[4].tops == ['concat1']
        assert layers[4].target == 'cpu'
        assert layers[4].subgraph is None
        assert layers[4].tops == ['concat1']
        assert layers[4].attrs['target'] == 'test'
        assert layers[4].attrs['input_names'] == ['xinput0', 'xinput1']
        assert layers[4].attrs['output_names'] == ['concat1']
        assert layers[4].attrs['input_layers']['xinput0'] == ['conv1']
        assert layers[4].attrs['input_layers']['xinput1'] == ['conv2']
        assert layers[4].attrs['output_layers']['concat1'] == ['concat1']
        assert (layers[4].attrs['__bottom_tensors'] == {
            'xinput0': ['conv1_bottom_NCHW-NHWC'],
            'xinput1': ['conv2_bottom_NCHW-NHWC']
        })
        assert (layers[4].attrs['orig_bottom_tensors'] == {
            'xinput0': ['in1'],
            'xinput1': ['in2']
        })
        assert layers[4].attrs['__top_tensors'] ==\
            {'concat1':
                ['merge_pool1_top_NHWC-NCHW_conv2_top_NHWC-NCHW']}
        assert layers[4].attrs['orig_top_tensors'] ==\
            {'concat1': ['dense1']}

        assert layers[5].type[0] == 'TupleGetItem'
        assert layers[5].name == 'concat1'
        assert layers[5].shapes == [1, 4, 2, 2]
        assert layers[5].bottoms == ['xp2']
        assert layers[5].tops == ['dense1']
        assert layers[5].attrs['transpose'] is True
        assert layers[5].attrs['axes'] == [0, 3, 1, 2]

        # assert layers[6].type[0] == 'Transpose'
        # assert layers[6].name ==\
        #     'merge_pool1_top_NHWC-NCHW_conv2_top_NHWC-NCHW'
        # assert layers[6].shapes == [1, 4, 2, 2]
        # assert layers[6].bottoms == ['concat1']
        # assert layers[6].tops == ['dense1']

        assert layers[6].type[0] == 'Dense'
        assert layers[6].name == 'dense1'
        assert layers[6].shapes == [1, 20]
        assert layers[6].bottoms == ['concat1']
        assert layers[6].tops == []

    def test_inception_like_block(self):
        net = [
            XLayer(name='in1',
                   type=['Input'],
                   shapes=[1, 1, 4, 4],
                   sizes=[16],
                   bottoms=[],
                   tops=['concat1'],
                   layer=['in1'],
                   targets=[]),
            XLayer(name='in2',
                   type=['Input'],
                   shapes=[1, 1, 4, 4],
                   sizes=[16],
                   bottoms=[],
                   tops=['concat1'],
                   layer=['in2'],
                   targets=[]),
            XLayer(name='concat1',
                   type=['Concat'],
                   shapes=[1, 2, 4, 4],
                   sizes=[32],
                   bottoms=['in1', 'in2'],
                   tops=['conv1', 'conv2'],
                   layer=['concat1'],
                   attrs={'axis': 1},
                   targets=[]),
            XLayer(name='conv1',
                   type=['Convolution'],
                   shapes=[1, 4, 3, 3],
                   sizes=[],
                   bottoms=['concat1'],
                   tops=['pool1'],
                   layer=['conv1'],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   attrs={
                       'data_layout': 'NCHW',
                       'padding': [[0, 0], [0, 0], [1, 1], [1, 1]]
                   },
                   targets=[]),
            XLayer(name='pool1',
                   type=['Pooling'],
                   shapes=[1, 4, 2, 2],
                   sizes=[],
                   bottoms=['conv1'],
                   tops=['concat2'],
                   layer=['pool1'],
                   attrs={
                       'data_layout': 'NCHW',
                       'padding': [[0, 0], [0, 0], [1, 1], [1, 1]]
                   },
                   targets=[]),
            XLayer(name='conv2',
                   type=['Convolution'],
                   shapes=[1, 4, 2, 2],
                   sizes=[],
                   bottoms=['concat1'],
                   tops=['concat2'],
                   layer=['conv2'],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   attrs={
                       'data_layout': 'NCHW',
                       'padding': [[0, 0], [0, 0], [1, 1], [1, 1]]
                   },
                   targets=[]),
            XLayer(name='concat2',
                   type=['Concat'],
                   shapes=[1, 8, 2, 2],
                   sizes=[32],
                   bottoms=['pool1', 'conv2'],
                   tops=['dense1'],
                   layer=['concat2'],
                   attrs={'axis': 1},
                   targets=[]),
            XLayer(name='dense1',
                   type=['Dense'],
                   shapes=[1, 20],
                   sizes=[20],
                   bottoms=['concat2'],
                   tops=[],
                   layer=['dense1'],
                   data=ConvData(np.array([1, 1]), np.array([0, 0])),
                   targets=[])
        ]
        xgraph = TestSubgraphBuildFunc.xgraph_factory.build_from_xlayer(net)
        p_xgraph = partition(xgraph, ['test'])
        dpu_xgraph = TestSubgraphBuildFunc.target_registry\
            .get_target_build_func('test')(p_xgraph)

        layers = dpu_xgraph.get_layers()
        assert len(dpu_xgraph) == 7

        assert layers[0].type[0] == 'Input'
        assert layers[0].name == 'in1'
        assert layers[0].shapes == [1, 1, 4, 4]
        assert layers[0].bottoms == []
        assert layers[0].tops ==\
            ['0_split_conv1_bottom_NCHW-NHWC_conv2_bottom_NCHW-NHWC']
        assert layers[0].target == 'cpu'
        assert layers[0].subgraph is None

        assert layers[1].type[0] == 'Transpose'
        assert layers[1].name ==\
            '0_split_conv1_bottom_NCHW-NHWC_conv2_bottom_NCHW-NHWC'
        assert layers[1].shapes == [1, 4, 4, 1]
        assert layers[1].bottoms == ['in1']
        assert layers[1].tops == ['xp0']
        assert layers[1].target == 'cpu'
        assert layers[1].subgraph is None

        assert layers[2].type[0] == 'Input'
        assert layers[2].name == 'in2'
        assert layers[2].shapes == [1, 1, 4, 4]
        assert layers[2].bottoms == []
        assert layers[2].tops ==\
            ['1_split_conv1_bottom_NCHW-NHWC_conv2_bottom_NCHW-NHWC']
        assert layers[2].target == 'cpu'
        assert layers[2].subgraph is None

        assert layers[3].type[0] == 'Transpose'
        assert layers[3].name ==\
            '1_split_conv1_bottom_NCHW-NHWC_conv2_bottom_NCHW-NHWC'
        assert layers[3].shapes == [1, 4, 4, 1]
        assert layers[3].bottoms == ['in2']
        assert layers[3].tops == ['xp0']
        assert layers[3].target == 'cpu'
        assert layers[3].subgraph is None

        assert layers[4].type[0] == 'TEST'
        assert layers[4].name == 'xp0'
        assert layers[4].shapes == [[1, 2, 2, 8]]
        assert layers[4].bottoms ==\
            ['0_split_conv1_bottom_NCHW-NHWC_conv2_bottom_NCHW-NHWC',
             '1_split_conv1_bottom_NCHW-NHWC_conv2_bottom_NCHW-NHWC']
        assert layers[4].tops == ['concat2']
        assert layers[4].target == 'cpu'
        assert layers[4].subgraph is None
        assert layers[4].tops == ['concat2']
        assert layers[4].attrs['target'] == 'test'
        assert layers[4].attrs['input_names'] == ['xinput0', 'xinput1']
        assert layers[4].attrs['output_names'] == ['concat2']
        assert layers[4].attrs['input_layers']['xinput0'] == ['concat1']
        assert layers[4].attrs['input_layers']['xinput1'] == ['concat1']
        assert layers[4].attrs['output_layers']['concat2'] == ['concat2']
        assert (layers[4].attrs['__bottom_tensors'] == {
            'xinput0':
            ['0_split_conv1_bottom_NCHW-NHWC_conv2_bottom'
             '_NCHW-NHWC'],
            'xinput1':
            ['1_split_conv1_bottom_NCHW-NHWC_conv2_bottom'
             '_NCHW-NHWC']
        })
        assert (layers[4].attrs['orig_bottom_tensors'] == {
            'xinput0': ['in1'],
            'xinput1': ['in2']
        })
        assert layers[4].attrs['__top_tensors'] ==\
            {'concat2':
                ['merge_pool1_top_NHWC-NCHW_conv2_top_NHWC-NCHW']}
        assert layers[4].attrs['orig_top_tensors'] ==\
            {'concat2': ['dense1']}

        assert layers[5].type[0] == 'TupleGetItem'
        assert layers[5].name == 'concat2'
        assert layers[5].shapes == [1, 8, 2, 2]
        assert layers[5].bottoms == ['xp0']
        assert layers[5].tops == ['dense1']

        # assert layers[6].type[0] == 'Transpose'
        # assert layers[6].name ==\
        #     'merge_pool1_top_NHWC-NCHW_conv2_top_NHWC-NCHW'
        # assert layers[6].shapes == [1, 8, 2, 2]
        # assert layers[6].bottoms == ['concat2']
        # assert layers[6].tops == ['dense1']

        assert layers[6].type[0] == 'Dense'
        assert layers[6].name == 'dense1'
        assert layers[6].shapes == [1, 20]
        assert layers[6].bottoms == ['concat2']
        assert layers[6].tops == []
Example #20
0
def xgraph_build_func(xgraph: XGraph,
                      target: str,
                      xtype,
                      layout='NCHW',
                      **kwargs) -> XGraph:

    fancy_logger.banner("Subgraph build func, target: {}, layout: {}".format(
        target, layout))

    compiler_output = xgraph.get_compiler_output() if xgraph.is_compiled() \
        else None
    compiler_output_keys = list(compiler_output.keys()) \
        if compiler_output else []
    logger.debug("Compiler output keys: {}".format(compiler_output_keys))

    if layout not in ['NCHW', 'NHWC']:
        raise ValueError(
            "Supported layouts are [NCHW, NHWC] but got: {}".format(layout))

    layout_transform_pass = \
        XGraphLayoutTransformationPass(layout, target=target)
    xgraph = layout_transform_pass.execute(xgraph, subgraphs_only=False)

    xgraph_factory = XGraphFactory()
    xgraph_partitioner = XGraphPartitioner()

    subgraphs = {
        xp.name: xp
        for xp in xgraph_partitioner.get_subgraphs(xgraph)
    }

    # Retrieve CompilerOutput if available
    # compiler_output = xgraph.get_compiler_output() if xgraph.is_compiled() \
    #     else None
    # compiler_output_keys = list(compiler_output.keys()) \
    #     if compiler_output else []
    # logger.debug("Compiler output keys: {}".format(compiler_output_keys))
    # Keep track of the visited partitions/subgraphs and the layers
    #   inside the partition
    visited_xps = {}

    # Keep track of the subgraph output tensors and the corresponding
    #   new layers (TupleGetItem or Transpose)
    xp_out_tensors_2_layers = {}

    name_changes = {}
    net_map = {}
    net = []
    for X in xgraph.get_layers():

        if X.subgraph is not None and X.subgraph not in visited_xps:

            Xp = subgraphs[X.subgraph]

            if 'target' in Xp.attrs and Xp.attrs['target'] == target:

                visited_xps[Xp.name] = set([X.name])

                logger.debug("XSHAPES: {}".format(X.shapes))

                bottoms = Xp.bottoms

                # Keep track of subgraph input and output names
                sub_xgraph = xgraph_factory.build_from_xlayer(Xp.subgraph_data)

                input_names = Xp.attrs['input_names'][:]
                output_names = Xp.attrs['output_names'][:]
                input_layers = \
                    [sub_xgraph.get(in_name) for in_name in input_names]
                output_layers = \
                    [sub_xgraph.get(out_name) for out_name in output_names]

                attrs = {
                    'input_names': input_names,
                    'output_names': output_names,
                    'input_layers':
                    {il.name: il.layer[:]
                     for il in input_layers},
                    'output_layers':
                    {ol.name: ol.layer[:]
                     for ol in output_layers}
                }
                for k, v in kwargs.items():
                    if k in attrs:
                        raise ValueError("Provided claimed subgraph layer"
                                         " key: {}".format(k))
                    attrs[k] = v

                if Xp.name in compiler_output_keys:
                    attrs['rt_in_map'] = compiler_output.get_in_map(Xp.name)
                    for in_name in input_names:
                        for merged_layer in attrs['input_layers'][in_name]:
                            attrs['rt_in_map'][merged_layer] = \
                                attrs['rt_in_map'][in_name]
                    attrs['rt_out_map'] = compiler_output.get_out_map(Xp.name)
                    for out_name in output_names:
                        for merged_layer in attrs['output_layers'][out_name]:
                            attrs['rt_out_map'][merged_layer] = \
                                attrs['rt_out_map'][out_name]

                Xp.attrs.update(attrs)

                shapes = Xp.shapes[:]

                subgraph_X = Xp._replace(
                    # name = X.name,
                    type=[xtype],
                    shapes=shapes,
                    bottoms=bottoms,
                    # Fill tops later
                    tops=[],
                    subgraph_data=[])
                net.append(subgraph_X.name)
                net_map[Xp.name] = subgraph_X

                # Subgraph layers have multiple outputs (Tuple) so we
                #   retrieve the different subgraph outputs
                #   (see output_names variable) using a TupleGetItem
                #   layer
                top_tensors = Xp.attrs['__top_tensors']

                for i, output_name in enumerate(output_names):
                    # Handle merged layers
                    out_tensor = Xp.attrs['output_layers'][output_name][-1]
                    tgi_name = out_tensor
                    # tgi_name = subgraph_X.name + '_tgi' + str(i)

                    top_tensor = top_tensors[output_name]

                    shapes = subgraph_X.shapes[i][:]
                    X_tgi = defaultXLayer()
                    X_tgi = X_tgi._replace(name=tgi_name,
                                           type=['TupleGetItem'],
                                           shapes=shapes,
                                           sizes=shapes.get_size(),
                                           layer=[tgi_name],
                                           tops=top_tensor[:],
                                           bottoms=[subgraph_X.name],
                                           internal=1,
                                           attrs={'index': i})
                    net.append(X_tgi.name)
                    # Keep track of TGI layer for both last merged layer and output name
                    net_map[tgi_name] = X_tgi
                    net_map[output_name] = X_tgi

                    subgraph_X.tops.append(tgi_name)

                    xp_out_tensors_2_layers[output_name] = tgi_name

            else:
                net.append(X.name)
                net_map[X.name] = X

        elif X.subgraph is not None and X.subgraph in visited_xps:
            # Remove layer
            visited_xps[X.subgraph].add(X.name)
        elif 'Transpose' in X.type:
            # Possibly merge transpose in TupleGetItem layer
            bX = net_map[X.bottoms[0]]
            new_tops = []
            for t in bX.tops:
                if t != X.name:
                    new_tops.append(t)
                elif len(X.tops) > 0:
                    new_tops.append(X.tops[0])
            if 'TupleGetItem' in bX.type:
                new_X = bX._replace(tops=new_tops)
                new_X.attrs['transpose'] = True
                new_X.attrs['axes'] = X.attrs['axes']
                new_X.shapes[:] = TensorShape(X.shapes[:])
                net_map[new_X.name] = new_X
                name_changes[X.name] = bX.name
            else:
                net.append(X.name)
                net_map[X.name] = X
        else:
            net.append(X.name)
            net_map[X.name] = X

        # Reflect possibly merged layers
        new_bottoms = [
            b if b not in name_changes else name_changes[b] for b in X.bottoms
        ]
        if new_bottoms != X.bottoms:
            new_X = X._replace(bottoms=new_bottoms)
            net_map[X.name] = new_X

    # Set tops and bottoms  & enforce topological sequence
    for xp in visited_xps.keys():
        Xp = subgraphs[xp]

        for b in Xp.bottoms:
            top_name = Xp.name
            bX = xgraph.get(b)
            bX.tops = [(bXt if bXt not in visited_xps[Xp.name] else top_name)
                       for bXt in bX.tops]

        for t in Xp.tops:
            tX = xgraph.get(t)
            tX.bottoms = [(tXb if tXb not in visited_xps[Xp.name] else
                           xp_out_tensors_2_layers[tXb]) for tXb in tX.bottoms]

    # Topological sorting
    X_net = [net_map[e] for e in net]
    top_net = sort_topologically(X_net)

    sub_xgraph = xgraph_factory.build_from_xlayer(top_net)

    # Merge transposes if they are cancelling out
    # optimizer = XGraphTransposesOptimizer(sub_xgraph)
    # optimizer.optimize()

    return sub_xgraph
Example #21
0
class DECENTQuantizer(XGraphBaseSubgraphQuantizer):

    # try:
    #     if hasattr(tf.contrib, 'decent_q'):
    #         from tensorflow.contrib import decent_q
    # except Exception as e:
    #     warnings.warn("Could not import decent_q module")
    try:
        #     from tensorflow.contrib import decent_q
        import tensorflow as tf
        if hasattr(tf, 'contrib') and hasattr(tf.contrib, 'decent_q'):
            from tensorflow.contrib import decent_q
        else:
            warnings.warn("Could not import decent_q module. Please check"
                          " if installed.")
    except ImportError:
        warnings.warn("Could not import decent_q module. Please check"
                      " if installed.")

    xgraph_factory = XGraphFactory()
    xgraph_partitioner = XGraphPartitioner()

    def __init__(self,
                 xgraph,
                 inputs_func,
                 work_dir=os.path.join(os.getcwd(), 'work'),
                 quant_iter=1,
                 **kwargs):

        super(DECENTQuantizer, self).__init__(xgraph, inputs_func, work_dir)

        self.quant_iter = quant_iter
        self.gen = TfGenerator()
        self.partition_graphs = {}
        self.res = {}
        self.kwargs = kwargs

        self.q_output = QuantizerOutput(name=xgraph.get_name())

    def quantize_subgraph(self, xgraph, inputs, input_names, output_names):
        # type: (XGraph, Dict[str, numpy.ndarray])
        """ Quantize subgraph with inputs """

        # Import Tensorflow only when needed to avoid strict dependency
        import tensorflow as tf

        frozen_graph = self.partition_graphs[xgraph.get_name()]
        logger.info("Load frozen graph from: {}".format(frozen_graph))
        input_graph_def = tf.compat.v1.GraphDef()
        with tf.io.gfile.GFile(frozen_graph, "rb") as f:
            input_graph_def.ParseFromString(f.read())

        logger.info("Quantization input: {} and output names: {}".format(
            input_names, output_names))
        input_shapes = [X.shapes.tolist() for X in xgraph.get_input_layers()]

        def inputs_func(iter):
            import numpy as np
            nonlocal inputs

            return inputs

        logger.info("START decent quantization for graph partition: {}".format(
            xgraph.get_name()))
        q_config = self.decent_q.QuantizeConfig(input_nodes=input_names,
                                                output_nodes=output_names,
                                                input_shapes=input_shapes,
                                                output_dir=self.work_dir,
                                                method='1',
                                                calib_iter=self.quant_iter)
        self.decent_q.quantize_frozen(input_graph_def, inputs_func, q_config)

        netcfg = os.path.join(self.work_dir, "deploy_model.pb")
        q_eval_file = os.path.join(self.work_dir, "quantize_eval_model.pb")
        quant_info_file = os.path.join(
            self.work_dir, 'quant_info_{}.txt'.format(xgraph.get_name()))
        self._save_quant_info(netcfg, quant_info_file)

        self.q_output.add(xgraph.get_name(), netcfg, quant_info_file,
                          frozen_graph, q_eval_file)

        # TODO
        # Add quantization info to corresponding XLayers
        self._add_quant_info_to_xgraph(netcfg)

    def quantize(self) -> None:
        """Quantize the XGraph model using the decent_q quantizer"""

        # NOTE For Conv2Dtranspose layers we need the specific batch size in
        #   tensorflow 1.13
        batch_size = list(self.inputs_func(0).values())[0].shape[0]

        fs = self.gen.generate(self.xgraph,
                               'graph',
                               subgraphs_only=True,
                               layout='NHWC',
                               batch_size=batch_size,
                               out_dir=self.work_dir,
                               **self.kwargs)

        if len(fs) != 1:
            raise ValueError("DECENT quantization currently only supports"
                             " models with one DPU compatible partition,"
                             " but got: {}".format(len(fs)))

        partition_key = list(fs.keys())[0]
        pb_path = list(fs.values())[0]

        self.partition_graphs[partition_key] = pb_path

        q_xgraph = super(DECENTQuantizer, self).quantize()

        self.xgraph.meta_attrs["is_quantized"] = True
        for qkey in self.q_output.keys():
            if 'quant_keys' not in self.xgraph.meta_attrs:
                self.xgraph.meta_attrs['quant_keys'] = [qkey]
            else:
                self.xgraph.meta_attrs['quant_keys'].append(qkey)
            quant_file = self.q_output.get_q_file(qkey)
            quant_info_file = self.q_output.get_q_info(qkey)
            quant_orig_pb = self.q_output.get_orig_pb(qkey)
            quant_eval_file = self.q_output.get_q_eval(qkey)
            self.xgraph.meta_attrs[qkey] = {
                'q_file': quant_file,
                'q_info': quant_info_file,
                'orig_pb': quant_orig_pb,
                'q_eval': quant_eval_file
            }

        self.xgraph.set_quantizer_output(self.q_output)
        # import pdb; pdb.set_trace()

        return q_xgraph

    def _add_quant_info_to_xgraph(self, deploy_frozen_graph: str) -> None:
        """
        Retrieve the quantization info from the provided quantized model and
        add the information to the corresponding XLayers
        """

        # Import tensorflow only when needed to avoid strict dependency
        import tensorflow as tf

        quant_info = []

        input_graph_def = tf.compat.v1.GraphDef()
        with tf.io.gfile.GFile(deploy_frozen_graph, "rb") as f:
            input_graph_def.ParseFromString(f.read())

            for idx, node in enumerate(input_graph_def.node):

                if node.name in self.xgraph:
                    X = self.xgraph.get(node.name)
                    X.attrs['vai_quant_idx'] = idx + 1

                    if 'ipos' in node.attr.keys():
                        X.attrs['vai_quant'] = ['vai_quant_in']
                        X.attrs['vai_quant_in'] = \
                            [int(v) for v in node.attr['ipos'].list.i]
                    if 'opos' in node.attr.keys():
                        X.attrs['vai_quant'].append('vai_quant_out')
                        X.attrs['vai_quant_out'] = \
                            [int(v) for v in node.attr['opos'].list.i]
                    if 'wpos' in node.attr.keys():
                        X.attrs['vai_quant'].append('vai_quant_weights')
                        X.attrs['vai_quant_weights'] = \
                            [int(v) for v in node.attr['wpos'].list.i]
                    if 'bpos' in node.attr.keys():
                        X.attrs['vai_quant'].append('vai_quant_biases')
                        X.attrs['vai_quant_biases'] = \
                            [int(v) for v in node.attr['bpos'].list.i]

    def _save_quant_info(self, deploy_frozen_graph, filename):
        # type: (str) -> None
        """
        Retrieve the quantization info from the provided quantized model
        """
        quant_info = self._get_quant_info(deploy_frozen_graph)

        lines = [[q_op['idx']] + [q_op['name']] +
                 [str(i)
                  for i in q_op['ipos']] + [str(i) for i in q_op['opos']] +
                 [str(i)
                  for i in q_op['wpos']] + [str(i) for i in q_op['bpos']]
                 for q_op in quant_info]
        s = '\n'.join([' '.join(line) for line in lines])

        with open(filename, 'w') as f:
            f.write(s)

    def _get_quant_info(self, deploy_frozen_graph):
        # type: (str) -> List[dict]
        """
        Retrieve the quantization info from the provided quantized model
        """

        # import tensorflow only when needed to avoid strict dependency
        import tensorflow as tf

        quant_info = []

        input_graph_def = tf.compat.v1.GraphDef()
        with tf.io.gfile.GFile(deploy_frozen_graph, "rb") as f:
            input_graph_def.ParseFromString(f.read())

            for idx, node in enumerate(input_graph_def.node):

                q_op = {
                    'idx': str(idx + 1),
                    'name': node.name,
                    'ipos': [],
                    'opos': [],
                    'wpos': [],
                    'bpos': []
                }

                if 'ipos' in node.attr.keys():
                    q_op['ipos'].extend(
                        [int(v) for v in node.attr['ipos'].list.i])
                if 'opos' in node.attr.keys():
                    q_op['opos'].extend(
                        [int(v) for v in node.attr['opos'].list.i])
                if 'wpos' in node.attr.keys():
                    q_op['wpos'].extend(
                        [int(v) for v in node.attr['wpos'].list.i])
                if 'bpos' in node.attr.keys():
                    q_op['bpos'].extend(
                        [int(v) for v in node.attr['bpos'].list.i])

                quant_info.append(q_op)

        return quant_info

    def eval(self,
             val_dir,
             gold_file,
             synset_words,
             batch_size,
             nb_batches,
             class_num=1000,
             gpu=0):
        #
        """
        """

        input_fn_data = {
            "prep_key": self.data_prep_key,
            "dir": val_dir,
            "batch": batch_size,
            "inputs": self.xgraph.get_input_names()
        }

        with open(os.path.join(FILE_PATH, 'calibration.json'), 'w') as f:
            json.dump(input_fn_data, f)

        with open(gold_file) as f:
            val_set = [line.strip('\n').split(' ') for line in f.readlines()]

        # frozen_graph_file = os.path.join(os.getcwd(), 'test.pb')
        frozen_graph_file = os.path.join(self.output_dir,
                                         "quantize_eval_model.pb")
        # TODO
        assert (len(self.xgraph.get_input_names()) == 1)
        assert (len(self.xgraph.get_output_names()) == 1)
        input_node = self.xgraph.get_input_names()[0]
        output_node = self.xgraph.get_output_names()[0]

        os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu)
        input_graph_def = tf.Graph().as_graph_def()
        input_graph_def.ParseFromString(
            tf.gfile.FastGFile(frozen_graph_file, "rb").read())

        tf.import_graph_def(input_graph_def, name='')

        # Get input tensors
        input_tensor = tf.get_default_graph()\
            .get_tensor_by_name(input_node+':0')
        input_labels = tf.compat.v1.placeholder(tf.float32,
                                                shape=[None, class_num])

        # Calculate accuracy
        output = tf.get_default_graph().get_tensor_by_name(output_node + ':0')
        prediction = tf.reshape(output, [batch_size, class_num])
        # correct_labels = tf.argmax(input_labels, 1)
        # top1_prediction = tf.nn.in_top_k(prediction, correct_labels, k = 1)
        # top5_prediction = tf.nn.in_top_k(prediction, correct_labels, k = 5)
        # top1_accuracy = tf.reduce_mean(tf.cast(top1_prediction,'float'))
        # top5_accuracy = tf.reduce_mean(tf.cast(top5_prediction,'float'))

        # Start evaluation
        logger.info("Start Evaluation for {} Batches...".format(nb_batches))
        with tf.Session() as sess:
            progress = ProgressBar()

            top1_sum_acc = 0
            top5_sum_acc = 0

            for iter in progress(range(0, nb_batches)):
                input_data = decent_prepfn.input_fn(iter)
                images = input_data[input_node]
                # labels = input_data['labels']
                logger.debug("IMAGES", images)
                labels = [
                    elem[1] for elem in val_set[iter * batch_size:(iter + 1) *
                                                batch_size]
                ]
                feed_dict = {input_tensor: images}
                raw_predictions = sess.run(prediction, feed_dict)
                logger.debug(raw_predictions)

                # logger.debug("Predictions shape: {}"
                #              .format(raw_predictions.shape))
                # logger.debug("Labels length: {}".format(len(labels)))
                top_1 = classification.get_top_k_accuracy(
                    raw_predictions, synset_words, 1, labels)
                top_5 = classification.get_top_k_accuracy(
                    raw_predictions, synset_words, 5, labels)
                top1_sum_acc += top_1
                top5_sum_acc += top_5
                logger.debug("int: {}, {}".format(top_1, top_5))

        final_top1_acc = top1_sum_acc / nb_batches
        final_top5_acc = top5_sum_acc / nb_batches

        print("Accuracy: Top1: {}, Top5: {}".format(final_top1_acc,
                                                    final_top5_acc))

    def dump(self, img_dir, input_names, max_dump_batches=1, dump_float=0):
        #
        """
        TODO: inupt_names
        """
        input_fn_data = {
            "prep_key": self.data_prep_key,
            "dir": img_dir,
            "batch": 1,
            "inputs": input_names
        }

        with open(os.path.join(FILE_PATH, 'calibration.json'), 'w') as f:
            json.dump(input_fn_data, f)

        frozen_graph = os.path.join(self.output_dir, 'quantize_eval_model.pb')

        command = """
        decent_q dump \
            --input_frozen_graph {} \
            --input_fn decent_prepfn.input_fn \
            --max_dump_batches {} \
            --dump_float {} \
            --output_dir {}
        """.format(frozen_graph, max_dump_batches, dump_float, self.output_dir)

        print("COMMAND", command)

        process = subprocess.Popen(command.split(),
                                   cwd=FILE_PATH,
                                   stdout=subprocess.PIPE)
        output, error = process.communicate()

        print(output, error)