Saving Models - NGBoost User Guide

from ngboost import NGBRegressor

from sklearn.datasets import load_diabetes
from sklearn.model_selection import train_test_split

X, Y = load_diabetes(return_X_y=True)
X_reg_train, X_reg_test, Y_reg_train, Y_reg_test = train_test_split(X, Y, test_size=0.2)

Saving ngboost models is easy with the pickle package:

ngb = NGBRegressor().fit(X_reg_train, Y_reg_train)

[iter 0] loss=5.7647 val_loss=0.0000 scale=1.0000 norm=65.8357

[iter 100] loss=5.4150 val_loss=0.0000 scale=1.0000 norm=45.1825

[iter 200] loss=5.1782 val_loss=0.0000 scale=1.0000 norm=36.4965

[iter 300] loss=5.0259 val_loss=0.0000 scale=1.0000 norm=32.7998

[iter 400] loss=4.9240 val_loss=0.0000 scale=2.0000 norm=61.5887

import pickle
from pathlib import Path

file_path = Path('ngbtest.p')

with file_path.open("wb") as f:
    pickle.dump(ngb, f)

with file_path.open("rb") as f:
    ngb_unpickled = pickle.load(f)

Y_preds = ngb_unpickled.predict(X_reg_test)
Y_dists = ngb_unpickled.pred_dist(X_reg_test)

Y_dists[0:5].params

{'loc': array([ 96.99460752, 149.95843629, 180.66874368,  84.50476373,
         88.24486158]),
 'scale': array([29.43302582, 36.20735616, 45.68591978, 25.37584676, 32.04836874])}

To avoid pickle, you can use skops, which does not run arbitrary code when it loads a file. Compress the file to keep it small, and check the list of types before you trust them:

from zipfile import ZIP_DEFLATED

import skops.io as sio

sio.dump(ngb, "ngbtest.skops", compression=ZIP_DEFLATED)

unknown_types = sio.get_untrusted_types(file="ngbtest.skops")
unknown_types

['ngboost.api.NGBRegressor',
 'ngboost.distns.normal.Normal',
 'ngboost.scores.LogScore',
 'sklearn.tree._tree.Tree']

ngb_skops = sio.load("ngbtest.skops", trusted=unknown_types)

ngb_skops.pred_dist(X_reg_test)[0:5].params

{'loc': array([ 96.99460752, 149.95843629, 180.66874368,  84.50476373,
         88.24486158]),
 'scale': array([29.43302582, 36.20735616, 45.68591978, 25.37584676, 32.04836874])}