CODEWITHSUNDEEP
dockertensorflowtensorflow-xlaxla
Fred
Fred

Reputation: 666

Tensorflow w/ XLA causing a memory leak

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

Answers (0)

Related Questions