tf.keras.utils.Sequence

1、前言

网上有很多这方面的介绍，我这里以加载文件来验证。

使用 tf.keras.utils.Sequence 必须重写__init__， __len__， __getitem__
可选重写 on_epoch_end
还可以添加自定义的函数。

2、代码

class DataSet(tf.keras.utils.Sequence):
    def __init__(self, img_path, lab_path, batch_size):     # 必要函数，初始化参数
        super(DataSet, self).__init__()
        self.img_path = img_path
        self.lab_path = lab_path
        self.batch_size = batch_size
        self.img_nums = len(os.listdir(self.img_path))
        self.soft_list = np.arange(0, self.img_nums)

        self.load_data()

    def __len__(self):      # 必要函数，确定每一个EPOCH需要多少步
        return math.ceil(self.img_nums / self.batch_size)

    def __getitem__(self, item):    # 必要函数，在每个EPOCH内，依次返回一个batch_size的数据，item指当前是第几个batch
        # 封装数据
        return self.img_arr[self.batch_size * item : self.batch_size * (item + 1)], \
               self.lab_arr[self.batch_size * item : self.batch_size * (item + 1)]

    def on_epoch_end(self):         # 可选函数，在每个EPOCH结束时要做的事情，比如重新打乱数据
        random.shuffle(self.soft_list)

    def load_data(self):            # 自定义函数，读取数据、数据增强等 可在这里进行。
        self.img_arr = []
        self.lab_arr = []
        for i in np.array(os.listdir(self.img_path))[self.soft_list]:
            self.img_arr.append(os.path.join(self.img_path, i))
            self.lab_arr.append(os.path.join(self.lab_path, i))


if __name__ == '__main__':
    dataset = DataSet(img_path=r'D:\dataset\my_test\image',
                      lab_path=r'D:\dataset\my_test\label',
                      batch_size=2)

    for idx, i in enumerate(dataset):
        print('第{}个batch，数据为：{}'.format(idx, i))

3、输出结果为：

第0个batch，数据为：(['D:\\dataset\\my_test\\image\\val_0.tif', 'D:\\dataset\\my_test\\image\\val_1.tif'], ['D:\\dataset\\my_test\\label\\val_0.tif', 'D:\\dataset\\my_test\\label\\val_1.tif'])
第1个batch，数据为：(['D:\\dataset\\my_test\\image\\val_10.tif', 'D:\\dataset\\my_test\\image\\val_11.tif'], ['D:\\dataset\\my_test\\label\\val_10.tif', 'D:\\dataset\\my_test\\label\\val_11.tif'])
第2个batch，数据为：(['D:\\dataset\\my_test\\image\\val_12.tif', 'D:\\dataset\\my_test\\image\\val_2.tif'], ['D:\\dataset\\my_test\\label\\val_12.tif', 'D:\\dataset\\my_test\\label\\val_2.tif'])
第3个batch，数据为：(['D:\\dataset\\my_test\\image\\val_3.tif', 'D:\\dataset\\my_test\\image\\val_4.tif'], ['D:\\dataset\\my_test\\label\\val_3.tif', 'D:\\dataset\\my_test\\label\\val_4.tif'])
第4个batch，数据为：(['D:\\dataset\\my_test\\image\\val_5.tif', 'D:\\dataset\\my_test\\image\\val_6.tif'], ['D:\\dataset\\my_test\\label\\val_5.tif', 'D:\\dataset\\my_test\\label\\val_6.tif'])
第5个batch，数据为：(['D:\\dataset\\my_test\\image\\val_7.tif', 'D:\\dataset\\my_test\\image\\val_8.tif'], ['D:\\dataset\\my_test\\label\\val_7.tif', 'D:\\dataset\\my_test\\label\\val_8.tif'])
第6个batch，数据为：(['D:\\dataset\\my_test\\image\\val_9.tif'], ['D:\\dataset\\my_test\\label\\val_9.tif'])