habana_frameworks.mediapipe.fn.Split

Class:
  • habana_frameworks.mediapipe.fn.Split(**kwargs)

Define graph call:
  • __call__(input)

Parameter:
  • input - Input tensor to operator. Supported dimensions: minimum = 1, maximum = 5. Supported data types: INT32, BFLOAT16, FLOAT32.

Description:

Splits a tensor into a list of tensors. The split is done along the specified ‘axis’. Lengths of the parts are specified in the sizes of the output parts.

Supported backend:
  • HPU

Keyword Arguments

kwargs

Description

axis

Axis along which tensors to be split.

  • Type: int

  • Default: 0

  • Optional: no

dtype

Output data type.

  • Type: habana_frameworks.mediapipe.media_types.dtype

  • Default: UINT8

  • Optional: yes

  • Supported data types:

    • INT32

    • FLOAT32

Note

  1. Input and output tensors must have the same data type.

  2. Aggregate size along aggregate dimension must match input tensor size in that dimension.

  3. Shape of output tensors must match input tensor in all dimensions except the split dimension.

  4. Currently, splitting into two tensors is supported. Split dimension should be divisible by 2.

Example: Split Operator

The following code snippet shows usage of Split operator.

from habana_frameworks.mediapipe import fn
from habana_frameworks.mediapipe.mediapipe import MediaPipe
from habana_frameworks.mediapipe.media_types import dtype as dt

# Create media pipeline derived class
class myMediaPipe(MediaPipe):
    def __init__(self, device, dir, queue_depth, batch_size):
        super(
            myMediaPipe,
            self).__init__(
            device,
            queue_depth,
            batch_size,
            self.__class__.__name__)

        self.input0 = fn.ReadNumpyDatasetFromDir(num_outputs=1,
                                                shuffle=False,
                                                dir=dir,
                                                pattern='x_*_split_inp.npy',
                                                dense=True,
                                                dtype=dt.FLOAT32)

        self.split = fn.Split(axis=4, dtype=dt.FLOAT32)

    def definegraph(self):
        input = self.input0()
        output0, output1 = self.split(input)
        return input, output0, output1

def main():
    batch_size = 2
    img_dir = '/path/to/numpyfiles'
    queue_depth = 2

    # Create media pipeline object
    pipe = myMediaPipe('hpu', img_dir, queue_depth, batch_size)

    # Build media pipeline
    pipe.build()

    # Initialize media pipeline iterator
    pipe.iter_init()

    # Run media pipeline
    input, output0, output1 = pipe.run()

    # Copy data to host from device as numpy array
    input = input.as_cpu().as_nparray()
    output0 = output0.as_cpu().as_nparray()
    output1 = output1.as_cpu().as_nparray()

    # Display shapes
    print('input shape:', input.shape)
    print('output0 shape:', output0.shape)
    print('output1 shape:', output1.shape)

if __name__ == "__main__":
    main()

The following is the output of Split operator:

input shape: (2, 4, 3, 128, 128)
output0 shape: (1, 4, 3, 128, 128)
output1 shape: (1, 4, 3, 128, 128)