41 lines
1.4 KiB
Python
41 lines
1.4 KiB
Python
import os.path
|
|
import random
|
|
|
|
import torch
|
|
from torch.utils.data import Dataset
|
|
|
|
|
|
|
|
class RandomNDataset(Dataset):
|
|
def __init__(self, latent_shape=(4, 64, 64), num_classes=1000, selected_classes:list=None, seeds=None, max_num_instances=50000, ):
|
|
self.selected_classes = selected_classes
|
|
if selected_classes is not None:
|
|
num_classes = len(selected_classes)
|
|
max_num_instances = 10*num_classes
|
|
self.num_classes = num_classes
|
|
self.seeds = seeds
|
|
if seeds is not None:
|
|
self.max_num_instances = len(seeds)*num_classes
|
|
self.num_seeds = len(seeds)
|
|
else:
|
|
self.num_seeds = (max_num_instances + num_classes - 1) // num_classes
|
|
self.max_num_instances = self.num_seeds*num_classes
|
|
|
|
self.latent_shape = latent_shape
|
|
|
|
|
|
def __getitem__(self, idx):
|
|
label = idx // self.num_seeds
|
|
if self.selected_classes:
|
|
label = self.selected_classes[label]
|
|
seed = random.randint(0, 1<<31) #idx % self.num_seeds
|
|
if self.seeds is not None:
|
|
seed = self.seeds[idx % self.num_seeds]
|
|
|
|
# cls_dir = os.path.join(self.root, f"{label}")
|
|
filename = f"{label}_{seed}.png",
|
|
generator = torch.Generator().manual_seed(seed)
|
|
latent = torch.randn(self.latent_shape, generator=generator, dtype=torch.float32)
|
|
return latent, label, filename
|
|
def __len__(self):
|
|
return self.max_num_instances |