| 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.transforms as tvtrans |
| import PIL.Image |
| PIL.Image.MAX_IMAGE_PIXELS = None |
| import math |
| import json |
| import copy |
| import pickle |
| from multiprocessing import shared_memory |
| import time |
| from .common import * |
| from ..log_service import print_log |
|
|
| from lib import visual_service as vis |
| from .. import sync |
|
|
| import webdataset as wds |
|
|
| |
| |
| |
|
|
| @regdataset() |
| class laion2b_dummy(ds_base): |
| def init_load_info(self): |
| self.load_info = [] |
|
|
| @regdataset() |
| class laion2b_webdataset(ds_base): |
| def init_load_info(self): |
| self.load_info = [] |
|
|
| def make_loader(self, batch_size, num_workers, train=True): |
| cfg = self.cfg |
| self.root_dir = cfg.root_dir |
|
|
| interpolation_mode = tvtrans.InterpolationMode.BICUBIC |
| if train: |
| trans = [ |
| tvtrans.Resize(cfg.scale, interpolation=interpolation_mode), |
| tvtrans.RandomCrop(cfg.scale), |
| tvtrans.ToTensor(),] |
| else: |
| trans = [ |
| tvtrans.Resize(cfg.scale, interpolation=interpolation_mode), |
| tvtrans.CenterCrop(cfg.scale), |
| tvtrans.ToTensor(),] |
|
|
| trans = tvtrans.Compose(trans) |
|
|
| trans_dict = {'jpg': trans} |
| postprocess = customized_postprocess |
|
|
| shuffle = cfg.get('shuffle', 10000) |
| shardshuffle = shuffle > 0 |
| node_world_size = sync.get_world_size('node') |
| nodesplitter = wds.shardlists.split_by_node \ |
| if node_world_size==1 else wds.shardlists.single_node_only |
|
|
| tars = [osp.join(self.root_dir, 'data', i) for i in os.listdir(osp.join(self.root_dir, 'data')) |
| if osp.splitext(i)[1]=='.tar'] |
| tars = sorted(tars) |
|
|
| dset = wds.WebDataset( |
| tars, |
| nodesplitter=nodesplitter, |
| shardshuffle=shardshuffle, |
| handler=wds.warn_and_continue).repeat().shuffle(shuffle) |
|
|
| print_log(f'Loading webdataset with {len(dset.pipeline[0].urls)} shards.') |
| self.min_size = cfg.get('min_size', None) |
| self.max_pwatermark = cfg.get('max_pwatermark', None) |
| dset = (dset |
| .select(self.filter_keys) |
| .decode('pil', handler=wds.warn_and_continue) |
| .select(self.filter_size) |
| .map_dict(**trans_dict, handler=wds.warn_and_continue)) |
|
|
| if postprocess is not None: |
| dset = dset.map(postprocess) |
| |
| dset.batched(batch_size, partial=False) |
| |
| loader = wds.WebLoader( |
| dset, |
| batch_size=None, |
| shuffle=False, |
| num_workers=num_workers, ) |
| return loader |
|
|
| def filter_size(self, x): |
| try: |
| valid = True |
| if self.min_size is not None and self.min_size > 1: |
| try: |
| valid = valid and x['json']['original_width'] >= self.min_size and \ |
| x['json']['original_height'] >= self.min_size |
| except Exception: |
| valid = False |
| if self.max_pwatermark is not None and self.max_pwatermark < 1.0: |
| try: |
| valid = valid and x['json']['pwatermark'] <= self.max_pwatermark |
| except Exception: |
| valid = False |
| return valid |
| except Exception: |
| return False |
|
|
| def filter_keys(self, x): |
| try: |
| return ("jpg" in x) and ("txt" in x) |
| except Exception: |
| return False |
|
|
| def train_dataloader(self): |
| return self.make_loader(self.train) |
|
|
| def val_dataloader(self): |
| return self.make_loader(self.validation, train=False) |
|
|
| def test_dataloader(self): |
| return self.make_loader(self.test, train=False) |
|
|
| def customized_postprocess(element): |
| return element['jpg']*2-1, element['txt'], element['__key__'] |
|
|
| def dict_collation_fn(samples, combine_tensors=True, combine_scalars=True): |
| keys = set.intersection(*[set(sample.keys()) for sample in samples]) |
| batched = {key: [] for key in keys} |
|
|
| for s in samples: |
| [batched[key].append(s[key]) for key in batched] |
|
|
| result = {} |
| for key in batched: |
| if isinstance(batched[key][0], (int, float)): |
| if combine_scalars: |
| result[key] = np.array(list(batched[key])) |
| elif isinstance(batched[key][0], torch.Tensor): |
| if combine_tensors: |
| result[key] = torch.stack(list(batched[key])) |
| elif isinstance(batched[key][0], np.ndarray): |
| if combine_tensors: |
| result[key] = np.array(list(batched[key])) |
| else: |
| result[key] = list(batched[key]) |
| return result |
|
|
| |
| |
| |
|
|
| def customized_postprocess_sdofficial(element): |
| return { |
| 'jpg': element['jpg']*2-1, |
| 'txt': element['txt'], } |
|
|
| @regdataset() |
| class laion2b_webdataset_sdofficial(laion2b_webdataset): |
| def make_loader(self, batch_size, num_workers, train=True): |
| cfg = self.cfg |
| self.root_dir = cfg.root_dir |
|
|
| interpolation_mode = tvtrans.InterpolationMode.BICUBIC |
| if train: |
| trans = [ |
| tvtrans.Resize(cfg.scale, interpolation=interpolation_mode), |
| tvtrans.RandomCrop(cfg.scale), |
| tvtrans.ToTensor(),] |
| else: |
| trans = [ |
| tvtrans.Resize(cfg.scale, interpolation=interpolation_mode), |
| tvtrans.CenterCrop(cfg.scale), |
| tvtrans.ToTensor(),] |
|
|
| trans = tvtrans.Compose(trans) |
|
|
| trans_dict = {'jpg': trans} |
| postprocess = customized_postprocess_sdofficial |
|
|
| shuffle = 10000 |
| shardshuffle = shuffle > 0 |
| node_world_size = 1 |
| nodesplitter = wds.shardlists.split_by_node \ |
| if node_world_size==1 else wds.shardlists.single_node_only |
|
|
| tars = [osp.join(self.root_dir, 'data', i) for i in os.listdir(osp.join(self.root_dir, 'data')) |
| if osp.splitext(i)[1]=='.tar'] |
| tars = sorted(tars) |
|
|
| dset = wds.WebDataset( |
| tars, |
| nodesplitter=nodesplitter, |
| shardshuffle=shardshuffle, |
| handler=wds.warn_and_continue).repeat().shuffle(shuffle) |
|
|
| print(f'Loading webdataset with {len(dset.pipeline[0].urls)} shards.') |
| self.min_size = cfg.get('min_size', None) |
| self.max_pwatermark = cfg.get('max_pwatermark', None) |
| dset = (dset |
| .select(self.filter_keys) |
| .decode('pil', handler=wds.warn_and_continue) |
| .select(self.filter_size) |
| .map_dict(**trans_dict, handler=wds.warn_and_continue)) |
|
|
| if postprocess is not None: |
| dset = dset.map(postprocess) |
| |
| dset.batched(batch_size, partial=False, collation_fn=dict_collation_fn) |
| |
| loader = wds.WebLoader( |
| dset, |
| batch_size=None, |
| shuffle=False, |
| num_workers=num_workers, ) |
| return loader |
|
|
