File size: 1,403 Bytes
8cc5633
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
# create_dataloaders.py

import logging

from momentfm.utils.utils import control_randomness
from torch.utils.data import DataLoader

from transformer_model.scripts.config_transformer import (BATCH_SIZE,
                                                          FORECAST_HORIZON)
from transformer_model.scripts.utils.informer_dataset_class import \
    InformerDataset


def create_dataloaders():
    logging.info("Setting random seeds...")
    control_randomness(seed=13)

    logging.info("Loading training dataset...")
    train_dataset = InformerDataset(
        data_split="train", random_seed=13, forecast_horizon=FORECAST_HORIZON
    )
    logging.info(
        "Train set loaded — Samples: %d | Features: %d",
        len(train_dataset),
        train_dataset.n_channels,
    )

    logging.info("Loading test dataset...")
    test_dataset = InformerDataset(
        data_split="test", random_seed=13, forecast_horizon=FORECAST_HORIZON
    )
    logging.info(
        "Test set loaded — Samples: %d | Features: %d",
        len(test_dataset),
        test_dataset.n_channels,
    )

    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=True)

    logging.info("Dataloaders created successfully.")
    return train_loader, test_loader


if __name__ == "__main__":
    create_dataloaders()