Spaces:
Runtime error
Runtime error
File size: 9,228 Bytes
fb53ec8 |
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 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 |
import os
import os.path as osp
import numpy as np
import numpy.random as npr
import torch
import torch.distributed as dist
import torchvision
import copy
import itertools
from ... import sync
from ...cfg_holder import cfg_unique_holder as cfguh
from ...log_service import print_log
import torch.distributed as dist
from multiprocessing import shared_memory
import pickle
import hashlib
import random
class ds_base(torch.utils.data.Dataset):
def __init__(self,
cfg,
loader = None,
estimator = None,
transforms = None,
formatter = None):
self.cfg = cfg
self.load_info = None
self.init_load_info()
self.loader = loader
self.transforms = transforms
self.formatter = formatter
if self.load_info is not None:
load_info_order_by = getattr(self.cfg, 'load_info_order_by', 'default')
if load_info_order_by == 'default':
self.load_info = sorted(self.load_info, key=lambda x:x['unique_id'])
else:
try:
load_info_order_by, reverse = load_info_order_by.split('|')
reverse = reverse == 'reverse'
except:
reverse = False
self.load_info = sorted(
self.load_info, key=lambda x:x[load_info_order_by], reverse=reverse)
load_info_add_idx = getattr(self.cfg, 'load_info_add_idx', True)
if (self.load_info is not None) and load_info_add_idx:
for idx, info in enumerate(self.load_info):
info['idx'] = idx
if estimator is not None:
self.load_info = estimator(self.load_info)
self.try_sample = getattr(self.cfg, 'try_sample', None)
if self.try_sample is not None:
try:
start, end = self.try_sample
except:
start, end = 0, self.try_sample
self.load_info = self.load_info[start:end]
self.repeat = getattr(self.cfg, 'repeat', 1)
pick = getattr(self.cfg, 'pick', None)
if pick is not None:
self.load_info = [i for i in self.load_info if i['filename'] in pick]
#########
# cache #
#########
self.cache_sm = getattr(self.cfg, 'cache_sm', False)
self.cache_cnt = 0
if self.cache_sm:
self.cache_pct = getattr(self.cfg, 'cache_pct', 0)
cache_unique_id = sync.nodewise_sync().random_sync_id()
self.cache_unique_id = hashlib.sha256(pickle.dumps(cache_unique_id)).hexdigest()
self.__cache__(self.cache_pct)
#######
# log #
#######
if self.load_info is not None:
console_info = '{}: '.format(self.__class__.__name__)
console_info += 'total {} unique images, '.format(len(self.load_info))
console_info += 'total {} unique sample. Cached {}. Repeat {} times.'.format(
len(self.load_info), self.cache_cnt, self.repeat)
else:
console_info = '{}: load_info not ready.'.format(self.__class__.__name__)
print_log(console_info)
def init_load_info(self):
# implement by sub class
pass
def __len__(self):
return len(self.load_info)*self.repeat
def __cache__(self, pct):
if pct == 0:
self.cache_cnt = 0
return
self.cache_cnt = int(len(self.load_info)*pct)
if not self.cache_sm:
for i in range(self.cache_cnt):
self.load_info[i] = self.loader(self.load_info[i])
return
for i in range(self.cache_cnt):
shm_name = str(self.load_info[i]['unique_id']) + '_' + self.cache_unique_id
if i % self.local_world_size == self.local_rank:
data = pickle.dumps(self.loader(self.load_info[i]))
datan = len(data)
# self.print_smname_to_file(shm_name)
shm = shared_memory.SharedMemory(
name=shm_name, create=True, size=datan)
shm.buf[0:datan] = data[0:datan]
shm.close()
self.load_info[i] = shm_name
else:
self.load_info[i] = shm_name
dist.barrier()
def __getitem__(self, idx):
idx = idx%len(self.load_info)
# element = copy.deepcopy(self.load_info[idx])
# 0730 try shared memory
element = copy.deepcopy(self.load_info[idx])
if isinstance(element, str):
shm = shared_memory.SharedMemory(name=element)
element = pickle.loads(shm.buf)
shm.close()
else:
element = copy.deepcopy(element)
element['load_info_ptr'] = self.load_info
if idx >= self.cache_cnt:
element = self.loader(element)
if self.transforms is not None:
element = self.transforms(element)
if self.formatter is not None:
return self.formatter(element)
else:
return element
# 0730 try shared memory
def __del__(self):
# Clean the shared memory
for infoi in self.load_info:
if isinstance(infoi, str) and (self.local_rank==0):
shm = shared_memory.SharedMemory(name=infoi)
shm.close()
shm.unlink()
def print_smname_to_file(self, smname):
try:
log_file = cfguh().cfg.train.log_file
except:
try:
log_file = cfguh().cfg.eval.log_file
except:
raise ValueError
# a trick to use the log_file path
sm_file = log_file.replace('.log', '.smname')
with open(sm_file, 'a') as f:
f.write(smname + '\n')
def singleton(class_):
instances = {}
def getinstance(*args, **kwargs):
if class_ not in instances:
instances[class_] = class_(*args, **kwargs)
return instances[class_]
return getinstance
from .ds_loader import get_loader
from .ds_transform import get_transform
from .ds_estimator import get_estimator
from .ds_formatter import get_formatter
@singleton
class get_dataset(object):
def __init__(self):
self.dataset = {}
def register(self, ds):
self.dataset[ds.__name__] = ds
def __call__(self, cfg):
if cfg is None:
return None
t = cfg.type
if t is None:
return None
elif t in ['laion2b', 'laion2b_dummy',
'laion2b_webdataset',
'laion2b_webdataset_sdofficial', ]:
from .. import ds_laion2b
elif t in ['coyo', 'coyo_dummy',
'coyo_webdataset', ]:
from .. import ds_coyo_webdataset
elif t in ['laionart', 'laionart_dummy',
'laionart_webdataset', ]:
from .. import ds_laionart
elif t in ['celeba']:
from .. import ds_celeba
elif t in ['div2k']:
from .. import ds_div2k
elif t in ['pafc']:
from .. import ds_pafc
elif t in ['coco_caption']:
from .. import ds_coco
else:
raise ValueError
loader = get_loader() (cfg.get('loader' , None))
transform = get_transform()(cfg.get('transform', None))
estimator = get_estimator()(cfg.get('estimator', None))
formatter = get_formatter()(cfg.get('formatter', None))
return self.dataset[t](
cfg, loader, estimator,
transform, formatter)
def register():
def wrapper(class_):
get_dataset().register(class_)
return class_
return wrapper
# some other helpers
class collate(object):
"""
Modified from torch.utils.data._utils.collate
It handle list different from the default.
List collate just by append each other.
"""
def __init__(self):
self.default_collate = \
torch.utils.data._utils.collate.default_collate
def __call__(self, batch):
"""
Args:
batch: [data, data] -or- [(data1, data2, ...), (data1, data2, ...)]
This function will not be used as induction function
"""
elem = batch[0]
if not (elem, (tuple, list)):
return self.default_collate(batch)
rv = []
# transposed
for i in zip(*batch):
if isinstance(i[0], list):
if len(i[0]) != 1:
raise ValueError
try:
i = [[self.default_collate(ii).squeeze(0)] for ii in i]
except:
pass
rvi = list(itertools.chain.from_iterable(i))
rv.append(rvi) # list concat
else:
rv.append(self.default_collate(i))
return rv
|