habana_frameworks.mediapipe.fn.GatherND

Class:

• habana_frameworks.mediapipe.fn.GatherND(**kwargs)

Define graph call:

• __call__(input, indices)

Parameter:

• input - Input tensor to operator. Supported dimensions: minimum = 1, maximum = 5. Supported data types: UINT8, UINT16.

• indices - Tensor containing indices for gather. Supported dimensions: minimum = 1, maximum = 5. Supported data types: INT32.

Description:

This operator gathers slices of input tensor into an output tensor of a specific rank. Let the input tensor be of rank r >= 1, and index tensor be of rank q >= 1, then the output tensor will be of rank q + r - FCD of the indices - 1, where FCD is the Fastest Changing Dimension. This operator is ONNX 1.6 GatherND compliant which is functionally equivalent as TF 2.4.1 gather_nd.

Supported backend:

• HPU

Keyword Arguments

kwargs	Description
dtype	Output data type. Type: habana_frameworks.mediapipe.media_types.dtype Default: UINT8 Optional: yes Supported data type: UINT8

Note

1. The FCD size of the index tensor must be less than the rank of the input tensor.

2. According to ONNX: output shape = ‘Index_feature_map.shape[0:q-1] ++ Input_feature_map.shape[q:r-1], where ++ denotes the concatenation of shapes.’ Since FCD first notation is followed in TPC, this shape is reversed.

Example: GatherND Operator

The following code snippet shows usage of GatherND operator:

from habana_frameworks.mediapipe import fn
from habana_frameworks.mediapipe.mediapipe import MediaPipe
from habana_frameworks.mediapipe.media_types import imgtype as it
from habana_frameworks.mediapipe.media_types import dtype as dt
import matplotlib.pyplot as plt
import numpy as np

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

        self.input = fn.ReadImageDatasetFromDir(shuffle=False,
                                                dir=dir,
                                                format="jpg")

        # WHCN
        self.decode = fn.ImageDecoder(device="hpu",
                                    output_format=it.RGB_P,
                                    resize=[img_w, img_h])

        indices_data = np.array([[5], [4], [3], [2], [1], [0]], dtype='int32')

        self.indices = fn.MediaConst(data=indices_data,
                                    shape=[1, batch_size],
                                    dtype=dt.INT32)

        self.gather_nd = fn.GatherND(dtype=dt.UINT8)

        # WHCN -> CWHN
        self.transpose = fn.Transpose(permutation=[2, 0, 1, 3],
                                    tensorDim=4,
                                    dtype=dt.UINT8)

    def definegraph(self):
        images, labels = self.input()
        images = self.decode(images)
        indices = self.indices()
        images = self.gather_nd(images, indices)
        images = self.transpose(images)
        return images, labels

def display_images(images, batch_size, cols):
    rows = (batch_size + 1) // cols
    plt.figure(figsize=(10, 10))
    for i in range(batch_size):
        ax = plt.subplot(rows, cols, i + 1)
        plt.imshow(images[i])
        plt.axis("off")
    plt.show()

def main():
    batch_size = 6
    img_width = 200
    img_height = 200
    img_dir = "/path/to/images"
    queue_depth = 2
    columns = 3

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

    # Build media pipeline
    pipe.build()

    # Initialize media pipeline iterator
    pipe.iter_init()

    # Run media pipeline
    const_data, labels = pipe.run()

    # Copy data to host from device as numpy array
    images = const_data.as_cpu().as_nparray()
    labels = labels.as_cpu().as_nparray()

    # Display images
    display_images(images, batch_size, columns)

if __name__ == "__main__":
    main()

Images Gathered in Reverse Order from a Batch 1

image #6

image #5

image #4

image #3

image #2

image #1

1

Licensed under a CC BY SA 4.0 license. The images used here are taken from https://data.caltech.edu/records/mzrjq-6wc02.