File size: 2,856 Bytes
8e1e3ba | 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 | import lightning as L
import numpy as np
from torch.utils.data import DataLoader
from datasets import load_dataset
from src.dataset import DecoderDataset
class StyleTransferDataModule(L.LightningDataModule):
def __init__(self, batch_size=8, image_size=256, val_split=0.05, test_split=0.05, max_samples=None):
super().__init__()
self.save_hyperparameters()
self.content_name = "wangwangxuebing/train_COCO_data2014_jpg"
self.style_name = "newsletter/HiDream-I1-Artists"
def setup(self, stage=None):
full_content = load_dataset(self.content_name, split="test")
full_style = load_dataset(self.style_name, split="train")
if self.hparams.max_samples:
full_content = full_content.select(range(min(len(full_content), self.hparams.max_samples)))
full_style = full_style.select(range(min(len(full_style), self.hparams.max_samples)))
seed = 42
content_indexes = np.random.RandomState(seed).permutation(len(full_content))
style_indexes = np.random.RandomState(seed).permutation(len(full_style))
def get_splits(all_indexes, val_ratio, test_ratio):
total_count = len(all_indexes)
n_val = int(total_count * val_ratio)
n_test = int(total_count * test_ratio)
n_train = total_count - n_val - n_test
train_end = n_train
val_end = n_train + n_val
train_indices = all_indexes[:train_end]
val_indices = all_indexes[train_end:val_end]
test_indices = all_indexes[val_end:]
return train_indices, val_indices, test_indices
content_train, content_val, content_test = get_splits(content_indexes, self.hparams.val_split, self.hparams.test_split)
style_train, style_val, style_test = get_splits(style_indexes, self.hparams.val_split, self.hparams.test_split)
if stage == "fit" or stage is None:
self.train_ds = DecoderDataset(full_content, full_style, content_train, style_train, self.hparams.image_size, is_train=True)
self.val_ds = DecoderDataset(full_content, full_style, content_val, style_val, self.hparams.image_size, is_train=False)
if stage == "test":
self.test_ds = DecoderDataset(full_content, full_style, content_test, style_test, self.hparams.image_size, is_train=False)
def train_dataloader(self):
return DataLoader(self.train_ds, batch_size=self.hparams.batch_size, shuffle=True, num_workers=12, pin_memory=True)
def val_dataloader(self):
return DataLoader(self.val_ds, batch_size=self.hparams.batch_size, shuffle=False, num_workers=4, pin_memory=True)
def test_dataloader(self):
return DataLoader(self.test_ds, batch_size=self.hparams.batch_size, shuffle=False, num_workers=4) |