test_boosted_trees_regression_numeric.py 文件源码

def _train_convert_evaluate(self, bt_params = {}, **params):
        """
        Set up the unit test by loading the dataset and training a model.
        """
        # Train a model
        xgb_model = xgboost.XGBRegressor(**params)
        xgb_model.fit(self.X, self.target)

        # Convert the model (feature_names can't be given because of XGboost)
        spec = xgb_converter.convert(xgb_model, self.feature_names, self.output_name, force_32bit_float = False)

        # Get predictions
        df = pd.DataFrame(self.X, columns=self.feature_names)
        df['prediction'] = xgb_model.predict(self.X)

        # Evaluate it
        metrics = evaluate_regressor(spec, df, target = 'target', verbose = False)
        return metrics