Orin NX tensorflow - high memory use

Overview

I am trying to run a tensorflow model to do predictions on the Orin NX 8GB platform. I have Nvidia’s tensorflow==2.12.0+nv23.6 installed so it is using the integrated GPU. This has been confirmed via tegrastats showing non-zero GPU usage when the model runs.

As a baseline, it seems like Linux and normal processes consume about 1.4GB of memory.

After the first call to model.predict(), I see memory use jump to 6GB and then it never seems to reduce. I suspect the extra memory is from the GPU since it doesn’t track to a specific process (at least exactly). See below htop screenshot :

38.5% MEM is ~2.8GB. With the 1.4GB baseline I measured, that leaves about 2GB unaccounted for from the 6.21GB shown utilized above.

My Setup

  • An Orin NX 8GB with L4T R35.4.1 (the minimal filesystem with no desktop)
  • Python 3.8 that L4T provides
  • tensorflow==2.12.0+nv23.6 from Nvidia
  • libcudnn.so.8.6.0
  • model is loaded from a .hdf5 file if that matters

Questions

  1. Is this a known issue with Nvidia’s tensorflow implementation for Jetson…that the GPU does not free up memory once allocated? If so, is there a way to limit the memory or tell it to behave differently?

  2. From other posts, I’ve seen recommendations to use tensorrt rather than normal tensorflow. I am fairly new to regular tensorflow so am not sure how difficult switching is. Is this a path worth pursuing?

Hi,

1. Since our TensorFlow is built with CUDA support, the CUDA library will be loaded to the memory when initializing the TensorFlow.
This will be improved in JetPack 6 since CUDA 12 has the new lazy loading feature.

2. Not sure which model you use.
Please find below for an examples:

Thanks.

Thanks for the response.

Is it expected that “lazy loading” in CUDA 12 will free up memory more often / automatically? It seems like the issue I’m observing is related to both loading the model (at which point memory increases from ~1.4GB to ~4 GB) but also running the model (memory stays around ~6-6.5GB).

I’ll take a look at the repo you linked for tensorrt

This topic was automatically closed 14 days after the last reply. New replies are no longer allowed.