TensorFlow の shuffle=batch 問題

大学の研究で，HDF5形式の学習用データを使って深層学習モデルを生成する必要がありました．そこで，tensorflow で shuffle=batch というオプションを指定してモデルの生成を試みたのですが，loss が全く減りませんでした泣． その原因と解決策についてまとめてみました！！

実行環境

tensorflow : 2.6.0

shuffle=batch とは

学習の実行時に呼び出す Model.fit() のオプションとなっています．

で説明されているように，suffle=batch を指定すると，バッチサイズのチャンクで学習用データをシャッフルします（下図）．

suffle=batch は，今回学習用データとして使用した HDF5 形式でよく使われるシャッフルのやり方みたいです ※（記事の最下部参照）

shuffle=batch を指定して tensorflow で学習させてみた

shuffle=batch を指定して tensorflow で学習させてみました．

autoencoder.fit(dl_train, dl_train,
                    epochs=200,
                    validation_data=(dl_val, dl_val),
                    batch_size=50,
                    shuffle='batch',
                    callbacks=[tensorboard, reduce_lr, checkpoint])

loss がほとんど変わらない泣

学習の途中経過を見てみると，以下のようになっており loss がほとんど変わりません泣

80/80 [==============================] - 47s 523ms/step - loss: 0.4781 - accuracy: 0.4976 - val_loss: 0.4853 - val_accuracy: 0.4978
Epoch 2/200
80/80 [==============================] - 41s 518ms/step - loss: 0.4781 - accuracy: 0.5016 - val_loss: 0.4853 - val_accuracy: 0.4978
Epoch 3/200
80/80 [==============================] - 42s 522ms/step - loss: 0.4781 - accuracy: 0.5032 - val_loss: 0.4853 - val_accuracy: 0.5027
...
Epoch 60/200
80/80 [==============================] - 42s 525ms/step - loss: 0.4780 - accuracy: 0.4996 - val_loss: 0.4854 - val_accuracy: 0.4978

原因

色々と調べてみたところ，tensorflow のバージョン に原因がありそうでした．

tensorflow の Github に以下のような Issue があり，これによると，

tensorflow 1.15, 2.0.0 では shuffle=batch はうまく動くけど，tensorflow 2.1.0 , 2.2.0, 2.3.0, 2.5.0 ではうまく動かないということが報告されていました．
私が使用していた tensorflow は 2.6.0 だったので，バージョンが原因かもしれないと思いました．

解決策 1 : tensorflow のバージョンを落とす

tensorflow のバージョンを 2.0.0 に落として再実行してみました．すると...

4000/4000 [==============================] - 56s 14ms/sample - loss: 0.3748 - accuracy: 0.5942 - val_loss: 0.3382 - val_accuracy: 0.6750
Epoch 2/200
4000/4000 [==============================] - 52s 13ms/sample - loss: 0.2189 - accuracy: 0.8033 - val_loss: 0.1352 - val_accuracy: 0.8672
Epoch 3/200
4000/4000 [==============================] - 53s 13ms/sample - loss: 0.1021 - accuracy: 0.8826 - val_loss: 0.0902 - val_accuracy: 0.8820
Epoch 4/200
4000/4000 [==============================] - 53s 13ms/sample - loss: 0.0761 - accuracy: 0.8939 - val_loss: 0.0732 - val_accuracy: 0.8976
Epoch 5/200
4000/4000 [==============================] - 53s 13ms/sample - loss: 0.0642 - accuracy: 0.9047 - val_loss: 0.0556 - val_accuracy: 0.9133
Epoch 6/200
4000/4000 [==============================] - 53s 13ms/sample - loss: 0.0535 - accuracy: 0.9119 - val_loss: 0.0477 - val_accuracy: 0.9237
Epoch 7/200
4000/4000 [==============================] - 53s 13ms/sample - loss: 0.0519 - accuracy: 0.9175 - val_loss: 0.0437 - val_accuracy: 0.9294

loss が減るようになりました！！

しかし，新しいバージョンの方が嬉しいことも沢山あると思います．例えば 1epoch あたりの学習時間は，2.6.0 の方が 2.0.0 よりも短いです（42s/epoch < 53s/epoch)．
そこで，2.6.0 で shuffle=batch と同等の機能を実装する方法を調べてみました！！

解決策 2 : fit_generator を使う

fit_generator を使うと，学習用データからミニバッチを生成する処理を自分で書くことができます．詳しい説明は省略しますが，この記事が参考になりました．

この fit_generator を使って，shuffle=batch と同等の機能を実装しました．

以下のように fit_generator を実装して，

my_generator.py

from tensorflow.keras.utils import Sequence
import h5py
import numpy as np


class MyGenerator(Sequence):
    """Custom generator"""

    def __init__(self, data_path, data_size, key_name,  
                 batch_size):
        """construction   
        :param data_paths: 
        :param data_size: Num of dataset sample
        :param batch_size: Batch size  
        """

        self.data_path = data_path
        self.data_size = data_size
        self.key_name = key_name
        self.batch_size = batch_size
        self.steps_per_epoch = int((self.data_size - 1) / batch_size) + 1
        self.shuffle_indices = None 
        self.step_cnt = 0
        

    def __getitem__(self, idx):
        """Get batch data   

        :param idx: Index of batch  

        :return imgs: numpy array of images 
        :return labels: numpy array of label  
        """
        if self.step_cnt == 0:
            index_array = np.array(range(self.data_size))
            last_batch = index_array[self.steps_per_epoch * self.batch_size:]
            index_array = index_array[:self.steps_per_epoch * self.batch_size]
            index_array = index_array.reshape((self.steps_per_epoch, self.batch_size))
            np.random.shuffle(index_array)
            index_array = index_array.flatten()
            self.shuffle_indices = np.append(index_array, last_batch)

        start_pos = self.batch_size * idx
        end_pos = start_pos + self.batch_size
        if end_pos > self.data_size:
            end_pos = self.data_size
        
        self.step_cnt += 1
        if self.step_cnt >= self.steps_per_epoch:
            self.step_cnt = 0
        inputs = h5py.File(self.data_path, 'r')[self.key_name][self.shuffle_indices[start_pos] : self.shuffle_indices[end_pos-1]+1, :, :, :, :]

        return inputs, inputs


    def __len__(self):
        """Batch length"""
        return self.steps_per_epoch


    def on_epoch_end(self):
        """Task when end of epoch"""
        pass

以下のように使うと，

train_gen = MyGenerator(train_data_path, train_data_size, 'data_train', batch_size)
val_gen = MyGenerator(val_data_path, val_data_size, 'data_val', batch_size)

autoencoder.fit_generator(
        train_gen, 
        steps_per_epoch=train_gen.steps_per_epoch, 
        validation_data=val_gen, 
        validation_steps=val_gen.steps_per_epoch, 
        epochs=200,
        callbacks=[tensorboard, reduce_lr, checkpoint])

上手く loss が減るようになりました！！

80/80 [==============================] - 43s 485ms/step - loss: 0.3824 - accuracy: 0.6273 - val_loss: 0.3511 - val_accuracy: 0.6535
Epoch 2/200
80/80 [==============================] - 38s 481ms/step - loss: 0.2702 - accuracy: 0.6870 - val_loss: 0.2303 - val_accuracy: 0.7449
Epoch 3/200
80/80 [==============================] - 38s 481ms/step - loss: 0.1763 - accuracy: 0.8063 - val_loss: 0.1357 - val_accuracy: 0.8527
Epoch 4/200
80/80 [==============================] - 42s 522ms/step - loss: 0.1093 - accuracy: 0.8637 - val_loss: 0.0988 - val_accuracy: 0.8765
Epoch 5/200
80/80 [==============================] - 43s 537ms/step - loss: 0.0799 - accuracy: 0.8846 - val_loss: 0.0771 - val_accuracy: 0.8873
Epoch 6/200
80/80 [==============================] - 43s 537ms/step - loss: 0.0594 - accuracy: 0.9026 - val_loss: 0.0517 - val_accuracy: 0.9140
Epoch 7/200
80/80 [==============================] - 43s 537ms/step - loss: 0.0494 - accuracy: 0.9114 - val_loss: 0.0454 - val_accuracy: 0.9197
Epoch 8/200
80/80 [==============================] - 43s 537ms/step - loss: 0.0438 - accuracy: 0.9179 - val_loss: 0.0419 - val_accuracy: 0.9250
Epoch 9/200
80/80 [==============================] - 43s 537ms/step - loss: 0.0407 - accuracy: 0.9209 - val_loss: 0.0368 - val_accuracy: 0.9291

tensorflow 2.6.0 で実行しているので，1epoch あたりの学習時間も解決策1より短く済んでいます！

まとめ

tensorflow 2.6.0 では，shuffle=batch をすると loss が減らなくなるという現象が起こります．そのため，tensorflow のバージョンを 2.0.0 以下に下げるか，shuffle=batch と同等の機能を fit_generator を使って実装する必要があります．

早く tensorflow 2.0.0 以降でもshuffle=batchが使えると良いなーと思います．

※ shffle=batch を指定しない場合は，shuffle=true と同じになります．shuffle=true では，学習用データ全体をシャッフルして一括で読み込みます．しかし，HDF5 形式のデータを読み込むときはデータのインデックスが昇順になっていないとダメらしく，shuffle=true としてしまうと，Indexing elements must be in increasing order というエラーが起こります．