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
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