Example #1
0
    def test_pixel_shuffle(self):
        scale_factor = 2
        input_shape = (1, 8, 2, 2)
        output_shape = (input_shape[0], int(
            input_shape[1] / (scale_factor**2)), input_shape[2] * scale_factor,
                        input_shape[3] * scale_factor)

        inputs = [('input0', input_shape)]
        outputs = [('output0', output_shape)]

        node_0 = helper.make_node("Reshape",
                                  inputs=[inputs[0][0]],
                                  outputs=['node0'],
                                  shape=[
                                      output_shape[0], output_shape[1],
                                      scale_factor, scale_factor,
                                      input_shape[2], input_shape[3]
                                  ])
        node_1 = helper.make_node("Transpose",
                                  inputs=['node0'],
                                  outputs=['node1'],
                                  perm=[0, 1, 4, 2, 5, 3])
        node_2 = helper.make_node("Reshape",
                                  inputs=['node1'],
                                  outputs=[outputs[0][0]],
                                  shape=list(output_shape))
        model = _onnx_create_model([node_0, node_1, node_2], inputs, outputs)
        _test_onnx_model(model, decimal=7)
    def test_pixel_shuffle(self):  # type: () -> None
        scale_factor = 2
        input_shape = (1, 8, 2, 2)
        output_shape = (
            input_shape[0],
            int(input_shape[1] / (scale_factor ** 2)),
            input_shape[2] * scale_factor,
            input_shape[3] * scale_factor
        )

        inputs = [('input0', input_shape)]
        outputs = [('output0', output_shape, TensorProto.FLOAT)]

        shape1 = [
            output_shape[0],
            output_shape[1],
            scale_factor,
            scale_factor,
            input_shape[2],
            input_shape[3]
        ]

        shape1 = numpy_helper.from_array(np.asarray(shape1), name="shape1")
        shape2 = numpy_helper.from_array(np.asarray(list(output_shape)), name="shape2")

        node_0 = helper.make_node(
            "Reshape",
            inputs=[inputs[0][0], 'shape1'],
            outputs=['node0'],
        )
        node_1 = helper.make_node(
            "Transpose",
            inputs=['node0'],
            outputs=['node1'],
            perm=[0, 1, 4, 2, 5, 3]
        )
        node_2 = helper.make_node(
            "Reshape",
            inputs=['node1','shape2'],
            outputs=[outputs[0][0]],
        )
        model = _onnx_create_model(
            [node_0, node_1, node_2], inputs, outputs, initializer=[shape1, shape2]
        )
        _test_onnx_model(model, decimal=7)