datamodule.py

from typing import List
import torch
from torch.utils.data import TensorDataset
from torch.utils.data import DataLoader
import pytorch_lightning as pl

class RoutesDataset():
    def __init__(self, datapath: str, label_names: List) -> None:
        super().__init__()
        dataset = torch.load(datapath)
        self.X_train = dataset['routes_train'].float()
        self.X_val = dataset['routes_val'].float()
        self.X_test = dataset['routes_test'].float()
        self.Y_train = [dataset["labels_train"][label_name] for label_name in label_names]
        self.Y_val = [dataset["labels_val"][label_name] for label_name in label_names]
        self.Y_test = [dataset["labels_test"][label_name] for label_name in label_names]
        self.feature_names = dataset['feature_columns']
        self.label_names = dataset['label_columns']
        self.routes_length = dataset['lengths']

class RoutesDataModule(pl.LightningDataModule):
    def __init__(self, datapath: str, label_names: str, batch_size: int, num_workers: int, mode: str = "train") -> None:
        super().__init__()
        assert mode in ["train", "test"]
        self.datapath = datapath
        self.label_names = label_names
        self.batch_size = batch_size
        self.num_workers = num_workers
        self.mode = mode
        self.dataset = RoutesDataset(self.datapath, self.label_names)
        self.trainset = TensorDataset(self.dataset.X_train, *self.dataset.Y_train)
        self.valset = TensorDataset(self.dataset.X_val, *self.dataset.Y_val)
        self.testset = TensorDataset(self.dataset.X_test, *self.dataset.Y_test)
    
    def setup(self, stage):
        pass

    def train_dataloader(self) -> None:
        if self.mode == "train":
            shuffle=True
        elif self.mode == "test":
            shuffle=False
        return DataLoader(self.trainset, self.batch_size, num_workers=self.num_workers, pin_memory=True, shuffle=shuffle)

    def val_dataloader(self) -> None:
        return DataLoader(self.valset, self.batch_size, num_workers=self.num_workers, pin_memory=True)

    def test_dataloader(self) -> None:
        return DataLoader(self.testset, self.batch_size, num_workers=self.num_workers, pin_memory=True)

if __name__ == "__main__":
    #RoutesDataset('data/routes_preprocessed_with_labels.pt', ["summed_travel_time", "summed_length", "mean_curvature", "abs_mean_grade", "roadType_share"])
    RoutesDataModule('data/routes_preprocessed_with_labels.pt', ["summed_travel_time", "summed_length", "mean_curvature", "abs_mean_grade", "roadType_share"], 64, 4, "train")