Memory leak after running onnx model numerous times #22271
Open
Description
Describe the issue
In then following scenario I see that memory consumption of python process is continuously growing:
- Train pipeline of CountVectorizer and RandomForestClassifier using sklearn.
- Then convert pipeline to onnxruntime.InferenceSession.
- Run model many times on random data and measure memory consumption.
To reproduce
The following code run run_model_many_times function two twice. Once with running onnxruntime.InferenceSession method and once without. While code is running it is printing memory consumption in megabytes, which is close to macOs activity monitor numbers.
In the first case memory consumption is continuously growing and in the second case is not.
import uuid
import onnxruntime
import psutil
from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import StringTensorType
from sklearn.ensemble import RandomForestClassifier
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.pipeline import Pipeline
def create_trained_model():
x = [str(uuid.uuid4()) for _ in range(100)]
y = ["a" for _ in range(50)] + ["b" for _ in range(50)]
pipeline = Pipeline(
steps=[
("vectorizer", CountVectorizer()),
("classifier", RandomForestClassifier())
]
)
pipeline.fit(x, y)
return pipeline
def convert_sklearn_model_to_onnx_session(sklearn_model):
onnx_model_proto_string = convert_sklearn(
sklearn_model,
initial_types=[('features', StringTensorType((None,)))],
verbose=False
).SerializeToString()
return onnxruntime.InferenceSession(onnx_model_proto_string, providers=["CPUExecutionProvider"])
def get_one_prediction(onnx_model, input_data):
return onnx_model.run(
[onnx_model.get_outputs()[1].name],
{onnx_model.get_inputs()[0].name: input_data}
)
def get_used_megabytes():
# https://psutil.readthedocs.io/en/latest/#psutil.Process.memory_full_info
# uss (Linux, macOS, Windows): aka “Unique Set Size”,
# this is the memory which is unique to a process and which would be freed
# if the process was terminated right now.
return psutil.Process().memory_full_info().uss / (1024 * 1024)
def run_model_many_times(onnx_model, count, *, dummy):
print("run_model_many_times, dummy = ", dummy)
print("Memory in the beginning = ", get_used_megabytes())
for i in range(count):
input_data = [str(uuid.uuid4()) for _ in range(20)]
if not dummy:
get_one_prediction(onnx_model, input_data)
if i % 10000 == 0:
print("Memory in the middle = ", get_used_megabytes())
print("Memory in the end = ", get_used_megabytes())
def main():
sklearn_model = create_trained_model()
onnx_model = convert_sklearn_model_to_onnx_session(sklearn_model)
count = 100000
run_model_many_times(onnx_model, count, dummy=False)
run_model_many_times(onnx_model, count, dummy=True)
if __name__ == '__main__':
main()Example of printing:
run_model_many_times, dummy = False
Memory in the beginning = 126.875
Memory in the middle = 126.875
Memory in the middle = 140.96875
Memory in the middle = 156.296875
Memory in the middle = 171.671875
Memory in the middle = 187.015625
Memory in the middle = 202.328125
Memory in the middle = 217.640625
Memory in the middle = 232.96875
Memory in the middle = 248.3125
Memory in the middle = 263.625
Memory in the end = 278.96875
run_model_many_times, dummy = True
Memory in the beginning = 278.96875
Memory in the middle = 278.96875
Memory in the middle = 278.96875
Memory in the middle = 278.96875
Memory in the middle = 278.96875
Memory in the middle = 278.96875
Memory in the middle = 278.96875
Memory in the middle = 278.96875
Memory in the middle = 278.96875
Memory in the middle = 278.96875
Memory in the middle = 278.96875
Memory in the end = 278.96875
Urgency
On production server python process consumed 100 GB of RAM memory after two months. The only solution is restart the process. So yes, it's urgent.
Platform
Linux
OS Version
both of MacOs 15.0 and Ubuntu 20
ONNX Runtime Installation
Released Package
ONNX Runtime Version or Commit ID
1.19.2
ONNX Runtime API
Python
Architecture
X64
Execution Provider
Default CPU
Execution Provider Library Version
No response