Habana Media Python package

habana_media_loader is a package designed for easy integration of media processing on Gaudi2. Main entry point (Python import) is habana_frameworks.mediapipe module that contains all the necessary functions to work with Gaudi2.

Structure

Properly built wheel contains:

• habana_frameworks python namespace (with all the folder structure inside).
• mediapipe folder catering media execution on device and medialoader folder catering pre-built mediapipe for pytorch framework.
• proper licensing.

Media package (habana_frameworks.mediapipe and habana_frameworks.medialoaders)

First part of media package contains media pipe which is responsible for media processing on device.

Following are the steps to create mediapipe

• Create a class derived from habana_frameworks.mediapipe super class.
• In the class constructor initialize super class.
• Create nodes required for execution along with it's parameters.
• Define a method definegraph() which defines the data flow between nodes created in constructor.

Following are the steps to execute a standalone mediapipe

• Instantiate an object of defined mediapipe class.
• Build the mediapipe by executing build() method of mediapipe object.
• Initialize the iterator by calling iter_init() method of mediapipe object.
• To produce one batch of dataset, execute run() method of mediapipe object. Each run() method call executes and produces one batch of device tensors.
• To view or manipulate tensors on host as_cpu() method of device tensor object can be called, which yields host tensor object.
• For numpy manipulation as_nparray() method of host tensor object can be called to get a numpy host array.

Example:

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
from habana_frameworks.mediapipe.media_types import layout as lt
import time

class myMediaPipe(MediaPipe):
    def __init__(self, device, queue_depth, batch_size, channel, height, width):
        super(
            myMediaPipe,
            self).__init__(
            device,
            queue_depth,
            batch_size,
            channel,
            height,
            width,
            self.__class__.__name__,
            layout=lt.NHWC)
        mediapipe_seed = int(time.time_ns() % (2**31 - 1))
        # create reader node and setting it's params
        self.input = fn.ReadImageDatasetFromDir(dir="/path/to/jpeg/dir/",
                                                format="JPEG",
                                                shuffle=True,
                                                seed=mediapipe_seed)
        # create decoder node and set it's params
        self.decode = fn.ImageDecoder(output_format=it.RGB_P,
                                      resize=[224, 224])

        # create transpose node and set it's params
        self.transpose = fn.Transpose(permutation=[2, 0, 1, 3], tensorDim=4)

    def definegraph(self):
        # define actual data flow of nodes
        jpegs, data = self.input()
        images = self.decode(jpegs)
        images = self.transpose(images)
        # return output nodes of the graph
        return images, data


# test specific params
batch_size = 4
img_width = 224
img_height = 224
channels = 3
queue_depth = 3
iterations = 5

# instantiating defined class
pipe = myMediaPipe("hpu", queue_depth, batch_size,
                   channels, img_height, img_width)
# build the pipe
pipe.build()
# initialize iterator
pipe.iter_init()

batch_count = 0
while(batch_count < iterations):
    try:
        # exectute and produce one batch of dataset.
        images, labels = pipe.run()
        # images and labels are device tensors.
    except StopIteration:
        print("stop iteration")
        break
    # as cpu will bring the device data to host and produce host tensors
    # as_nparray will convert host tensors to numpy array.
    images = images.as_cpu().as_nparray()
    labels = labels.as_cpu().as_nparray()
    batch_count = batch_count + 1

Second part of media package contains pre built media pipe for pytorch.

torch folder contains media_dataloader_mediapipe containing HPUMediaPipe which can be used to create resnet and SSD media pipe for pytorch

Following are the steps to use HPUMediaPipe for pytorch

• Import HPUMediaPipe from habana_frameworks.medialoaders.torch.media_dataloader_mediapipe
• Instantiate an object of HPUMediaPipe with following parameters:
  • a_torch_transforms: transforms to be applied on mediapipe.
  • a_root: directory path from which to load the images.
  • a_annotation_file: path from which to load annotation file for SSD.
  • a_batch_size: mediapipe output batch size.
  • a_shuffle: whether images have to be shuffled. <True/False>
  • a_drop_last: whether to drop the last incomplete batch or round up.<True/False>
  • a_prefetch_count: queue depth for media processing.
  • a_num_instances: number of devices.
  • a_instance_id: instance id of current device.
  • a_model_ssd: whether mediapipe is to be created for SSD. <True/False>
  • a_device: media device to run mediapipe on.
• Separate HPUMediaPipe objects can be created for training and validation.
• Instantiate an object of HPUResnetPytorchIterator (for resnet) or HPUSsdPytorchIterator (for SSD) with following parameters
  • mediapipe: media pipe object.

Example for resnet media pipe:

from habana_frameworks.medialoaders.torch.media_dataloader_mediapipe import HPUMediaPipe
from habana_frameworks.mediapipe.plugins.iterator_pytorch import HPUResnetPytorchIterator

normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225])
torch_transforms = transforms.Compose([
            transforms.RandomResizedCrop(224),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            normalize,
        ])

root = "/JPEG/path"
batch_size = 256
shuffle = True
drop_last = False
prefetch_factor = 3
num_instances = 1
instance_id = 0

pipeline = HPUMediaPipe(a_torch_transforms=torch_transforms, a_root=root, a_batch_size=batch_size,
                        a_shuffle=shuffle, a_drop_last=drop_last, a_prefetch_count=prefetch_factor,
                        a_num_instances=num_instances, a_instance_id=instance_id, a_device="hpu")

iterator = HPUResnetPytorchIterator(mediapipe=pipeline)