neural_compressor.data.dataloaders.tensorflow_dataloader

TensorFlow Dataloader implementation.

Module Contents

Classes

TFDataDataLoader	Tensorflow dataloader class.
TensorflowBertDataLoader	Subclass of DefaultDataLoader.
TensorflowModelZooBertDataLoader	Subclass of DefaultDataLoader.
TensorflowDataLoader	DataLoader for framework Tensorflow.

class neural_compressor.data.dataloaders.tensorflow_dataloader.TFDataDataLoader(dataset, batch_size=1, last_batch='rollover')

Tensorflow dataloader class.

In tensorflow1.x dataloader is coupled with the graph, but it also support feed_dict method to do session run, this dataloader is designed to satisfy the usage of feed dict in tf1.x. Although it’s a general dataloader and can be used in MXNet and PyTorch.

Parameters:

• dataset – obj. wrapper of needed data.

• batch_size – int. batch size

class neural_compressor.data.dataloaders.tensorflow_dataloader.TensorflowBertDataLoader

Subclass of DefaultDataLoader.

this dataloader is designed to satisfy the usage of Bert models.

class neural_compressor.data.dataloaders.tensorflow_dataloader.TensorflowModelZooBertDataLoader

Subclass of DefaultDataLoader.

this dataloader is designed to satisfy the usage of Model Zoo Bert models.

class neural_compressor.data.dataloaders.tensorflow_dataloader.TensorflowDataLoader(dataset, batch_size=1, last_batch='rollover', collate_fn=None, sampler=None, batch_sampler=None, num_workers=0, pin_memory=False, shuffle=False, distributed=False)

DataLoader for framework Tensorflow.

if it’s a tf.data.Dataset we will directly use the dataloader in the other case will use DefaultDataLoader instead.