Habana Media Loader

The Habana Media Loader is used to dataload and pre-process inputs for deep learning frameworks. It consists of pre-enabled dataloaders for commonly used datasets, currently only ImageNet is supported, and building blocks to assemble a generic dataloader. The loader decides internally if part of operations can be offloaded to Habana accelerator. If the offload cannot be executed, it will result in using the passing alternative dataloader function to run.

Habana Media Loader can operate in different modes, the optimal one is selected based on the underlying hardware:

• In first-gen Gaudi, it uses either the framework default dataloader or AEON based dataloader, depending on the use case. Both are done on the host CPU.

• In Gaudi2, the dataloader uses hardware-based decoders for acceleration, lowering the load on the host CPU.

Setting Up the Environment

To install Habana Media Loader, run the following command:

pip install habana_media_loader-1.13.0-463-py3-none-any.whl

Note

If you are using Habana docker image, skip this step as Habana Media Loader is pre-installed.

Using Media Loader with TensorFlow

• Import Habana ImageNet dataset object:

from habana_frameworks.tensorflow.media import habana_imagenet_dataset

• Create the dataset object that exposes the same APIs that are aligned with TensorFlow DataLoader APIs. Refer to the example in TensorFlow Model References GitHub page.

return habana_imagenet_dataset(fallback,
                         is_training,
                         tf_data_dir,
                         jpeg_data_dir,
                         batch_size,
                         num_channels,
                         img_size,
                         dtype,
                         **fallback_kwargs)

Parameter	Description
Fallback	Fallback function that would run if Habana media dataloader cannot be used
is_training	Boolean variable indicates if the dataloader is in training or evaluation mode
tf_data_dir	Path to dataset with tf_records and is used for fallback
jpeg_data_dir	Path to dataset with JPEGs and is used for media pipe
batch_size	Batch size
num_channels	Number of channels
img_size	Image size
dtype	Data loader image type
Fallback_kwargs	Other kwargs that should be passed to fallback function