diff --git a/qlib/contrib/model/pytorch_gats_ts.py b/qlib/contrib/model/pytorch_gats_ts.py index 633da075a..25fdef5f4 100644 --- a/qlib/contrib/model/pytorch_gats_ts.py +++ b/qlib/contrib/model/pytorch_gats_ts.py @@ -22,6 +22,7 @@ import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader +from torch.utils.data import Sampler from ...model.base import Model from ...data.dataset import DatasetH @@ -30,6 +31,21 @@ from ...contrib.model.pytorch_lstm import LSTMModel from ...contrib.model.pytorch_gru import GRUModel +class DailyBatchSampler(Sampler): + def __init__(self, data_souce): + self.data_source = data_source + self.data = self.data_source.loc[self.data_source.get_index()] + self.daily_count = self.data.groupby(level=0).size().values + self.daily_index = np.roll(np.cumsum(self.daily_count), 1) + + def __iter__(self): + for idx, count in zip(self.daily_index, self.daily_count): + yield slice(idx, idx + count) + + def __len__(self): + return len(self.data_source) + + class GATs(Model): """GATs Model @@ -235,8 +251,11 @@ class GATs(Model): dl_train.config(fillna_type="ffill+bfill") # process nan brought by dataloader dl_valid.config(fillna_type="ffill+bfill") # process nan brought by dataloader - train_loader = DataLoader(dl_train, batch_size=self.batch_size, shuffle=True, num_workers=self.n_jobs) - valid_loader = DataLoader(dl_valid, batch_size=self.batch_size, shuffle=False, num_workers=self.n_jobs) + sampler_train = DailyBatchSampler(dl_train) + sampler_valid = DailyBatchSampler(dl_valid) + + train_loader = DataLoader(dl_train, sampler=sampler_train, num_workers=self.n_jobs) + valid_loader = DataLoader(dl_valid, sampler=sampler_valid, num_workers=self.n_jobs) if save_path == None: save_path = create_save_path(save_path) @@ -307,7 +326,8 @@ class GATs(Model): dl_test = dataset.prepare("test", data_key=DataHandlerLP.DK_I) dl_test.config(fillna_type="ffill+bfill") - test_loader = DataLoader(dl_test, batch_size=self.batch_size, num_workers=self.n_jobs) + sampler_test = DailyBatchSampler(dl_test) + test_loader = DataLoader(dl_test, sampler=sampler_test, num_workers=self.n_jobs) self.ALSTM_model.eval() preds = []