allow user to pass test_size = 0 and avoid using eval sets in prediction models

This commit is contained in:
robcaulk
2022-07-25 19:40:13 +02:00
parent 55cf378ec2
commit 56b17e6f3c
4 changed files with 67 additions and 44 deletions

View File

@@ -28,17 +28,22 @@ class CatboostPredictionModel(BaseRegressionModel):
label=data_dictionary["train_labels"],
weight=data_dictionary["train_weights"],
)
test_data = Pool(
data=data_dictionary["test_features"],
label=data_dictionary["test_labels"],
weight=data_dictionary["test_weights"],
)
if self.freqai_info.get('data_split_parameters', {}).get('test_size', 0.1) == 0:
test_data = None
else:
test_data = Pool(
data=data_dictionary["test_features"],
label=data_dictionary["test_labels"],
weight=data_dictionary["test_weights"],
)
model = CatBoostRegressor(
allow_writing_files=False,
**self.model_training_parameters,
)
if self.freqai_info.get('data_split_parameters', {}).get('test_size', 0.1) == 0:
test_data = None
model.fit(X=train_data, eval_set=test_data)
return model

View File

@@ -36,7 +36,9 @@ class CatboostPredictionMultiModel(BaseRegressionModel):
model = MultiOutputRegressor(estimator=cbr)
model.fit(X=X, y=y, sample_weight=sample_weight) # , eval_set=eval_set)
train_score = model.score(X, y)
test_score = model.score(*eval_set)
logger.info(f"Train score {train_score}, Test score {test_score}")
if self.freqai_info.get('data_split_parameters', {}).get('test_size', 0.1) != 0:
train_score = model.score(X, y)
test_score = model.score(*eval_set)
logger.info(f"Train score {train_score}, Test score {test_score}")
return model

View File

@@ -25,11 +25,15 @@ class LightGBMPredictionModel(BaseRegressionModel):
all the training and test data/labels.
"""
eval_set = (data_dictionary["test_features"], data_dictionary["test_labels"])
if self.freqai_info.get('data_split_parameters', {}).get('test_size', 0.1) == 0:
eval_set = None
else:
eval_set = (data_dictionary["test_features"], data_dictionary["test_labels"])
X = data_dictionary["train_features"]
y = data_dictionary["train_labels"]
model = LGBMRegressor(**self.model_training_parameters)
model.fit(X=X, y=y, eval_set=eval_set)
return model