Init ptb

PytorchBoot/dataset/__init__.py

@@ -0,0 +1 @@
+from PytorchBoot.dataset.base_dataset import BaseDataset

PytorchBoot/dataset/base_dataset.py

@@ -0,0 +1,44 @@
+from abc import ABC
+
+import numpy as np
+
+from torch.utils.data import Dataset
+from torch.utils.data import DataLoader, Subset
+
+from PytorchBoot.component import Component
+
+class BaseDataset(ABC, Dataset, Component):
+    def __init__(self, config):
+        super(BaseDataset, self).__init__()
+        self.config = config       
+
+    @staticmethod
+    def process_batch(batch, device):
+        for key in batch.keys():
+            if isinstance(batch[key], list):
+                continue
+            batch[key] = batch[key].to(device)
+        return batch
+    
+    def get_collate_fn(self):
+        return None
+    
+    def get_loader(self, shuffle=False):
+        ratio = self.config["ratio"]
+        if ratio > 1 or ratio <= 0:
+            raise ValueError(
+                f"dataset ratio should be between (0,1], found {ratio} in {self.config['name']}"
+            )
+        subset_size = max(1,int(len(self) * ratio))
+        indices = np.random.permutation(len(self))[:subset_size]
+        subset = Subset(self, indices)
+        return DataLoader(
+            subset,
+            batch_size=self.config["batch_size"],
+            num_workers=self.config["num_workers"],
+            shuffle=shuffle,
+            collate_fn=self.get_collate_fn(),
+        )
+    
+    
+