Use cloudpickle throughout

This commit is contained in:
Matthias
2023-11-25 19:36:14 +01:00
parent e8d0b01991
commit eeb460e55c
+3 -2
View File
@@ -12,7 +12,7 @@ import numpy as np
import pandas as pd import pandas as pd
import psutil import psutil
import rapidjson import rapidjson
from joblib import dump, load from joblib import load
from joblib.externals import cloudpickle from joblib.externals import cloudpickle
from numpy.typing import NDArray from numpy.typing import NDArray
from pandas import DataFrame from pandas import DataFrame
@@ -471,7 +471,8 @@ class FreqaiDataDrawer:
# Save the trained model # Save the trained model
if self.model_type == 'joblib': if self.model_type == 'joblib':
dump(model, save_path / f"{dk.model_filename}_model.joblib") with (save_path / f"{dk.model_filename}_model.joblib").open("wb") as fp:
cloudpickle.dump(model, fp)
elif self.model_type == 'keras': elif self.model_type == 'keras':
model.save(save_path / f"{dk.model_filename}_model.h5") model.save(save_path / f"{dk.model_filename}_model.h5")
elif self.model_type in ["stable_baselines3", "sb3_contrib", "pytorch"]: elif self.model_type in ["stable_baselines3", "sb3_contrib", "pytorch"]: