Enormous memory usage after batched forward passes with TensorFlow 2.16.1
As a workaround (I just found), one can replace `model(images).numpy()` with `model.predict(images)`. This tames the memory usage. @Dobiasd,
As you mentioned there was memory usage with tensorflow 2.16.1 while executing the mentioned code. When compared with 2.15, I can see the difference in memory usage. Could you please allow some time to check with the developer and provide the solution for the same. Meanwhile you can try using model.predict(images) where the memory is the same as 2.15.
Kindly find the gist of it [here](https://colab.research.google.com/gist/tilakrayal/56402dc07a59baf4d0b694686851d799/untitled1761.ipynb).

<details><summary>Click to expand for the screenshots!</summary>
2.15:

![Screenshot 2024-03-15 2 26 23 PM](https://github.com/tensorflow/tensorflow/assets/81610181/039ebf94-1c87-4025-b769-b6e4a02e4762)

2.16:

![Screenshot 2024-03-15 2 11 36 PM](https://github.com/tensorflow/tensorflow/assets/81610181/2ad65c65-d6dc-415b-885d-6c994a63d8ca)
 This issue is stale because it has been open for 7 days with no activity. It will be closed if no further activity occurs. Thank you. Thanks for reproducing the effect. I know that using `model.predict(images)` is a workaround for the problematic `model(images)` (wrote it already above). However, I don't feel that this is a reason to close this issue. `model(images)` is still broken. @Dobiasd,
By default tensorflow v2.16 uses keras 3.0 version. I tried to import the keras as **import tf_keras as keras** and the code given execution time with the **model(images)** is similar as tensorflow v2.15. Kindly find the gist of it [here](https://colab.research.google.com/gist/tilakrayal/eb64aa5f382a486d1355a7bc2550309b/untitled1793.ipynb). 

```python
Memory usage after 0 run(s) (in MiB): 1279.699
Memory usage after 1 run(s) (in MiB): 1518.086
Memory usage after 2 run(s) (in MiB): 1518.867
Memory usage after 3 run(s) (in MiB): 1488.008
Memory usage after 4 run(s) (in MiB): 1518.637
Memory usage after 5 run(s) (in MiB): 1518.578
Memory usage after 6 run(s) (in MiB): 1518.777
Memory usage after 7 run(s) (in MiB): 1518.613
Memory usage after 8 run(s) (in MiB): 1503.199
Memory usage after 9 run(s) (in MiB): 1518.297
```
Thank you! Thanks! Does this mean this issue should be moved from the `tensorflow/tensorflow` repo to the `keras-team/keras` repo? Btw., `model.predict(images)` also leaks a bit of memory:

```python3
import numpy as np
import psutil
import tensorflow as tf

model = tf.keras.applications.ResNet152V2()
for run in range(10000):
    images = np.zeros([run % 100 + 1, 224, 224, 3], dtype=np.uint8)
    memory_usage_in_MiB = psutil.Process().memory_info().rss / (1024 * 1024)
    print(f"Memory usage after {run} run(s) (in MiB): {memory_usage_in_MiB:.3f}", flush=True)
    model.predict(images, verbose=0)
```

```
Memory usage after 0 run(s) (in MiB): 783.898
Memory usage after 1 run(s) (in MiB): 843.965
Memory usage after 2 run(s) (in MiB): 905.809
Memory usage after 3 run(s) (in MiB): 926.613
Memory usage after 4 run(s) (in MiB): 929.820
Memory usage after 5 run(s) (in MiB): 988.938
Memory usage after 6 run(s) (in MiB): 1022.199
Memory usage after 7 run(s) (in MiB): 1030.844
Memory usage after 8 run(s) (in MiB): 1035.547
Memory usage after 9 run(s) (in MiB): 1075.227
Memory usage after 10 run(s) (in MiB): 1080.125
Memory usage after 11 run(s) (in MiB): 1085.094
Memory usage after 12 run(s) (in MiB): 1064.848
Memory usage after 13 run(s) (in MiB): 1082.512
Memory usage after 14 run(s) (in MiB): 1094.887
Memory usage after 15 run(s) (in MiB): 1102.629
Memory usage after 16 run(s) (in MiB): 1122.195
Memory usage after 17 run(s) (in MiB): 1101.715
Memory usage after 18 run(s) (in MiB): 1118.234
Memory usage after 19 run(s) (in MiB): 1120.234
Memory usage after 20 run(s) (in MiB): 1136.480
...
Memory usage after 409 run(s) (in MiB): 1254.375
Memory usage after 410 run(s) (in MiB): 1260.383
Memory usage after 713 run(s) (in MiB): 1297.750
Memory usage after 714 run(s) (in MiB): 1300.977
Memory usage after 715 run(s) (in MiB): 1303.609
...
Memory usage after 3117 run(s) (in MiB): 1329.902
Memory usage after 3118 run(s) (in MiB): 1357.523
...
```

([Dockerfile](https://gist.github.com/Dobiasd/d37045cf3f4413457ff2427edeb69931) to reproduce) @Dobiasd,
Looks like this is also happening with the tensorflow v.2.16 & keras3.0. With the tensorflow v2.15 and tf-keras I was not able to find any issue with the memory. Kindly please raise the request in keras-team/keras repo. Thank you! Thanks. As suggested, I've opened [this issue](https://github.com/keras-team/keras/issues/19500) in `keras-team/keras`. The conclusion so far of the discussion in the Keras repo [is](https://github.com/keras-team/keras/issues/19500#issuecomment-2053705772), that it's a problem with TensorFlow, not with Keras. This is what I get when running @Dobiasd's dockerfile image on my machine:
```
Memory usage after 0 run(s) (in MiB): 993.359
Memory usage after 1 run(s) (in MiB): 1021.312
Memory usage after 2 run(s) (in MiB): 1115.410
Memory usage after 3 run(s) (in MiB): 1166.316
Memory usage after 4 run(s) (in MiB): 1233.500
Memory usage after 5 run(s) (in MiB): 1271.699
Memory usage after 6 run(s) (in MiB): 1360.941
Memory usage after 7 run(s) (in MiB): 1388.930
Memory usage after 8 run(s) (in MiB): 1464.305
Memory usage after 9 run(s) (in MiB): 1591.188
...
Memory usage after 37 run(s) (in MiB): 3212.566
Memory usage after 38 run(s) (in MiB): 3280.664
Memory usage after 39 run(s) (in MiB): 3320.461
Memory usage after 40 run(s) (in MiB): 2003.023   <--- big drop
Memory usage after 41 run(s) (in MiB): 2137.840
Memory usage after 42 run(s) (in MiB): 2198.230
...
Memory usage after 67 run(s) (in MiB): 4361.066
Memory usage after 68 run(s) (in MiB): 4424.414
Memory usage after 69 run(s) (in MiB): 4457.262
Memory usage after 70 run(s) (in MiB): 2415.262   <--- big drop
Memory usage after 71 run(s) (in MiB): 2599.285
Memory usage after 72 run(s) (in MiB): 2723.152
...
```
This pattern seems to repeat: slowly growing, then a drop.  It never seems to monotonically increase: only up to some limit before dropping again.

I'm not sure it's a leak, since it stays below a certain (perhaps machine-specific) limit before it drops down.  Outside of docker, running locally on my machine, I see more stable results on CPU:
```
Memory usage after 0 run(s) (in MiB): 988.746
Memory usage after 1 run(s) (in MiB): 1042.246
Memory usage after 2 run(s) (in MiB): 1050.246
Memory usage after 3 run(s) (in MiB): 1046.680
Memory usage after 4 run(s) (in MiB): 1050.180
Memory usage after 5 run(s) (in MiB): 1052.312
Memory usage after 6 run(s) (in MiB): 1048.418
Memory usage after 7 run(s) (in MiB): 1052.418
Memory usage after 8 run(s) (in MiB): 1049.156
Memory usage after 9 run(s) (in MiB): 1054.156
Memory usage after 10 run(s) (in MiB): 1054.156
...
Memory usage after 96 run(s) (in MiB): 1074.281
Memory usage after 97 run(s) (in MiB): 1072.508
Memory usage after 98 run(s) (in MiB): 1076.008
Memory usage after 99 run(s) (in MiB): 1070.641
```
which is very similar to @fchollet's [colab](https://colab.research.google.com/drive/1FesdhTGp65FmL1ZqNj28yJ8yqSPnG3Mt?usp=sharing).

I don't even think it's a keras 3 vs legacy keras issue, since using [legacy `tf-keras`](https://keras.io/getting_started/#tensorflow--keras-2-backwards-compatibility) leads to the same memory profile for me both in docker and my local machine.

My guess is its just TensorFlow's [BFC allocator](https://github.com/google/tsl/blob/556552d6c0c97c7fdb8689cdc75768f0cfa0be7c/tsl/framework/bfc_allocator.h#L51) slowly growing.  Something might have changed in TF 2.16 related to memory allocations that reduces memory re-use before the allocator consolidates freed blocks.  If you have a lot of memory available, it might grow more before doing this on your particular machine. Thanks for testing!

> If you have a lot of memory available, it might grow more before doing this on your particular machine.

That's very unfortunate if you're running something like that in a Kubernetes pod with a (cgroups-managed) memory limit, and TensorFlow uses the memory as if it were the only process on the whole node. At least that's how I ran into this problem, i.e., my K8s pods were constantly OOMKilled. > Thanks for testing!
> 
> > If you have a lot of memory available, it might grow more before doing this on your particular machine.
> 
> That's very unfortunate if you're running something like that in a Kubernetes pod with a (cgroups-managed) memory limit, and TensorFlow uses the memory as if it were the only process on the whole node. At least that's how I ran into this problem, i.e., my K8s pods were constantly OOMKilled.

Sorry, I made a mistake - using legacy keras _does_ fix the issue.  Setting the environment variable:
```
export TF_USE_LEGACY_KERAS=1
```
_does_ restore the original memory profile.  So keras 3 is a contributing factor.  I think it has to do with tracking information for training (since this is a trainable model).  Note that the initial jump in memory use is just the model loading the weights for the first time, since you report your first memory usage _before_ your first predict call.

I still don't think it's a "leak".  Modifying your script a bit to use `tracemalloc` and recording memory use _after_ the first predict iteration:
```python
import numpy as np
import psutil
import tensorflow as tf
import tracemalloc
import gc

def main():
  model = tf.keras.applications.ResNet152V2()
  images = np.zeros([20, 224, 224, 3], dtype=np.uint8)

  for run in range(20):
    model(images)
    gc.collect()
    memory_usage_in_MiB = psutil.Process().memory_info().rss / (1024 * 1024)
    print(f"Memory usage after {run} run(s) (in MiB): {memory_usage_in_MiB:.3f}", flush=True)

tracemalloc.start()
main()
snapshot = tracemalloc.take_snapshot()
top_stats = snapshot.statistics('lineno')
print("[ Top 10 ]")
for stat in top_stats[:10]:
    print(stat)
```
I get the following with Keras 3 by exporting `TF_USE_LEGACY_KERAS=0`:
```
Memory usage after 0 run(s) (in MiB): 6094.586
Memory usage after 1 run(s) (in MiB): 7081.043
Memory usage after 2 run(s) (in MiB): 7133.891
Memory usage after 3 run(s) (in MiB): 7132.301
Memory usage after 4 run(s) (in MiB): 7139.891
Memory usage after 5 run(s) (in MiB): 7219.988
Memory usage after 6 run(s) (in MiB): 7156.012
Memory usage after 7 run(s) (in MiB): 7247.742
Memory usage after 8 run(s) (in MiB): 7218.574
Memory usage after 9 run(s) (in MiB): 7219.082
Memory usage after 10 run(s) (in MiB): 7222.090
Memory usage after 11 run(s) (in MiB): 7245.625
Memory usage after 12 run(s) (in MiB): 7245.625
Memory usage after 13 run(s) (in MiB): 7259.473
Memory usage after 14 run(s) (in MiB): 7266.531
Memory usage after 15 run(s) (in MiB): 7266.871
Memory usage after 16 run(s) (in MiB): 7242.562
Memory usage after 17 run(s) (in MiB): 7243.641
Memory usage after 18 run(s) (in MiB): 7194.695
Memory usage after 19 run(s) (in MiB): 7247.301
[ Top 10 ]
/tmp/venv/lib/python3.10/site-packages/tensorflow/python/framework/constant_op.py:96: size=188 MiB, count=1272, average=151 KiB
/tmp/venv/lib/python3.10/site-packages/keras/src/utils/tracking.py:65: size=634 KiB, count=3405, average=191 B
/tmp/venv/lib/python3.10/site-packages/tensorflow/python/trackable/autotrackable.py:70: size=444 KiB, count=1144, average=398 B
/tmp/venv/lib/python3.10/site-packages/keras/src/utils/tracking.py:139: size=339 KiB, count=6676, average=52 B
/tmp/venv/lib/python3.10/site-packages/tensorflow/python/ops/resource_variable_ops.py:2026: size=318 KiB, count=1631, average=200 B
/tmp/venv/lib/python3.10/site-packages/keras/src/utils/tracking.py:95: size=235 KiB, count=3339, average=72 B
/tmp/venv/lib/python3.10/site-packages/tensorflow/python/trackable/base.py:361: size=229 KiB, count=1085, average=216 B
<frozen importlib._bootstrap_external>:672: size=226 KiB, count=2131, average=109 B
/usr/local/google/home/cantonios/.pyenv/versions/3.10.12/lib/python3.10/abc.py:123: size=224 KiB, count=2834, average=81 B
/tmp/venv/lib/python3.10/site-packages/keras/src/utils/tracking.py:100: size=191 KiB, count=874, average=224 B
```
The `constant_op` is allocating the weight and shape tensors.  Then there's a bunch of "tracking" allocations.

If I use legacy keras by exporting `TF_USE_LEGACY_KERAS=1`, I get:
```
Memory usage after 0 run(s) (in MiB): 1341.215
Memory usage after 1 run(s) (in MiB): 1313.246
Memory usage after 2 run(s) (in MiB): 1313.789
Memory usage after 3 run(s) (in MiB): 1497.645
Memory usage after 4 run(s) (in MiB): 1618.266
Memory usage after 5 run(s) (in MiB): 1624.480
Memory usage after 6 run(s) (in MiB): 1625.363
Memory usage after 7 run(s) (in MiB): 1625.152
Memory usage after 8 run(s) (in MiB): 1625.297
Memory usage after 9 run(s) (in MiB): 1625.344
Memory usage after 10 run(s) (in MiB): 1624.699
Memory usage after 11 run(s) (in MiB): 1622.949
Memory usage after 12 run(s) (in MiB): 1625.496
Memory usage after 13 run(s) (in MiB): 1624.801
Memory usage after 14 run(s) (in MiB): 1625.109
Memory usage after 15 run(s) (in MiB): 1625.328
Memory usage after 16 run(s) (in MiB): 1624.387
Memory usage after 17 run(s) (in MiB): 1625.062
Memory usage after 18 run(s) (in MiB): 1625.141
Memory usage after 19 run(s) (in MiB): 1625.152
[ Top 10 ]
/tmp/venv/lib/python3.10/site-packages/tensorflow/python/framework/constant_op.py:96: size=107 MiB, count=1311, average=83.8 KiB
/usr/local/google/home/cantonios/.pyenv/versions/3.10.12/lib/python3.10/abc.py:123: size=982 KiB, count=10222, average=98 B
/tmp/venv/lib/python3.10/site-packages/tf_keras/src/engine/base_layer.py:3223: size=671 KiB, count=311, average=2208 B
/tmp/venv/lib/python3.10/site-packages/tf_keras/src/engine/base_layer.py:375: size=330 KiB, count=4532, average=75 B
/tmp/venv/lib/python3.10/site-packages/tensorflow/python/ops/resource_variable_ops.py:2026: size=319 KiB, count=1634, average=200 B
/tmp/venv/lib/python3.10/site-packages/tf_keras/src/engine/base_layer.py:3134: size=314 KiB, count=822, average=392 B
/tmp/venv/lib/python3.10/site-packages/tensorflow/python/trackable/base.py:361: size=269 KiB, count=1274, average=216 B
/tmp/venv/lib/python3.10/site-packages/tf_keras/src/utils/object_identity.py:135: size=263 KiB, count=5608, average=48 B
/tmp/venv/lib/python3.10/site-packages/tf_keras/src/utils/object_identity.py:141: size=254 KiB, count=313, average=830 B
/tmp/venv/lib/python3.10/site-packages/tensorflow/python/framework/tensor_shape.py:830: size=158 KiB, count=2817, average=57 B
```

The main difference to me looks like mostly `tracking`-related allocations in keras 3.  Every keras layer is `Trackable`, so my guess is the increased memory usage lies there.  I believe converting to a `tf.function` essentially disables this trainable tracking, so you get the similar results as your `model.predict(...)` call.

I still don't get unbounded growth though, just vastly increased memory usage.

---

As for what to do: can you live with legacy keras and set the environment variable `TF_USE_LEGACY_KERAS=1`?  Can you freeze your model into an inference-only TF model and use that?  Otherwise I think the keras team might be better at debugging. Actually, rewriting your reproducer in terms of `keras` alone and setting  it to use JAX as the backend, I get similar high memory usage:
```python
import os
os.environ["KERAS_BACKEND"] = "jax"

import numpy as np
import psutil
import keras
import tracemalloc
import gc

def main():
  model = keras.applications.ResNet152V2()
  images = np.zeros([20, 224, 224, 3], dtype=np.uint8)

  for run in range(20):
    model(images)
    # model.predict(images)
    gc.collect()
    memory_usage_in_MiB = psutil.Process().memory_info().rss / (1024 * 1024)
    print(f"Memory usage after {run} run(s) (in MiB): {memory_usage_in_MiB:.3f}", flush=True)

tracemalloc.start()
main()
snapshot = tracemalloc.take_snapshot()
top_stats = snapshot.statistics('lineno')
print("[ Top 10 ]")
for stat in top_stats[:10]:
    print(stat)

```
```
An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.
Memory usage after 0 run(s) (in MiB): 5612.281
Memory usage after 1 run(s) (in MiB): 5612.195
Memory usage after 2 run(s) (in MiB): 5610.016
Memory usage after 3 run(s) (in MiB): 5611.594
Memory usage after 4 run(s) (in MiB): 5609.293
Memory usage after 5 run(s) (in MiB): 5611.176
Memory usage after 6 run(s) (in MiB): 5611.828
Memory usage after 7 run(s) (in MiB): 5609.570
Memory usage after 8 run(s) (in MiB): 5611.676
Memory usage after 9 run(s) (in MiB): 5612.074
Memory usage after 10 run(s) (in MiB): 5611.574
Memory usage after 11 run(s) (in MiB): 5610.383
Memory usage after 12 run(s) (in MiB): 5613.406
Memory usage after 13 run(s) (in MiB): 5610.793
Memory usage after 14 run(s) (in MiB): 5610.465
Memory usage after 15 run(s) (in MiB): 5610.535
Memory usage after 16 run(s) (in MiB): 5611.539
Memory usage after 17 run(s) (in MiB): 5612.266
Memory usage after 18 run(s) (in MiB): 5610.262
Memory usage after 19 run(s) (in MiB): 5611.531
[ Top 10 ]
/tmp/venv/lib/python3.10/site-packages/h5py/_hl/dataset.py:1057: size=230 MiB, count=2238, average=105 KiB
/tmp/venv/lib/python3.10/site-packages/jax/_src/dispatch.py:333: size=1014 KiB, count=1632, average=636 B
/tmp/venv/lib/python3.10/site-packages/keras/src/utils/tracking.py:65: size=634 KiB, count=3405, average=191 B
<string>:1: size=480 KiB, count=6086, average=81 B
/tmp/venv/lib/python3.10/site-packages/jax/_src/source_info_util.py:222: size=369 KiB, count=694, average=544 B
/tmp/venv/lib/python3.10/site-packages/keras/src/layers/layer.py:1361: size=341 KiB, count=1146, average=305 B
/tmp/venv/lib/python3.10/site-packages/jax/_src/config.py:208: size=312 KiB, count=1816, average=176 B
<frozen importlib._bootstrap_external>:672: size=297 KiB, count=3060, average=99 B
/tmp/venv/lib/python3.10/site-packages/keras/src/utils/tracking.py:139: size=250 KiB, count=4930, average=52 B
<string>:3: size=244 KiB, count=4343, average=58 B
```
Changing the model to use `model.predict(...)` instead yields much lower memory usage:
```
An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.
1/1 ━━━━━━━━━━━━━━━━━━━━ 6s 6s/step
Memory usage after 0 run(s) (in MiB): 656.660
1/1 ━━━━━━━━━━━━━━━━━━━━ 4s 4s/step
Memory usage after 1 run(s) (in MiB): 693.199
1/1 ━━━━━━━━━━━━━━━━━━━━ 1s 508ms/step
Memory usage after 2 run(s) (in MiB): 697.055
1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 492ms/step
Memory usage after 3 run(s) (in MiB): 697.496
1/1 ━━━━━━━━━━━━━━━━━━━━ 1s 503ms/step
Memory usage after 4 run(s) (in MiB): 697.465
1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 499ms/step
Memory usage after 5 run(s) (in MiB): 698.285
1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 498ms/step
Memory usage after 6 run(s) (in MiB): 698.219
1/1 ━━━━━━━━━━━━━━━━━━━━ 1s 508ms/step
Memory usage after 7 run(s) (in MiB): 698.145
1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 499ms/step
Memory usage after 8 run(s) (in MiB): 698.293
1/1 ━━━━━━━━━━━━━━━━━━━━ 1s 501ms/step
Memory usage after 9 run(s) (in MiB): 698.617
1/1 ━━━━━━━━━━━━━━━━━━━━ 1s 502ms/step
Memory usage after 10 run(s) (in MiB): 698.617
1/1 ━━━━━━━━━━━━━━━━━━━━ 1s 505ms/step
Memory usage after 11 run(s) (in MiB): 698.055
1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 500ms/step
Memory usage after 12 run(s) (in MiB): 698.348
1/1 ━━━━━━━━━━━━━━━━━━━━ 1s 501ms/step
Memory usage after 13 run(s) (in MiB): 698.781
1/1 ━━━━━━━━━━━━━━━━━━━━ 1s 501ms/step
Memory usage after 14 run(s) (in MiB): 698.602
1/1 ━━━━━━━━━━━━━━━━━━━━ 1s 500ms/step
Memory usage after 15 run(s) (in MiB): 698.922
1/1 ━━━━━━━━━━━━━━━━━━━━ 1s 505ms/step
Memory usage after 16 run(s) (in MiB): 698.758
1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 492ms/step
Memory usage after 17 run(s) (in MiB): 698.801
1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 496ms/step
Memory usage after 18 run(s) (in MiB): 698.977
1/1 ━━━━━━━━━━━━━━━━━━━━ 1s 503ms/step
Memory usage after 19 run(s) (in MiB): 698.875
[ Top 10 ]
/tmp/venv/lib/python3.10/site-packages/h5py/_hl/dataset.py:1057: size=230 MiB, count=1335, average=176 KiB
/tmp/venv/lib/python3.10/site-packages/jax/_src/source_info_util.py:222: size=1832 KiB, count=3449, average=544 B
<string>:1: size=892 KiB, count=11325, average=81 B
/tmp/venv/lib/python3.10/site-packages/jax/_src/dispatch.py:333: size=639 KiB, count=1029, average=636 B
/tmp/venv/lib/python3.10/site-packages/keras/src/utils/tracking.py:65: size=634 KiB, count=3405, average=191 B
/tmp/venv/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py:1932: size=453 KiB, count=10520, average=44 B
/tmp/venv/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py:1931: size=452 KiB, count=10496, average=44 B
/tmp/venv/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py:1912: size=404 KiB, count=7370, average=56 B
/tmp/venv/lib/python3.10/site-packages/keras/src/layers/layer.py:1361: size=341 KiB, count=1146, average=305 B
<frozen importlib._bootstrap_external>:672: size=297 KiB, count=3060, average=99 B
```
which seems to confirm that all the memory is the keras model storing training information.

So I'm pretty confident this isn't a TensorFlow bug, and it may or may not be a Keras bug - depending on if that memory use for training is actually needed by keras. > As for what to do: can you live with legacy keras and set the environment variable `TF_USE_LEGACY_KERAS=1`? Can you freeze your model into an inference-only TF model and use that? Otherwise I think the keras team might be better at debugging.

Can conform, in my original minimal example, `TF_USE_LEGACY_KERAS=1` [helps](https://gist.github.com/Dobiasd/c13cfd6697010125e03143021189ce85). I'll try it in PROD next week. :+1: Thanks! Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/63701">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/63701">No</a>

The following minimal example reproduces the effect. (No GPU, just CPU.)

```python3
import numpy as np
import psutil
import tensorflow as tf

model = tf.keras.applications.ResNet152V2()

images = np.zeros([20, 224, 224, 3], dtype=np.uint8)

for run in range(10):
    memory_usage_in_MiB = psutil.Process().memory_info().rss / (1024 * 1024)
    print(f"Memory usage after {run} run(s) (in MiB): {memory_usage_in_MiB:.3f}", flush=True)
    model(images)
```

```
Memory usage after 0 run(s) (in MiB): 790.715
Memory usage after 1 run(s) (in MiB): 5969.941
Memory usage after 2 run(s) (in MiB): 7112.559
Memory usage after 3 run(s) (in MiB): 7787.625
Memory usage after 4 run(s) (in MiB): 7827.727
Memory usage after 5 run(s) (in MiB): 6785.316
Memory usage after 6 run(s) (in MiB): 5971.680
Memory usage after 7 run(s) (in MiB): 6615.105
Memory usage after 8 run(s) (in MiB): 7151.496
Memory usage after 9 run(s) (in MiB): 6679.250
```

([Dockerfile](https://gist.github.com/Dobiasd/4e6747f3c057b939cb2672d142f612fc) to reproduce)

And it's not just `tf.keras.applications.ResNet152V2()`. It also happens (for example) with `tf.keras.applications.inception_resnet_v2.InceptionResNetV2` and `tf.keras.applications.inception_v3.InceptionV3`. And it also happens when using Python `3.12.2` instead of `3.11.8`.

With TensorFlow `2.15.1` (instead of `2.16.1`), however, the memory usage does not explode:

```
Memory usage after 0 run(s) (in MiB): 1038.891
Memory usage after 1 run(s) (in MiB): 1154.062
Memory usage after 2 run(s) (in MiB): 1154.508
Memory usage after 3 run(s) (in MiB): 1154.684
Memory usage after 4 run(s) (in MiB): 1215.840
Memory usage after 5 run(s) (in MiB): 1154.664
Memory usage after 6 run(s) (in MiB): 1154.816
Memory usage after 7 run(s) (in MiB): 1155.059
Memory usage after 8 run(s) (in MiB): 1154.715
Memory usage after 9 run(s) (in MiB): 1155.367
```

([Dockerfile](https://gist.github.com/Dobiasd/0acf434d521cabf7fe3102b551bf36b0) to reproduce)