Reputation: 666
I'm training the EfficientDet neural network with Tensorflow 2.9 in a Docker container.
Without XLA compilation, everything runs fine. With XLA, I'm getting a 4x performance boost! However, there is a memory leak that increases RAM usage until things crash.
I enabled XLA by adding an environment variable to my Dockerfile like so:
ENV TF_XLA_FLAGS="--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit"
This is happening on EfficientDet but not EfficientNetV2. The leak does not affect GPU memory, only RAM.
Any ideas?
Upvotes: 1
Views: 259