Class MyDataLoader

test/tensorflow/quantization/test_smooth_quant.py:53–73

class MyDataLoader:
    def __init__(self, dataset, batch_size=1):
        self.dataset = dataset
        self.batch_size = batch_size
        self.length = math.ceil(len(dataset) / self.batch_size)

    def __iter__(self):
        images_list = []
        labels_list = []
        for _, (images, labels) in enumerate(self.dataset):
            images = np.expand_dims(images, axis=0)
            labels = np.expand_dims(labels, axis=0)
            images_list.append(images[0])
            labels_list.append(labels[0])
            if self.batch_size == len(images_list):
                yield (images_list, labels_list)
                images_list = []
                labels_list = []

    def __len__(self):
        return self.length


class TestSmoothQuantTF3xNewApi(unittest.TestCase):