Spaces:
Running
Running
File size: 2,755 Bytes
abd15df |
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 |
import pytorch_lightning as L
from torch.utils.data import DataLoader, random_split
import torch
import time
class ImageDataModule(L.LightningDataModule):
def __init__(
self,
train_dataset,
val_dataset,
test_dataset,
global_batch_size,
num_workers,
num_nodes=1,
num_devices=1,
val_proportion=0.1,
):
super().__init__()
self._builders = {
"train": train_dataset,
"val": val_dataset,
"test": test_dataset,
}
self.num_workers = num_workers
self.batch_size = global_batch_size // (num_nodes * num_devices)
print(f"Each GPU will receive {self.batch_size} images")
self.val_proportion = val_proportion
@property
def num_classes(self):
if hasattr(self, "train_dataset"):
return self.train_dataset.num_classes
else:
return self._builders["train"]().num_classes
def setup(self, stage=None):
"""Setup the datamodule.
Args:
stage (str): stage of the datamodule
Is be one of "fit" or "test" or None
"""
print("Stage", stage)
start_time = time.time()
if stage == "fit" or stage is None:
self.train_dataset = self._builders["train"]()
self.val_dataset = self._builders["val"]()
print(f"Train dataset size: {len(self.train_dataset)}")
print(f"Val dataset size: {len(self.val_dataset)}")
else:
self.test_dataset = self._builders["test"]()
print(f"Test dataset size: {len(self.test_dataset)}")
end_time = time.time()
print(f"Setup took {(end_time - start_time):.2f} seconds")
def train_dataloader(self):
return DataLoader(
self.train_dataset,
batch_size=self.batch_size,
shuffle=True,
pin_memory=False,
drop_last=True,
num_workers=self.num_workers,
collate_fn=self.train_dataset.collate_fn_density,
)
def val_dataloader(self):
return DataLoader(
self.val_dataset,
batch_size=self.batch_size,
shuffle=False,
pin_memory=False,
num_workers=self.num_workers,
collate_fn=self.val_dataset.collate_fn,
)
def test_dataloader(self):
return DataLoader(
self.test_dataset,
batch_size=self.batch_size,
shuffle=False,
pin_memory=False,
num_workers=self.num_workers,
collate_fn=self.test_dataset.collate_fn,
)
|