google ColaboratoryでシステムRAMがクラッシュしないようメモリ解放したい

Question

実現したいこと

「ランタイムを再起動」をせずに繰り返し色々なモデルを学習させたい。

前提

画像分類タスクを行いたく、google Colaboratoryにてtensorflowを使ってニューラルネットワークを実装しようとしています。
使用言語はPythonです。

コードを実行していくとシステムRAMが激しく消費され、すぐに上限の12.7GBを超えてセッションがクラッシュしてしまいます。
1回分の学習はできますが、パラメータ探索のために、「ランタイムを再起動」で毎回リセットせずとも連続の学習ができるように、メモリの節約や解放の手段が必要だと考えています。

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

システムRAMが急激に消費されるポイントがいくつかあります。それぞれでメモリの節約や解放の方法があれば教えて頂きたいです。

①学習データをマウントしているgoogle Driveから読み込む
ここで2.8GB使っています。この学習データはnumpyのndarray配列であり予めnpz形式で保存しているものです。数千枚の130KB程度の画像が入っていてX_train.npzのサイズは400MBです。400MBの読み込みで7倍メモリを使うものなのか理解できず…。

import numpy as np
X_train = np.load('/content/drive/MyDrive/X_train.npz')['arr_0']

②定義したモデルで学習を行う
モデルの中身で変動しますが3GB前後使います。全結合層や畳み込み層が3層くらいでパラメータのサイズは400KB～10MBです。一回の学習でこれくらい使うのはまあ良いとして、modelを再定義し最初から学習しても消費した分は戻らず増加する一方です。

epochs = 10
batch_size = 32
result = model.fit(tr_x,tr_y, epochs=epochs, batch_size=batch_size, validation_data=(va_x,va_y))

試したこと

②については不要になった段階でmodelやresultを削除すれば良いのでは？と思い、ネットで検索し使えそうな以下のコードを実行してみましたが、システムRAMは減りませんでした。

del model, result
keras.backend.clear_session()
import gc
gc.collect()

補足情報（FW/ツールのバージョンなど）

google Colabは無料版で、ハードウェアアクセラレータはT4 GPUを選択しています。

Comments

[退会済みユーザー]

X_train.npz は圧縮されているのではないでしょうか。だから実行時に展開されてメモリを圧迫するのだと思います。
gc（Gabage Colection)は、ローカルで実行する場合と異なり、クラウドのGoogle Colaboではgcを実行したからといって確実に開放されるとは限りません。

個人的に、やりたいことがあるなら、素直にProプランに課金してメモリを積んだ方が精神衛生的によいのではないかなと思います。
（「なんとかタダでメモリを節約する方法を見つけたい」という、手段と目的が逆転している状況なら何も言いませんが。。。）

[退会済みユーザー]

試したことがないので投げやりなコメントですが、コラボラトリーはスワップできるんですかね？
あとは妥協してモデルを小さくするとかコンパクトなモバイルネットに頼るとか。

[debon]

X_train.npzは圧縮されてる訳ではないようです。元のjpgファイル群が入ったフォルダのサイズも380MBでしたので。gcの事は承知しました。
小さいモデルにしたり有料版にするとかはアリなんですが、簡単に数GB飛ぶ状況ではすぐ同じ問題にぶつかりそうなので、先にこの問題を解決したいと思っています。

[jbpb0]

> X_train.npzのサイズは400MB

> 元のjpgファイル群が入ったフォルダのサイズも380MB

jpegファイルはたいてい圧縮されてるので、npzファイルのサイズがそれに近いなら、npzファイルも圧縮されてる可能性があります

> 400MBの読み込みで7倍メモリを使う

X_train = np.load(...
の下に下記を追加して実行したら、何て表示されますか？

print(X_train.shape)
print(X_train.size)
print(X_train.dtype)
print(X_train.itemsize)
print(X_train.nbytes)

また、
X_train.nbytes (= X_train.size * X_train.itemsize)
は、2.8GBに近い数値ではありませんでしょうか？

[debon]

皆さんコメントありがとうございます。上記コードを実行したところ以下の通りでした。
(3000, 200, 200, 3)
360000000
float64
8
2880000000
これにより、X_trainが実際に2.8GB程度あるのが分かりました。①では誤りでメモリ消費が増大しているわけではないと納得しました。

[退会済みユーザー]

古いtf（v1.4位）だと訳の分からないコードを捏ね回さないとメモリリークしてましたが、最近のTFはどうなんでしょうね。
当時はメモリの最大量を縛るとかそういうパラメータを捏ね回していたいた気がします。

[退会済みユーザー]

きようびのTFで使えそうな記事がmediumにありました。
https://medium.com/dive-into-ml-ai/dealing-with-memory-leak-issue-in-keras-model-training-e703907a6501

Answer 1

上記fourteenlengthさんのコメントの記事を参考に②について解決の見込みが立ちました。ありがとうございます。

result = model.fit(tr_x,tr_y, epochs=epochs, batch_size=batch_size, validation_data=(va_x,va_y))

におけるtr_x,tr_y,va_x,va_yはX_trainから作ったnumpyのndarray配列でした。これらを

tr_x = tf.convert_to_tensor(tr_x)
tr_y = tf.convert_to_tensor(tr_y)
va_x = tf.convert_to_tensor(va_x)
va_y = tf.convert_to_tensor(va_y)

でtensorflowのテンソルに変換してからmodel.fit()に入力することで繰り返しmodel.fit()を実行しても際限なくメモリを消費することがなくなりました。

なお上記「試したこと」を再度試すと今回は100%ではないですがメモリがちゃんと解放されました。ただ効果が見られなかった前回との違いが分からず…。コードのシンプルさという観点からもテンソルへの変換の方で対応したいと思います。