liuyu/qlib
1
0
Fork 0
mirror of https://github.com/microsoft/qlib.git synced 2026-07-16 01:06:56 +08:00
Code Activity

fix pkl file not loading in StaticDataLoader (#1896)

Browse Source 
* fix pkl file not loading in StaticDataLoader

* resolve hard code

* resolve hard code
This commit is contained in:
Linlang
2025-03-18 16:05:24 +08:00
committed by GitHub
parent 82f1ef2def
commit 4d621bff99
9 changed files with 20 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
pyproject.toml
View File

@@ -44,6 +44,7 @@ dependencies = [
"matplotlib", "matplotlib",
"jupyter", "jupyter",
"nbconvert", "nbconvert",
"pyarrow",
] ]
[project.optional-dependencies] [project.optional-dependencies]

3
qlib/data/dataset/loader.py
View File

@@ -279,6 +279,9 @@ class StaticDataLoader(DataLoader, Serializable):
) )
self._data.sort_index(inplace=True) self._data.sort_index(inplace=True)
elif isinstance(self._config, (str, Path)): elif isinstance(self._config, (str, Path)):
if str(self._config).strip().endswith(".parquet"):
self._data = pd.read_parquet(self._config, engine="pyarrow")
else:
with Path(self._config).open("rb") as f: with Path(self._config).open("rb") as f:
self._data = pickle.load(f) self._data = pickle.load(f)
elif isinstance(self._config, pd.DataFrame): elif isinstance(self._config, pd.DataFrame):

2
qlib/rl/trainer/trainer.py
View File

@@ -200,7 +200,7 @@ class Trainer:
if ckpt_path is not None: if ckpt_path is not None:
_logger.info("Resuming states from %s", str(ckpt_path)) _logger.info("Resuming states from %s", str(ckpt_path))
self.load_state_dict(torch.load(ckpt_path)) self.load_state_dict(torch.load(ckpt_path, weights_only=False))
else: else:
self.initialize() self.initialize()

2
tests/rl/test_trainer.py
View File

@@ -194,7 +194,7 @@ def test_trainer_checkpoint():
assert (output_dir / "002.pth").exists() assert (output_dir / "002.pth").exists()
assert os.readlink(output_dir / "latest.pth") == str(output_dir / "002.pth") assert os.readlink(output_dir / "latest.pth") == str(output_dir / "002.pth")
trainer.load_state_dict(torch.load(output_dir / "001.pth")) trainer.load_state_dict(torch.load(output_dir / "001.pth", weights_only=False))
assert trainer.current_iter == 1 assert trainer.current_iter == 1
assert trainer.current_episode == 100 assert trainer.current_episode == 100