1
0
mirror of https://github.com/microsoft/qlib.git synced 2026-07-12 23:36:54 +08:00

fix tra dataset bug (#1050)

This commit is contained in:
you-n-g
2022-04-15 17:15:44 +08:00
committed by GitHub
parent e1271a83f7
commit 41447f320b
2 changed files with 12 additions and 5 deletions

View File

@@ -6,8 +6,7 @@ import torch
import numpy as np import numpy as np
import pandas as pd import pandas as pd
from qlib.utils import init_instance_by_config from qlib.data.dataset import DatasetH
from qlib.data.dataset import DatasetH, DataHandler
device = "cuda" if torch.cuda.is_available() else "cpu" device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -95,7 +94,7 @@ class MTSDatasetH(DatasetH):
shuffle=True, shuffle=True,
pin_memory=False, pin_memory=False,
drop_last=False, drop_last=False,
**kwargs **kwargs,
): ):
assert horizon > 0, "please specify `horizon` to avoid data leakage" assert horizon > 0, "please specify `horizon` to avoid data leakage"
@@ -150,8 +149,15 @@ class MTSDatasetH(DatasetH):
def _prepare_seg(self, slc, **kwargs): def _prepare_seg(self, slc, **kwargs):
fn = _get_date_parse_fn(self._index[0][1]) fn = _get_date_parse_fn(self._index[0][1])
start_date = fn(slc.start)
end_date = fn(slc.stop) if isinstance(slc, slice):
start, stop = slc.start, slc.stop
elif isinstance(slc, (list, tuple)):
start, stop = slc
else:
raise NotImplementedError(f"This type of input is not supported")
start_date = fn(start)
end_date = fn(stop)
obj = copy.copy(self) # shallow copy obj = copy.copy(self) # shallow copy
# NOTE: Seriable will disable copy `self._data` so we manually assign them here # NOTE: Seriable will disable copy `self._data` so we manually assign them here
obj._data = self._data obj._data = self._data

View File

@@ -171,6 +171,7 @@ class DatasetH(Dataset):
Parameters Parameters
---------- ----------
slc : please refer to the docs of `prepare` slc : please refer to the docs of `prepare`
NOTE: it may not be an instance of slice. It may be a segment of `segments` from `def prepare`
""" """
if hasattr(self, "fetch_kwargs"): if hasattr(self, "fetch_kwargs"):
return self.handler.fetch(slc, **kwargs, **self.fetch_kwargs) return self.handler.fetch(slc, **kwargs, **self.fetch_kwargs)