File size: 2,814 Bytes
3b99abb | 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 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | import abc
import torch
import lmdb
import pytorch_lightning as pl
import copy
from torch.utils.data import DataLoader
from tqdm import tqdm
_10TB = 10995116277760
class LMDBDataset(pl.LightningDataModule):
"""
Abstract class from which other datasets inherit. We use LMDB database for all subclasses.
"""
def __init__(self,
train_lmdb: str = None,
valid_lmdb: str = None,
test_lmdb: str = None,
dataloader_kwargs: dict = None):
"""
Args:
train_lmdb: path to train lmdb
valid_lmdb: path to valid lmdb
test_lmdb: path to test lmdb
dataloader_kwargs: kwargs for dataloader
"""
super().__init__()
self.train_lmdb = train_lmdb
self.valid_lmdb = valid_lmdb
self.test_lmdb = test_lmdb
self.dataloader_kwargs = dataloader_kwargs if dataloader_kwargs is not None else {}
self.env = None
self.operator = None
def is_initialized(self):
return self.env is not None
def _init_lmdb(self, path):
if self.env is not None:
self._close_lmdb()
# open lmdb
self.env = lmdb.open(path, lock=False, map_size=_10TB)
self.operator = self.env.begin()
def _close_lmdb(self):
if self.env is not None:
self.env.close()
self.env = None
self.operator = None
def _cursor(self):
return self.operator.cursor()
def _get(self, key: str or int):
value = self.operator.get(str(key).encode())
if value is not None:
value = value.decode()
return value
def _dataloader(self, stage):
self.dataloader_kwargs["shuffle"] = True if stage == "train" else False
lmdb_path = getattr(self, f"{stage}_lmdb")
dataset = copy.copy(self)
dataset._init_lmdb(lmdb_path)
setattr(dataset, "stage", stage)
return DataLoader(dataset, collate_fn=dataset.collate_fn, **self.dataloader_kwargs)
def train_dataloader(self):
return self._dataloader("train")
def test_dataloader(self):
return self._dataloader("test")
def val_dataloader(self):
return self._dataloader("valid")
@abc.abstractmethod
def __getitem__(self, index):
raise NotImplementedError
@abc.abstractmethod
def __len__(self):
raise NotImplementedError
@abc.abstractmethod
def collate_fn(self, batch):
"""
Datasets should implement it as the function will be set when initializing Dataloader
Returns:
inputs: dict
labels: dict
"""
raise NotImplementedError |