File size: 2,725 Bytes
69524d0 |
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 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
import torch
from torch.utils import data
import pytorch_lightning as pl
from torch.utils.data import DataLoader
import numpy as np
import constants as cst
import time
from torch.utils import data
from utils.utils_data import one_hot_encoding_type, tanh_encoding_type
class Dataset(data.Dataset):
"""Characterizes a dataset for PyTorch"""
def __init__(self, x, y, seq_size):
"""Initialization"""
self.seq_size = seq_size
self.length = y.shape[0]
self.x = x
self.y = y
if type(self.x) == np.ndarray:
self.x = torch.from_numpy(x).float()
if type(self.y) == np.ndarray:
self.y = torch.from_numpy(y).long()
self.data = self.x
def __len__(self):
"""Denotes the total number of samples"""
return self.length
def __getitem__(self, i):
input = self.x[i:i+self.seq_size, :]
return input, self.y[i]
class DataModule(pl.LightningDataModule):
def __init__(self, train_set, val_set, batch_size, test_batch_size, is_shuffle_train=True, test_set=None, num_workers=16):
super().__init__()
self.train_set = train_set
self.val_set = val_set
self.test_set = test_set
self.batch_size = batch_size
self.test_batch_size = test_batch_size
self.is_shuffle_train = is_shuffle_train
if train_set.data.device.type != cst.DEVICE: #this is true only when we are using a GPU but the data is still on the CPU
self.pin_memory = True
else:
self.pin_memory = False
self.num_workers = num_workers
def train_dataloader(self):
return DataLoader(
dataset=self.train_set,
batch_size=self.batch_size,
shuffle=self.is_shuffle_train,
pin_memory=self.pin_memory,
drop_last=False,
num_workers=self.num_workers,
persistent_workers=True
)
def val_dataloader(self):
return DataLoader(
dataset=self.val_set,
batch_size=self.test_batch_size,
shuffle=False,
pin_memory=self.pin_memory,
drop_last=False,
num_workers=self.num_workers,
persistent_workers=True
)
def test_dataloader(self):
return DataLoader(
dataset=self.test_set,
batch_size=self.test_batch_size,
shuffle=False,
pin_memory=self.pin_memory,
drop_last=False,
num_workers=self.num_workers,
persistent_workers=True
)
|