38 lines
1.1 KiB
Python
38 lines
1.1 KiB
Python
import xgboost as xgb
|
|
import numpy as np
|
|
from sklearn.datasets import load_digits
|
|
from sklearn.cross_validation import KFold, StratifiedKFold, train_test_split
|
|
from sklearn.metrics import mean_squared_error
|
|
import unittest
|
|
|
|
rng = np.random.RandomState(1994)
|
|
|
|
|
|
class TestCrossValidation(unittest.TestCase):
|
|
def test_cv(self):
|
|
digits = load_digits(3)
|
|
X = digits['data']
|
|
y = digits['target']
|
|
dm = xgb.DMatrix(X, label=y)
|
|
|
|
params = {
|
|
'max_depth': 2,
|
|
'eta': 1,
|
|
'silent': 1,
|
|
'objective':
|
|
'multi:softprob',
|
|
'num_class': 3
|
|
}
|
|
|
|
seed = 2016
|
|
nfolds = 5
|
|
skf = StratifiedKFold(y, n_folds=nfolds, shuffle=True, random_state=seed)
|
|
|
|
import pandas as pd
|
|
cv1 = xgb.cv(params, dm, num_boost_round=10, nfold=nfolds, seed=seed)
|
|
cv2 = xgb.cv(params, dm, num_boost_round=10, folds=skf, seed=seed)
|
|
cv3 = xgb.cv(params, dm, num_boost_round=10, nfold=nfolds, stratified=True, seed=seed)
|
|
assert cv1.shape[0] == cv2.shape[0] and cv2.shape[0] == cv3.shape[0]
|
|
assert cv2.iloc[-1,0] == cv3.iloc[-1,0]
|
|
|