habana_frameworks.mediapipe.fn.Split
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.
|
dtype |
Output data type.
|
Note
Input and output tensors must have the same data type.
Aggregate size along aggregate dimension must match input tensor size in that dimension.
Shape of output tensors must match input tensor in all dimensions except the split dimension.
Currently splitting into two tensor is supported along last dimension and it 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)