Description
System information.
- Have I written custom code (as opposed to using a stock example script provided in Keras):
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 22.04.1 LTS (GNU/Linux 5.16.10 x86_64)
- TensorFlow installed from (source or binary): source
- TensorFlow version (use command below): 2.11.0
- Python version: 3.10.6
- Bazel version (if compiling from source):
- GPU model and memory: NVIDIA RTX A5000, 24GB
- Exact command to reproduce:
Describe the problem.
Memory usage steadily increases when using tf.Model and tf.Model.fit() in a loop, and leads to Out Of Memory exception saturating the memory eventually. clear_session() does not help. The same code with TF version == 2.9.2 has an almost constant memory usage instead, and works as expected.
Describe the problem clearly here. Be sure to convey here why it's a bug in Keras or why the requested feature is needed.
Describe the current behavior.
Memory usage steadily increases when using tf.Model and tf.Model.fit() in a loop, and leads to Out Of Memory exception saturating the memory eventually.
Describe the expected behavior.
The memory usage remains almost the same.
Standalone code to reproduce the issue.
import tensorflow as tf
import time
class MyModel(tf.keras.Model):
def __init__(self):
super().__init__()
self.dense1 = tf.keras.layers.Dense(1000, activation=tf.nn.relu)
self.dense2 = tf.keras.layers.Dense(10000, activation=tf.nn.softmax)
self.dense3 = tf.keras.layers.Dense(10000, activation=tf.nn.softmax)
self.dense4 = tf.keras.layers.Dense(1000, activation=tf.nn.softmax)
def call(self, inputs):
x = self.dense1(inputs)
x = self.dense2(x)
x = self.dense3(x)
x = self.dense4(x)
return x
for r in range(0, 10000):
model = MyModel()
ds = tf.data.Dataset.from_tensor_slices((tf.random.uniform((64*4, 1000)), tf.ones((64*4))))
model.compile(optimizer='sgd', loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True))
model.fit(ds.batch(64))
tf.keras.backend.clear_session()
time.sleep(3)
print("round: ", r)
Provide a reproducible test case that is the bare minimum necessary to generate
the problem. If possible, please share a link to Colab/Jupyter/any notebook.
Source code / logs.
Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached. Try to provide a reproducible test case that is the bare minimum necessary to generate the problem.