tf.data.Dataset的一些小坑

我们使用数据的时候都是用batch来做输入,使用tf.data.Dataset的时候,一般会这样写:

dataset = dataset.batch(batch_size).repeat(epochs)

用来说明我们需要对整个数据集进行多少个epochs,每次的输入大小是多少个batch.
注意:
如果我们的数据集的数量为N,而N%batch_size刚好能整除的话,上述代码是没有任何bug的,但如果整除不了,那么在每个epoch的最后一个batch,其数据不再是batch_size个,而是N%batch_size个数据. 程序中如果有设定好tf.placeholder来修饰input,那么程序在运行到最后一个batch的时候就会报错,因为batch数量对不上. 所以,无论怎样,一个比较好的代码习惯是这么写:

dataset = dataset.batch(batch_size, drop_remainder=True).repeat(epochs)

这样程序会自动把最后一个不足batch_size的batch给忽略掉.