61286c6e8f
* Removed some data copying. * Make sure gpu_id is valid before any configuration is carried out.
52 lines
1.8 KiB
Python
52 lines
1.8 KiB
Python
'''Loading a pickled model generated by test_pickling.py, only used by
|
|
`test_gpu_with_dask.py`'''
|
|
import unittest
|
|
import os
|
|
import xgboost as xgb
|
|
import json
|
|
|
|
from test_gpu_pickling import build_dataset, model_path, load_pickle
|
|
|
|
|
|
class TestLoadPickle(unittest.TestCase):
|
|
def test_load_pkl(self):
|
|
'''Test whether prediction is correct.'''
|
|
assert os.environ['CUDA_VISIBLE_DEVICES'] == '-1'
|
|
bst = load_pickle(model_path)
|
|
x, y = build_dataset()
|
|
test_x = xgb.DMatrix(x)
|
|
res = bst.predict(test_x)
|
|
assert len(res) == 10
|
|
|
|
def test_predictor_type_is_auto(self):
|
|
'''Under invalid CUDA_VISIBLE_DEVICES, predictor should be set to
|
|
auto'''
|
|
assert os.environ['CUDA_VISIBLE_DEVICES'] == '-1'
|
|
bst = load_pickle(model_path)
|
|
config = bst.save_config()
|
|
config = json.loads(config)
|
|
assert config['learner']['gradient_booster']['gbtree_train_param'][
|
|
'predictor'] == 'auto'
|
|
|
|
def test_predictor_type_is_gpu(self):
|
|
'''When CUDA_VISIBLE_DEVICES is not specified, keep using
|
|
`gpu_predictor`'''
|
|
assert 'CUDA_VISIBLE_DEVICES' not in os.environ.keys()
|
|
bst = load_pickle(model_path)
|
|
config = bst.save_config()
|
|
config = json.loads(config)
|
|
assert config['learner']['gradient_booster']['gbtree_train_param'][
|
|
'predictor'] == 'gpu_predictor'
|
|
|
|
def test_wrap_gpu_id(self):
|
|
assert os.environ['CUDA_VISIBLE_DEVICES'] == '0'
|
|
bst = load_pickle(model_path)
|
|
config = bst.save_config()
|
|
config = json.loads(config)
|
|
assert config['learner']['generic_param']['gpu_id'] == '0'
|
|
|
|
x, y = build_dataset()
|
|
test_x = xgb.DMatrix(x)
|
|
res = bst.predict(test_x)
|
|
assert len(res) == 10
|