Keep old behavior of model loading

This commit is contained in:
Emre 2022-12-01 11:03:51 +03:00
parent 4a9982f86b
commit 396e666e9b
No known key found for this signature in database
GPG Key ID: 0EAD2EE11B666ABA
1 changed files with 2 additions and 2 deletions

View File

@ -503,7 +503,7 @@ class FreqaiDataDrawer:
dump(model, save_path / f"{dk.model_filename}_model.joblib")
elif self.model_type == 'keras':
model.save(save_path / f"{dk.model_filename}_model.h5")
elif self.model_type in ['stable_baselines3', 'sb3_contrib']:
elif 'stable_baselines' in self.model_type or 'sb3_contrib' == self.model_type:
model.save(save_path / f"{dk.model_filename}_model.zip")
if dk.svm_model is not None:
@ -589,7 +589,7 @@ class FreqaiDataDrawer:
elif self.model_type == 'keras':
from tensorflow import keras
model = keras.models.load_model(dk.data_path / f"{dk.model_filename}_model.h5")
elif self.model_type in ['stable_baselines3', 'sb3_contrib']:
elif 'stable_baselines' in self.model_type or 'sb3_contrib' == self.model_type:
mod = importlib.import_module(
self.model_type, self.freqai_info['rl_config']['model_type'])
MODELCLASS = getattr(mod, self.freqai_info['rl_config']['model_type'])