File size: 8,098 Bytes
3ce19a2 | 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 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 | import numpy as np
import pickle
import os
import torch
from torch.utils.data import TensorDataset, DataLoader
from torchvision.datasets import ImageFolder
import torchvision.transforms as transforms
from sklearn.model_selection import train_test_split
from models import parse_layer_string
class ImageFolderSkipDotDirs(ImageFolder):
"""ImageFolder that ignores dot-prefixed directories (e.g. ``.ada_eval_cache``)."""
@staticmethod
def find_classes(directory):
classes = sorted(
entry.name
for entry in os.scandir(directory)
if entry.is_dir() and not entry.name.startswith(".")
)
if not classes:
raise FileNotFoundError(
f"No non-hidden class subdirectories found under {directory!r}"
)
class_to_idx = {cls_name: i for i, cls_name in enumerate(classes)}
return classes, class_to_idx
def set_up_data(H):
blocks = parse_layer_string(H.dec_blocks)
H.block_res = [s[0] for s in blocks]
H.res = sorted(set([s[0] for s in blocks if s[0] <= H.max_hierarchy]))
shift_loss = -127.5
scale_loss = 1. / 127.5
if H.dataset == 'imagenet32':
trX, vaX, teX = imagenet32(H.data_root)
H.image_size = 32
H.image_channels = 3
shift = -116.2373
scale = 1. / 69.37404
elif H.dataset in ['fewshot', 'fewshot512']:
trX, vaX, teX = few_shot_image_folder(H.data_root, H.image_size)
H.image_channels = 3
shift = -116.2373
scale = 1. / 69.37404
elif H.dataset == 'imagenet64':
trX, vaX, teX = imagenet64(H.data_root)
H.image_size = 64
H.image_channels = 3
shift = -115.92961967
scale = 1. / 69.37404
elif H.dataset == 'ffhq_256':
trX, vaX, teX = ffhq256(H.data_root)
H.image_size = 256
H.image_channels = 3
shift = -112.8666757481
scale = 1. / 69.84780273
elif H.dataset == 'ffhq_1024':
trX, vaX, teX = ffhq1024(H.data_root)
H.image_size = 1024
H.image_channels = 3
shift = -0.4387
scale = 1.0 / 0.2743
shift_loss = -0.5
scale_loss = 2.0
elif H.dataset == 'cifar10':
(trX, _), (vaX, _), (teX, _) = cifar10(H.data_root, one_hot=False)
H.image_size = 32
H.image_channels = 3
shift = -120.63838
scale = 1. / 64.16736
elif H.dataset == 'celebahq256':
# CelebA-HQ-256 is loaded as a flat image folder (PNG/JPG). The data
# root should contain the images directly (or one subdirectory).
H.image_size = 256
trX, vaX, teX = few_shot_image_folder(H.data_root, H.image_size)
H.image_channels = 3
shift = -112.8666757481
scale = 1. / 69.84780273
else:
raise ValueError('unknown dataset: ', H.dataset)
do_low_bit = H.dataset in ['ffhq_256']
if H.test_eval:
print('DOING TEST')
eval_dataset = teX
else:
eval_dataset = vaX
shift = torch.tensor([shift]).cuda().view(1, 1, 1, 1)
scale = torch.tensor([scale]).cuda().view(1, 1, 1, 1)
shift_loss = torch.tensor([shift_loss]).cuda().view(1, 1, 1, 1)
scale_loss = torch.tensor([scale_loss]).cuda().view(1, 1, 1, 1)
if H.dataset == 'ffhq_1024':
train_data = ImageFolder(trX, transforms.ToTensor())
valid_data = ImageFolder(eval_dataset, transforms.ToTensor())
untranspose = True
elif H.dataset not in ['fewshot', 'fewshot512', 'celebahq256']:
train_data = TensorDataset(torch.as_tensor(trX))
valid_data = TensorDataset(torch.as_tensor(eval_dataset))
untranspose = False
else:
train_data = trX
# Load in chunks to avoid OOM when the dataset is large (e.g. 30K CelebA-HQ)
chunks = []
for data_train in DataLoader(train_data, batch_size=min(1000, len(train_data)), num_workers=0):
chunk = (data_train[0] * 255).to(torch.uint8).permute(0, 2, 3, 1)
chunks.append(chunk)
ds = torch.cat(chunks, dim=0)
del chunks
train_data = TensorDataset(ds)
valid_data = train_data
untranspose = False
def preprocess_func(x):
nonlocal shift
nonlocal scale
nonlocal shift_loss
nonlocal scale_loss
nonlocal do_low_bit
nonlocal untranspose
'takes in a data example and returns the preprocessed input'
'as well as the input processed for the loss'
if untranspose:
x[0] = x[0].permute(0, 2, 3, 1)
inp = x[0].cuda(non_blocking=True).float()
inp.mul_(1./127.5).add_(-1)
# out = inp.clone()
# inp.add_(shift).mul_(scale)
# if do_low_bit:
# 5 bits of precision
# out.mul_(1. / 8.).floor_().mul_(8.)
# out.add_(shift_loss).mul_(scale_loss)
return inp, inp
return H, train_data, valid_data, preprocess_func
def mkdir_p(path):
os.makedirs(path, exist_ok=True)
def flatten(outer):
return [el for inner in outer for el in inner]
def unpickle_cifar10(file):
fo = open(file, 'rb')
data = pickle.load(fo, encoding='bytes')
fo.close()
data = dict(zip([k.decode() for k in data.keys()], data.values()))
return data
def few_shot_image_folder(data_root, image_size):
transform_list = [
transforms.Resize((int(image_size), int(image_size))),
transforms.ToTensor(),
]
trans = transforms.Compose(transform_list)
train_data = ImageFolderSkipDotDirs(data_root, trans)
return train_data, train_data, train_data
def imagenet32(data_root):
trX = np.load(os.path.join(data_root, 'imagenet32-train.npy'), mmap_mode='r')
np.random.seed(42)
tr_va_split_indices = np.random.permutation(trX.shape[0])
train = trX[tr_va_split_indices[:-5000]]
valid = trX[tr_va_split_indices[-5000:]]
test = np.load(os.path.join(data_root, 'imagenet32-valid.npy'), mmap_mode='r')
return train, valid, test
def imagenet64(data_root):
trX = np.load(os.path.join(data_root, 'imagenet64-train.npy'), mmap_mode='r')
np.random.seed(42)
tr_va_split_indices = np.random.permutation(trX.shape[0])
train = trX[tr_va_split_indices[:-5000]]
valid = trX[tr_va_split_indices[-5000:]]
test = np.load(os.path.join(data_root, 'imagenet64-valid.npy'), mmap_mode='r') # this is test.
return train, valid, test
def ffhq1024(data_root):
# we did not significantly tune hyperparameters on ffhq-1024, and so simply evaluate on the test set
return os.path.join(data_root, 'ffhq1024/train'), os.path.join(data_root, 'ffhq1024/valid'), os.path.join(data_root, 'ffhq1024/valid')
def ffhq256(data_root):
trX = np.load(os.path.join(data_root, 'ffhq-256.npy'), mmap_mode='r')
np.random.seed(5)
tr_va_split_indices = np.random.permutation(trX.shape[0])
train = trX[tr_va_split_indices[:-7000]]
valid = trX[tr_va_split_indices[-7000:]]
return train, valid, valid
def cifar10(data_root, one_hot=True):
tr_data = [unpickle_cifar10(os.path.join(data_root, 'cifar-10-batches-py/', 'data_batch_%d' % i)) for i in range(1, 6)]
# NumPy 2.x requires an explicit sequence (not a generator) for stacking.
trX = np.vstack([data['data'] for data in tr_data])
trY = np.asarray(flatten([data['labels'] for data in tr_data]))
te_data = unpickle_cifar10(os.path.join(data_root, 'cifar-10-batches-py/', 'test_batch'))
teX = np.asarray(te_data['data'])
teY = np.asarray(te_data['labels'])
trX = trX.reshape(-1, 3, 32, 32).transpose(0, 2, 3, 1)
teX = teX.reshape(-1, 3, 32, 32).transpose(0, 2, 3, 1)
trX, vaX, trY, vaY = train_test_split(trX, trY, test_size=5000, random_state=11172018)
if one_hot:
trY = np.eye(10, dtype=np.float32)[trY]
vaY = np.eye(10, dtype=np.float32)[vaY]
teY = np.eye(10, dtype=np.float32)[teY]
else:
trY = np.reshape(trY, [-1, 1])
vaY = np.reshape(vaY, [-1, 1])
teY = np.reshape(teY, [-1, 1])
return (trX, trY), (vaX, vaY), (teX, teY)
|