Tensorflow tf.dataset.shuffle 非常慢

Question

我正在用 9100 張圖像（每個尺寸為 256 x 64）訓練 VAE model。 我用 Nvidia RTX 3080 訓練 model。首先，我將所有圖像加載到大小為 9100 x 256 x 64 的 numpy 數組中，稱為traindata 。 然后，為了形成一個訓練數據集，我使用

train_dataset = (tf.data.Dataset.from_tensor_slices(traindata).shuffle(len(traindata)).batch(batch_size))

這里我使用了65的batch_size 。關於訓練時看到的東西，我主要有2個問題：

問題一：

根據文檔，整個數據集在每個時期都被重新洗牌。 但是，這種方式的訓練非常慢（每個 epoch 大約 50 秒）。 我通過在創建數據集時不調用.shuffle(len(traindata))與沒有 shuffle 的訓練進行了比較，訓練速度要快得多（大約 20s/epoch）。 我想知道為什么.shuffle()操作如此緩慢，是否有任何方法可以使其更快？ 根據這個 StatsSE thread ，洗牌對於訓練非常重要，這就是我包含洗牌操作的原因。

問題2：

當我在創建數據集時調用.shuffle()時，Tensorflow 總是給出以下消息

I tensorflow/core/platform/windows/subprocess.cc:308] SubProcess ended with return code: 4294967295

我嘗試在網上搜索，但仍然無法理解這背后的含義。 這是否意味着有一些錯誤，或者這只是一個我可以忽略的警告？

Answer 1

這是因為將數據集的所有元素保存在緩沖區中是昂貴的。 除非您絕對需要完美的隨機性，否則您應該使用較小的buffer_size 。 最終將采用所有元素，但以更具確定性的方式。

這就是使用較小的buffer_size會發生的事情，比如 3。緩沖區是括號，Tensorflow 在這個括號中采樣一個隨機值。 隨機抽取的是^

1) [1 2 3]4 5 6 7 8 9 
      ^
2) [1 3 4]5 6 7 8
        ^
3) [1 3 5]6 7 8
        ^
4) [1 3 6]7 8
    ^
5) [3 6 7]8

ETC

因此，較早的值將在您的 epoch 中較早地獲取，但您仍將完成一些洗牌，並且最終將獲取所有樣本。

tl;博士將buffer_size減少了很多