Memory-mapped XLinear Model

[weiliw-amz]

Issue #, if available:
N/A

Description of changes:

• Memory-mapped PECOS XLinear model
  • Greatly reduce loading time.
  • Ideal for large models that user want to quickly try a few inferences without waiting for loading full model into memory.
  • Also capable for large model inference that could not be stored in memory.

Usage:
User needs to have a XLinear Model saved on disk (in original .npz format), and manually compile into mmap format by calling compile_mmap_model:

import sys
import pathlib
from pecos.xmc.xlinear.model import XLinearModel


npz_model_path = f"/path/to/xlinear/pecos-models/"
mmap_model_path = f"/path/to/xlinear/mmap-models/"

pathlib.Path(mmap_model_path).mkdir(parents=True, exist_ok=True)
print(f"Compiling mmap model from: {npz_model_path}, will save to : {mmap_model_path}...")
XLinearModel.compile_mmap_model(npz_model_path, mmap_model_path)
print("mmap model saved.")

Then user can load the memory-mapped model and do inference:

import sys
from pecos.xmc.xlinear.model import XLinearModel


mmap_model_path = f"/path/to/xlinear/mmap-models/"

# Load model
if sys.argv[2] == "--cmmap-lazy":
    print("Loading C/C++ mem map model lazy-loaded...")
    xlm = XLinearModel.load(mmap_model_path, is_predict_only=True, lazy_load=True)
elif sys.argv[2] == "--cmmap":
    print("Loading C/C++ mem map model...")
    xlm = XLinearModel.load(mmap_model_path, is_predict_only=True, lazy_load=False)
else:
    print(f"Wrong option: {sys.argv[2]}")

# Load test data
Xt = XLinearModel.load_feature_matrix(f"/test/data/validation/X.npz")
Yt = XLinearModel.load_label_matrix(f"/test/data/validation/Y.npz")

# Predict
Yt_pred = xlm.predict(Xt)
Yt_pred = Yt_pred.tocsr()
metric = smat_util.Metrics.generate(Yt, Yt_pred, topk=10)
print(metric)

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.