问题类型

Bug

来源

二进制文件

Tensorflow版本

2.10,2.9也进行了测试

自定义代码

无

OS平台和发行版

Ubuntu Linux 22.04和Google Colab测试过

移动设备

• 无响应*

Python版本

3.9

Bazel版本

• 无响应*

GCC/编译器版本

• 无响应*

CUDA/cuDNN版本

11.8

GPU型号和内存

• 无响应*

当前行为？

When dataset is generated via `Dataset.from_generator`, samples that have already been used are not removed from GPU memory (or at least this is what I presume), leading to an increased use of VRAM and finally an OOM error as the training proceeds.

This is an important bug, since very often (as in this case) someone might use `from_generator` to perform augmentation etc, and It can be easily seen by examining the `GPU_mem_usage`, that the memory usage indeed grows. For some reason, Colab allocates memory straight away to ~8 GB, so the growth is only noticeable after around 300-500 epochs. The problem is way worse when training on real, larger datasets.

For a simple summary, just view the graphs on the end of the notebook.

Notice (in the notebook), that even though total memory taken on the GPU grows, TensorFlow "thinks" it is consuming basically constant amount of memory, which points into the direction of a memory leak.

Also, it takes around 900 epochs to fill almost the whole memory (16 GB in case of Colab). Epochs consist of 1024 images of shape (64, 64, 3), float32 dtype, giving 1024*3*64*64*900*4 bytes allocated in total, which is around 45 GB, so I presume not all memory is leaked, OR leaking data is not the cause of the problem.

Notice, however, that MobileNetV3 (which I have used in this example) is roughly 18MB in size, and 18MB*900 epochs is basically the aforementioned 16 GB. This could possibly mean that the model's state is leaked somehow, but I have yet to test this hypothesis (for example by using a larger model and checking when the leak happens, i.e. if `num_epochs_until_crash*model_size==gpu_vram_capacity`).

Also, I have tested multiple scenarios locally, and this seems NOT to happen without the `os.environ['TF_GPU_ALLOCATOR'] = 'cuda_malloc_async'` set, however I have no relevant logs to prove it, since when the flag is absent, whole GPU memory is allocated, and I have no means of monitoring it relevantly, since it doesn't change. I have, however, performed 4 tests and none resulted in OOM error after 1500 epochs, while locally they would usually throw OOM after 400 epochs (RTX 3060Ti).

重现问题的独立代码

I have provided a minimal example on google colab, can be viewed here:
https://colab.research.google.com/drive/1-ANVp8KF9irKvqdR390QNlop-pPhGMrU?usp=sharing
I think the example is self-explanatory.

相关日志输出

• 无响应*