test_xgboost.py 文件源码

def test_train_minimum_params(df_train):
    params = {'num_rounds': 1}
    # TODO: Anything > 1 worker can easily get stuck, as a system
    # with 2 cpus will only spawn a single worker.
    model = mjolnir.training.xgboost.train(df_train, params, num_workers=1)

    # What else can we practically assert?
    df_transformed = model.transform(df_train)
    assert 'prediction' in df_transformed.columns
    assert 0.74 == pytest.approx(model.eval(df_train), abs=0.01)

    # make sure train didn't clobber the incoming params
    assert params['num_rounds'] == 1