How to get memory usage of my infer model via Python API?

[jmren168]

Hi,

I'm trying to get the memory usage of my infer model. Following ChatGPT to create a python script below, but the result seems to work incorrectly, because I always get 4-5 KB, even the onnx file is more than 50KB. Any comments are highly appreciated.

import tracemalloc
import onnxruntime

def load_onnx_model(model_path):
    session = onnxruntime.InferenceSession(model_path)
    return session

def main():
    tracemalloc.start()  # Start tracing memory

    onnx_model_path = 'your_model.onnx'  # Replace with the path to your ONNX file
    onnx_session = load_onnx_model(onnx_model_path)

    current, peak = tracemalloc.get_traced_memory()
    print(f"Current memory usage: {current / 10**3} KB")
    print(f"Peak memory usage: {peak / 10**3} KB")

    tracemalloc.stop()  # Stop tracing memory

if __name__ == "__main__":
    main()

Another similar question is #7205, but it is not clear how to get the memory usage via Python API. Thanks if any furthur comments.