tf.data.Datasetバッチサイズが指定と異なる

Question

tf.data.Dataset.from_tensor_slices().batch()
を用いてdatasetとしてデータをバッチ化しようとしています。
.batch()の引数で指定したバッチサイズが出力と一致しません。

import tensorflow as tf

BATCH_SIZES = [256, 128, 96, 64, 48, 32, 16]

sliced = tf.data.Dataset.from_tensor_slices(tf.range(60000))
for i in BATCH_SIZES:
    dataset = sliced.batch(i)
    for item in dataset:
        item
    print(item.shape)


(96,)
(96,)
(96,)
(32,)
(48,)
(32,)
(16,)

バッチサイズ指定がが96以上では、バッチサイズが96。
バッチサイズ指定がが64では、バッチサイズが32。
バッチサイズ指定がが48以下では、バッチサイズが指定通り。
バッチサイズの上限が96でも48でもないし不可解。

ドキュメントではバッチサイズ指定は32ビットまで可能な筈。

Answer 1

データサイズである60000が、指定のバッチサイズで割り切れるかどうかで挙動が変わっています。

以下の簡単な実験からも、整合する状況であることが確認できるかと思います。
(バッチサイズ96以上: 0または96なので最後のバッチサイズは必ず96、バッチサイズ64: 最後のバッチサイズは32、残り: キレイに割り切れるため指定サイズのバッチ)

In [1]: [(i, 60000 % i) for i in [256, 128, 96, 64, 48, 32, 16]]
Out[1]: [(256, 96), (128, 96), (96, 0), (64, 32), (48, 0), (32, 0), (16, 0)]

もし末尾の余ったデータを破棄したい場合、batch() に drop_remainder=True を設定する必要があります。
https://www.tensorflow.org/api_docs/python/tf/data/Dataset#batch

Comments

[iwase13]

明確なご回答ありがとうございました。
Forループでの最後の割り切れなかったバッチ数が出力されたという事ですね。
Forの途中でも指定したバッチ数になっていないのではないかと勘違いしていました。