TensorFlowで大量の画像や巨大な配列をpredictするとメモリーリークが生じることが多い。
model.fitでは問題ないのに、model.predictになるとOOMになってしまう報告が散見する。
model.predict leads to oom issues, but model.fit does not:
my training model is a modified version of U-Net with the goal being to add a filter to a 512x1808 image. the issue I'm having is that I keep getting out of m...
stackoverflow.com
OOM when calling model.predict()
Hi everyone, I’m having an issue with model.predict() causing OOM errors. Strangely, this doesn’t happen while training, only while predicting on the val datas...
discuss.tensorflow.org
out of memory when using model.predict() · Issue #5337 · keras-team/keras
hi~ I am now using keras to build my network. the training is normal. but when I use the model.predict() to predict the results, it happend to a error : out of ...
github.com
根本的な解決ではないが、predictを明示的にCPUで行うようにしたところOOMは起きなくなった。
TensorFlowでは、tf.deviceを使うことで処理を行うデバイスを指定することができる。
tf.device | TensorFlow v2.11.0
Specifies the device for ops created/executed in this context.
www.tensorflow.org
そのため、model.fitはGPUで、model.predictはCPUで行うことが可能となる。
以下のように処理を行うことで、OOMを回避することができた。
strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
model = ...
model.fit(...)
with tf.device("/cpu:0"):
pred = model.predict(...)
コメント