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

Bases: neural_compressor.data.dataloaders.base_dataloader.BaseDataLoader

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

batch(batch_size, last_batch='rollover')

Dataset return data per batch.

class neural_compressor.data.dataloaders.tensorflow_dataloader.TensorflowBertDataLoader(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)

Bases: neural_compressor.data.dataloaders.default_dataloader.DefaultDataLoader

Subclass of DefaultDataLoader.

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

class neural_compressor.data.dataloaders.tensorflow_dataloader.TensorflowModelZooBertDataLoader(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)

Bases: neural_compressor.data.dataloaders.default_dataloader.DefaultDataLoader

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)

Bases: neural_compressor.data.dataloaders.base_dataloader.BaseDataLoader

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.