9c9070aea2
* Do not derive from unittest.TestCase (not needed for pytest) * assertRaises -> pytest.raises * Simplify test_empty_dmatrix with test parametrization * setUpClass -> setup_class, tearDownClass -> teardown_class * Don't import unittest; import pytest * Use plain assert * Use parametrized tests in more places * Fix test_gpu_with_sklearn.py * Put back run_empty_dmatrix_reg / run_empty_dmatrix_cls * Fix test_eta_decay_gpu_hist * Add parametrized tests for monotone constraints * Fix test names * Remove test parametrization * Revise test_slice to be not flaky
59 lines
1.3 KiB
Python
59 lines
1.3 KiB
Python
import pickle
|
|
import numpy as np
|
|
import xgboost as xgb
|
|
import os
|
|
|
|
|
|
kRows = 100
|
|
kCols = 10
|
|
|
|
|
|
def generate_data():
|
|
X = np.random.randn(kRows, kCols)
|
|
y = np.random.randn(kRows)
|
|
return X, y
|
|
|
|
|
|
class TestPickling:
|
|
def run_model_pickling(self, xgb_params):
|
|
X, y = generate_data()
|
|
dtrain = xgb.DMatrix(X, y)
|
|
bst = xgb.train(xgb_params, dtrain)
|
|
|
|
dump_0 = bst.get_dump(dump_format='json')
|
|
assert dump_0
|
|
|
|
filename = 'model.pkl'
|
|
|
|
with open(filename, 'wb') as fd:
|
|
pickle.dump(bst, fd)
|
|
|
|
with open(filename, 'rb') as fd:
|
|
bst = pickle.load(fd)
|
|
|
|
with open(filename, 'wb') as fd:
|
|
pickle.dump(bst, fd)
|
|
|
|
with open(filename, 'rb') as fd:
|
|
bst = pickle.load(fd)
|
|
|
|
assert bst.get_dump(dump_format='json') == dump_0
|
|
|
|
if os.path.exists(filename):
|
|
os.remove(filename)
|
|
|
|
def test_model_pickling_binary(self):
|
|
params = {
|
|
'nthread': 1,
|
|
'tree_method': 'hist'
|
|
}
|
|
self.run_model_pickling(params)
|
|
|
|
def test_model_pickling_json(self):
|
|
params = {
|
|
'nthread': 1,
|
|
'tree_method': 'hist',
|
|
'enable_experimental_json_serialization': True
|
|
}
|
|
self.run_model_pickling(params)
|