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

num_outputs

Number of split outputs.

  • Type: int

  • Default: 0

  • Optional: no

dtype

Output data type for every split.

  • Type: list of habana_frameworks.mediapipe.media_types.dtype

  • Default: UINT8

  • Optional: no

  • 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 10 tensors maximum is supported.

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 MediaPipe derived class
class myMediaPipe(MediaPipe):
    def __init__(self, device, dir, queue_depth, batch_size, num_split):
        super(
            myMediaPipe,
            self).__init__(
            device,
            queue_depth,
            batch_size,
            self.__class__.__name__)
        self.num_split=num_split
        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, num_outputs=self.num_split, dtype=[dt.FLOAT32]*self.num_split)

    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
    num_split=2
    # Create MediaPipe object
    pipe = myMediaPipe('hpu', img_dir, queue_depth, batch_size, num_split)

    # Build MediaPipe
    pipe.build()

    # Initialize MediaPipe iterator
    pipe.iter_init()

    # Run MediaPipe
    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)