Class TestLightgbm

tests/test_lightgbm.py:8–76

from common import gpu_test

class TestLightgbm(unittest.TestCase):
    # Based on the "simple_example" from their documentation:
    # https://github.com/Microsoft/LightGBM/blob/master/examples/python-guide/simple_example.py
    def test_cpu(self):
        lgb_train, lgb_eval = self.load_datasets()

        params = {
            'task': 'train',
            'boosting_type': 'gbdt',
            'objective': 'regression',
            'metric': {'l2', 'auc'},
            'num_leaves': 31,
            'learning_rate': 0.05,
            'feature_fraction': 0.9,
            'bagging_fraction': 0.8,
            'bagging_freq': 5,
            'force_row_wise': True,
            'verbose': 0
        }

        # Run only one round for faster test
        gbm = lgb.train(params,
                        lgb_train,
                        num_boost_round=1,
                        valid_sets=lgb_eval,
                        callbacks=[lgb.early_stopping(stopping_rounds=1)])

        self.assertEqual(1, gbm.best_iteration)

    @gpu_test
    def test_gpu(self):
        lgb_train, lgb_eval = self.load_datasets()
        
        params = {
            'boosting_type': 'gbdt',
            'objective': 'regression',
            'metric': 'auc',
            'num_leaves': 31,
            'learning_rate': 0.05,
            'feature_fraction': 0.9,
            'bagging_fraction': 0.8,
            'bagging_freq': 5,
            'force_row_wise': True,
            'verbose': 1,
            'device': 'gpu'
        }
        
        # Run only one round for faster test
        gbm = lgb.train(params,
                        lgb_train,
                        num_boost_round=1,
                        valid_sets=lgb_eval,
                        callbacks=[lgb.early_stopping(stopping_rounds=1)])

        self.assertEqual(1, gbm.best_iteration)
    
    def load_datasets(self):
        df_train = pd.read_csv('/input/tests/data/lgb_train.csv', header=None, sep='\t')