画像のマルチラベル分類モデルを学習したいのですが、メモリ不足でうまくいきません

前提・実現したいこと

機械学習を勉強して数ヶ月経った学生です。
現在multi-labelの画像分類モデルを作成しようとしています。

発生している問題・エラーメッセージ

学習が進んでいくにつれてメモリを大量に消費してしまい、必要量を確保できずエラーになってしまいます。
500MB程度のデータ(計5万枚程度)しか使っていないのにも関わらず、メモリ使用量が32GBを超えています。
自前で用意しているデータセットに何か問題があるのではないかと思っています。

該当のソースコード

（簡略化したもの）

python
1import ...
2
3AUTOTUNE = tf.data.experimental.AUTOTUNE
4
5def macro_f1(y, y_hat, thresh=0.5):
6    y_pred = tf.cast(tf.greater(y_hat, thresh), tf.float32)
7    tp = tf.cast(tf.math.count_nonzero(y_pred * y, axis=0), tf.float32)
8    fp = tf.cast(tf.math.count_nonzero(y_pred * (1 - y), axis=0), tf.float32)
9    fn = tf.cast(tf.math.count_nonzero((1 - y_pred) * y, axis=0), tf.float32)
10    f1 = 2*tp / (2*tp + fn + fp + 1e-16)
11    macro_f1 = tf.reduce_mean(f1)
12    return macro_f1
13
14def create_dataset(filenames, labels, is_training=True, total_cnt=0, batch_size=8, train_flag=False):
15    dataset = tf.data.Dataset.from_tensor_slices((filenames, labels))
16    if train_flag:
17        dataset = dataset.map(parse_function, num_parallel_calls=AUTOTUNE)
18    else:
19        dataset = dataset.map(parse_function_val, num_parallel_calls=AUTOTUNE)
20    dataset = dataset.cache()
21    if train_flag:
22        dataset = dataset.shuffle(buffer_size=total_cnt)
23    dataset = dataset.repeat().batch(batch_size)
24    dataset = dataset.prefetch(buffer_size=AUTOTUNE)
25    return dataset
26
27def parse_function(filename, label):
28    image_string = tf.io.read_file(filename)
29    image_decoded = tf.image.decode_jpeg(image_string, channels=3)
30    image_resized = tf.image.resize(image_decoded, [224, 224])
31    image_normalized = (image_resized / 255) -1
32    image_aug = tf.image.random_flip_left_right(image=image_normalized)
33    image_aug = tf.image.random_flip_up_down(image=image_aug)
34    return image_aug, label
35
36def parse_function_val(filename, label):
37    image_string = tf.io.read_file(filename)
38    image_decoded = tf.image.decode_jpeg(image_string, channels=3)
39    image_resized = tf.image.resize(image_decoded, [224, 224])
40    image_normalized = (image_resized / 255) -1
41    return image_normalized, label
42
43
44def main():
45
46    ~略~
47
48    # データセット作成用リスト
49    a_train    = []
50    b_train    = []
51    c_train    = []
52    image_list_train = []
53
54    a_val   = []
55    b_val   = []
56    c_val   = []
57    asphalt_val    = []
58    image_list_val = []
59
60    DATA_DIR = os.path.join(DIR, 'data')
61    modes = ['train', 'val']
62    classes = ['a', 'b', 'c', 'a_b', 'b_c', 'c_a', 'a_b_c'] #フォルダ名
63
64    #データセット作成
65    for mode in modes:
66        for class in classes:
67            dir_path = os.path.join(DATA_DIR, mode, class)
68            file_list = glob.glob(os.path.join(dir_path, '*'))
69            file_list = sorted(file_list)
70
71            for file_name in file_list:
72                # フォルダ名(ex: a_b)を_で分割し、a,b,cの文字が含まれているかどうかでラベルを付ける
73                cls_list = cls_name.split('_')
74                if 'a' in cls_list: a = 1
75                else:               a = 0
76                if 'b' in cls_list: b = 1
77                else:               b = 0
78                if 'c' in cls_list: c = 1
79                else:               c = 0
80                
81                if mode == 'train':
82                    a_train.append(a)
83                    b_train.append(b)
84                    c_train.append(c)
85                    image_list_train.append(file_name)
86                elif mode == 'val':
87                    a_val.append(a)
88                    b_val.append(b)
89                    c_val.append(c)
90                    image_list_val.append(file_name)
91                else:
92                    print('ERROR')
93
94    num_train = len(image_list_train)
95    num_val = len(image_list_val)
96 
97   # labelの作成 ex) [[0,0,1],[1,1,0],[0,1,0]]
98    cls_id_train = []
99    cls_id_val = []
100    for i in range(num_train):
101        cls_id_train.append([a_train[i], b_train[i], c_train[i]])
102    for i in range(num_val):
103        cls_id_val.append([a_val[i], b_val[i], c_val[i]])
104
105    checkpoint_path = ckpt_dir + '/ckpt-{epoch}-{val_macro_f1:.2f}-{val_loss:.2f}'
106
107    # データセット作成
108    train_ds = create_dataset(image_list_train,
109                              cls_id_train,
110                              total_cnt=num_train,
111                              batch_size=batch_size,
112                              train_flag=True,
113                              )
114    val_ds = create_dataset(image_list_val,
115                            cls_id_val,
116                            total_cnt=num_val,
117                            batch_size=batch_size,
118                            train_flag=False,
119                            )
120
121    # 学習
122    IMG_SHAPE = (img_size, img_size, channels)
123    base_model = tf.keras.applications.MobileNetV2(input_shape=IMG_SHAPE,
124                                                include_top=False,
125                                                weights='imagenet')
126    global_average_layer = tf.keras.layers.GlobalAveragePooling2D()
127    prediction_layer = tf.keras.layers.Dense(n_classes, activation='sigmoid')
128    model = tf.keras.Sequential([base_model,global_average_layer,prediction_layer])
129    model.compile(optimizer=optimizers.SGD(lr=lr, momentum=momentum, nesterov=nesterov),loss='binary_crossentropy', metrics=[macro_f1])
130
131    # validationのlossやaccuracyをモニターしてモデルを保存するタイミングを定義
132    ckpt_cb = tf.keras.callbacks.ModelCheckpoint(checkpoint_path,
133                                                 save_weights_only=True,
134                                                 monitor='val_macro_f1',
135                                                 mode='max',
136                                                 verbose=1)
137    # logをcsvに書き出す
138    csv_logger = tf.keras.callbacks.CSVLogger(ckpt_dir+'/training.csv', separator=',')
139
140    history = model.fit(train_ds,
141                        steps_per_epoch= int(num_train//batch_size),
142                        validation_data=val_ds,
143                        validation_steps= int(num_val//batch_size),
144                        shuffle=True,
145                        epochs=epochs,
146                        callbacks=[ckpt_cb, csv_logger],)
147
148    model.save_weights(ckpt_dir + '/my_checkpoint')

試したこと

バッチサイズを変更したり、DataGeneratorを使った実装に切り替えてみました。
前者は効果がなく、後者についてはエラーが出てきてしまいました（コメントにて詳細を追記しました）

初心者のためどういったことを試せばよいのかもあまり良くわかっていないのですが,,,
アドバイスなどありましたらお知恵をお貸しいただけると非常に助かります。

meg_

2021/05/14 14:59

> バッチサイズを変更したり、 > 前者は効果がなくいくつからいくつにされましたか？そのときのコードはありますか？またメモリ使用量は全く変化ありませんでしたか？

orionza

2021/05/14 15:41

情報不足で大変恐縮です。補足させていただきます。バッチサイズに関しては色々変更してみたのですが、総ステップ数の1/10程度まで学習した時点でのメモリ使用量がほぼ一定でした。もともと8に設定していたのですが、1にしてみても学習は完了しないと思います。ーーーーーーーーーーーーーーーーーーーまた、一応DataGeneratorを用いた方法につきましても追記しておきます。以下のようにdatasetを用意し、train_generatorを model.fitのtrain_dsに置き換えたところ「Error occurred when finalizing GeneratorDataset iterator: Failed precondition: Python interpreter state is not initialized. The process may be terminated. 」というエラーが出ました train_datagen = ImageDataGenerator( rescale=1./255, zoom_range=0.2, horizontal_flip=True, vertical_flip=True, rotation_range=20, width_shift_range=0.2, height_shift_range=0.2, ) val_datagen = ImageDataGenerator(rescale=1./255) train_generator = train_datagen.flow_from_directory( os.path.join(data_dir, "train"), target_size=(img_size, img_size), batch_size=batch_size, # classes=classes, class_mode='categorical')

meg_

2021/05/14 16:11

バッチサイズというのは学習時の学習データ数のことで合っていますか？質問のコードではmodel.fit()でbatch_sizeの指定をしていないようですので確認です。

orionza

2021/05/14 16:30

model.fit()に渡しているのがNumPy配列ではなく tf.data Datasetオブジェクトであるためfit()の引数にbatch_sizeを直接渡していませんが、バッチサイズは8としてdatasetを作成しています。なおdataset作成に用いている関数は、コード上部で定義しているcreate_dataset()です。見当違いの答え方をしてしまっていたら、恐縮ですがご指摘いただければと思います。よろしくお願いいたします。

meg_

2021/05/14 17:28

データはtf.data.Datasetでしたね。失礼しました。