How to store an sklearn RF model such that it has fastest load time?

user_4232601 0
Viewed 149

I have an RF model that I am trying to load faster during deployment. The first time it was created and saved as a compressed joblib file. But the load time is too high for my need. Here's a snippet of my attempt to benchmark what I have:

# load the base file:
print("Loading compressed joblib...")
MODEL = f"{base_path}/data/model_files/{dataset}_RF_{C}.compressed"
tic = time.perf_counter()
with open(MODEL, "rb") as f:
    model = joblib.load(f)
    mem = sys.getsizeof(model) / 1024 / 1024
toc = time.perf_counter()
load_time = toc - tic
size = os.path.getsize(MODEL) / 1024 / 1024
print(f"file on disk: {size:0.2f} MB\nfile on RAM: {mem:0.2f} MB\nload time: {load_time:0.3f} seconds\n")

Output:

Loading compressed joblib...
file on disk: 52.22 MB
file on RAM: 0.00 MB
load time: 9.510 seconds

I am unsure why the file size loaded in RAM is not shown. Please help with that as well if you know.

Uncompressed joblib:

# create uncompressed joblib, save it, load it, benchmark
print("Saving uncompressed joblib...")
MODEL = f"{base_path}/data/model_files/{dataset}_RF_{C}.uncompressed.joblib"
tic = time.perf_counter()
with open(MODEL, "wb") as f:
    joblib.dump(model, f, compress=False)
toc = time.perf_counter()
save_time = toc - tic
size = os.path.getsize(MODEL) / 1024 / 1024
print(f"file on disk: {size:0.2f} MB\nsave time: {save_time:0.3f} seconds\n")
del model

print("Loading uncompressed joblib...")
tic = time.perf_counter()
with open(MODEL, "rb") as f:
    model = joblib.load(f)
    mem = sys.getsizeof(model) / 1024 / 1024
toc = time.perf_counter()
load_time = toc - tic
size = os.path.getsize(MODEL) / 1024 / 1024
print(f"file on disk: {size:0.2f} MB\nfile on RAM: {mem:0.2f} MB\nload time: {load_time:0.3f} seconds")

Output:

Saving uncompressed joblib...
file on disk: 4376.51 MB
save time: 8.718 seconds

Loading uncompressed joblib...
file on disk: 4376.51 MB
file on RAM: 0.00 MB
load time: 7.645 seconds

As it can be seen, by trading off 50MB compression to 4.3GB file size, I am not gaining much on load time.

joblib scikit-learn pickle random-forest
1 Answers

Based on your description, it sounds like you are loading the model into a deployment environment so that it can return live predictions.

If that is the case, there are a few things that I recommend investigating:

  • Load your model into memory before it needs to start serving live requests/at a time when 5-10 seconds does not matter. For example, load the model into memory when the service or lambda function you are using is starting up.
  • Convert the model into ONNX format, by following the sklearn-onnx guide. As a microsoft blog post describes, using ONNX can especially help to make the inference time faster.
  • Consider training smaller Random Forest models by increasing the max_depth, min_samples_split, and/or the min_samples_leaf parameters (4.3GB for an uncompressed model sounds fairly large)
user_6825250
Related