c589eff941
* Only define `gpu_id` and `n_gpus` in `LearnerTrainParam` * Pass LearnerTrainParam through XGBoost vid factory method. * Disable all GPU usage when GPU related parameters are not specified (fixes XGBoost choosing GPU over aggressively). * Test learner train param io. * Fix gpu pickling.
21 lines
550 B
Python
21 lines
550 B
Python
'''Loading a pickled model generated by test_pickling.py'''
|
|
import pickle
|
|
import unittest
|
|
import os
|
|
import xgboost as xgb
|
|
import sys
|
|
|
|
sys.path.append("tests/python")
|
|
from test_pickling import build_dataset, model_path
|
|
|
|
|
|
class TestLoadPickle(unittest.TestCase):
|
|
def test_load_pkl(self):
|
|
assert os.environ['CUDA_VISIBLE_DEVICES'] == ''
|
|
with open(model_path, 'rb') as fd:
|
|
bst = pickle.load(fd)
|
|
x, y = build_dataset()
|
|
test_x = xgb.DMatrix(x)
|
|
res = bst.predict(test_x)
|
|
assert len(res) == 10
|