| from __future__ import absolute_import |
| from __future__ import division |
| from __future__ import print_function |
|
|
| import logging |
| import os |
| import json |
| import pathlib |
| from os.path import basename |
|
|
| from timm.data import create_loader |
| import torch |
| import torch.utils.data |
| import torch.distributed as dist |
| import torchvision.datasets as datasets |
| from torchvision.io import read_image |
| import torch.distributed as dist |
| from pathlib import Path |
| from yacs.config import CfgNode as CN |
|
|
| from ..LangEncoder import build_tokenizer |
|
|
| from .tsv import TSVImageTextDatasetV2 |
| from .tsv import TSVMeta |
| from .transforms import build_transforms |
|
|
| logger = logging.getLogger(__name__) |
|
|
|
|
| def build_dataset(cfg, is_train): |
| if cfg['DATASET']['DATASET'] == 'image_text_pairs_v2': |
| dataset = _build_pairs_dataset_v2(cfg, is_train) |
| else: |
| raise ValueError(f'Unknown dataset: {cfg["DATASET"]["DATASET"]}') |
| return dataset |
|
|
|
|
| def _get_tsv_list(cfg, is_train): |
| tmp_list = [] |
| if is_train and 'TRAIN_TSV_LIST' in cfg['DATASET']: |
| tmp_list = cfg['DATASET']['TRAIN_TSV_LIST'] |
| elif 'TEST_TSV_LIST' in cfg['DATASET']: |
| tmp_list = cfg['DATASET']['TEST_TSV_LIST'] |
|
|
| tsv_list = [] |
| for l in tmp_list: |
| if l.endswith('.list'): |
| with open(l, 'r') as f: |
| tsv_list.extend([i.strip() for i in f]) |
| else: |
| tsv_list.append(l) |
|
|
| logger.info(f'tsv list: {tsv_list}') |
|
|
| return tsv_list |
|
|
|
|
| def _get_token_file(cfg): |
| num_nodes = dist.get_world_size() // torch.cuda.device_count() |
| if isinstance(cfg['DATASET']['TOKEN_FILE'], list): |
| if num_nodes == 1: |
| logger.warning('=> Multi token files are provided, but only one node is used for training') |
| sas_token_file = cfg['DATASET']['TOKEN_FILE'][0] |
| else: |
| rank = dist.get_rank() |
| node_idx = rank // torch.cuda.device_count() |
| num_token_files = len(cfg['DATASET']['TOKEN_FILE']) |
| sas_token_file = cfg['DATASET']['TOKEN_FILE'][node_idx % num_token_files] |
| else: |
| sas_token_file = cfg['DATASET']['TOKEN_FILE'] |
|
|
| sas_token_file = os.path.join(cfg['DATASET']['ROOT'], sas_token_file) |
|
|
| if ( |
| cfg['DATASET']['LOADER'] == 'blobfuse' |
| or not os.path.isfile(sas_token_file) |
| ): |
| sas_token_file = None |
|
|
| return sas_token_file |
|
|
|
|
| def _build_pairs_dataset_v2(cfg, is_train): |
| transforms = build_transforms(cfg, is_train) |
| logger.info('transforms: {}'.format(transforms)) |
|
|
| dataset_name = cfg['DATASET']['TRAIN_SET'] \ |
| if is_train else cfg['DATASET']['TEST_SET'] |
|
|
| tokenobj = build_tokenizer(cfg['LANG_ENCODER']) |
|
|
| if cfg['DATASET']['DATA_FORMAT'] != 'tsv': |
| raise ValueError('Only support tsv format for pairs dataset v2') |
|
|
| tsv_list = _get_tsv_list(cfg, is_train) |
|
|
| if len(tsv_list) > 0: |
| tsv_filenames = sorted( |
| [ |
| os.path.join(cfg['DATASET']['ROOT'], dataset_name, f) |
| for f in tsv_list |
| ] |
| ) |
| else: |
| dataset_path = os.path.join(cfg['DATASET']['ROOT'], dataset_name) |
| tsv_files = Path(dataset_path).glob('**/*.tsv') |
|
|
| tsv_filenames = sorted( |
| [ |
| str(path) |
| for path in tsv_files |
| ] |
| ) |
|
|
| image_tsv_files = [ |
| filename |
| for filename in tsv_filenames |
| if ( |
| 'image-' in basename(filename) |
| or 'image_' in basename(filename) |
| or '_image' in basename(filename) |
| or '-image' in basename(filename) |
| or 'images-' in basename(filename) |
| ) |
| ] |
| text_tsv_files = [ |
| filename |
| for filename in tsv_filenames |
| if ( |
| 'text-' in basename(filename) |
| or 'text_' in basename(filename) |
| or '_text' in basename(filename) |
| or '-text' in basename(filename) |
| or 'texts-' in basename(filename) |
| ) |
| ] |
|
|
| logger.info( |
| "=> found %d/%d tsv file(s) to load.", |
| len(image_tsv_files), len(text_tsv_files) |
| ) |
|
|
| num_captions = 1 \ |
| if is_train else cfg['DATASET'].get('NUM_CAPTIONS', 1) |
| text_format = cfg['DATASET'].get('TEXT_FORMAT', 'json') |
|
|
| sas_token_file = _get_token_file(cfg) |
| logger.info("=> SAS token path: %s", sas_token_file) |
|
|
| metas = [] |
| cfg_data = cfg['DATASET'] |
| if 'CLASSIFICATION_SETS' in cfg_data and 'NUM_CLASSES' in cfg_data: |
| for source, num_classes in zip(cfg_data['CLASSIFICATION_SETS'], cfg_data['NUM_CLASSES']): |
| metas.append( |
| TSVMeta( |
| source=source, |
| num_classes=num_classes, |
| task='classification' |
| ) |
| ) |
| logger.info('=> add meta: {}'.format(metas[-1])) |
|
|
| if 'coco-caption' in dataset_name: |
| logger.info('=> coco caption data is used') |
| logger.info('=> update num_captions: 5, text_format: json') |
| logger.warning('=> set sas token to None for coco evaluation') |
| sas_token_file = None |
| num_captions = 5 |
| text_format = 'json' |
|
|
| dataset = TSVImageTextDatasetV2( |
| image_tsv_files, text_tsv_files, |
| transform=transforms, |
| tokenize=tokenobj, |
| context_length=cfg['LANG_ENCODER']['CONTEXT_LENGTH'], |
| num_captions=num_captions, |
| text_format=text_format, |
| is_train=is_train, |
| sas_token_path=sas_token_file, |
| metas=metas, |
| prompt_engineering=cfg['DATASET'].get('PROMPT_ENGINEERING', True), |
| concat_queries=cfg['DATASET'].get('CONCAT_QUERIES', False) |
| ) |
|
|
| logger.info( |
| "=> %s set size: %d", 'train' |
| if is_train else 'val', len(dataset) |
| ) |
|
|
| return dataset |
|
|
|
|
| def build_dataloader(cfg, is_train=True, distributed=False): |
| dataset = build_dataset(cfg, is_train) |
|
|
| if ( |
| is_train |
| and 'TIMM_AUG' in cfg['AUG'] |
| and cfg['AUG']['TIMM_AUG']['USE_LOADER'] |
| ): |
| logger.info('=> use timm loader for training') |
| timm_cfg = CN(init_dict=cfg['AUG']['TIMM_AUG']) |
| data_loader = create_loader( |
| dataset, |
| input_size=cfg['IMAGE_ENCODER']['IMAGE_SIZE'][0], |
| batch_size=cfg['TRAIN']['BATCH_SIZE_PER_GPU'], |
| is_training=True, |
| use_prefetcher=True, |
| no_aug=False, |
| re_prob=timm_cfg.RE_PROB, |
| re_mode=timm_cfg.RE_MODE, |
| re_count=timm_cfg.RE_COUNT, |
| re_split=timm_cfg.RE_SPLIT, |
| scale=cfg['AUG']['SCALE'], |
| ratio=cfg['AUG']['RATIO'], |
| hflip=timm_cfg.HFLIP, |
| vflip=timm_cfg.VFLIP, |
| color_jitter=timm_cfg.COLOR_JITTER, |
| auto_augment=timm_cfg.AUTO_AUGMENT, |
| num_aug_splits=0, |
| interpolation=cfg['AUG']['INTERPOLATION'], |
| mean=cfg['IMAGE_ENCODER']['IMAGE_MEAN'], |
| std=cfg['IMAGE_ENCODER']['IMAGE_STD'], |
| num_workers=cfg['WORKERS'], |
| distributed=distributed, |
| collate_fn=None, |
| pin_memory=cfg['PIN_MEMORY'], |
| use_multi_epochs_loader=True |
| ) |
| else: |
| if is_train: |
| batch_size_per_gpu = cfg['TRAIN']['BATCH_SIZE_PER_GPU'] |
| shuffle = cfg['TRAIN'].get('SHUFFLE', True) |
| else: |
| batch_size_per_gpu = cfg['TEST']['BATCH_SIZE_PER_GPU'] |
| shuffle = cfg['TEST'].get('SHUFFLE', False) |
|
|
| if distributed or cfg.get('ALWAYS_ENABLE_SAMPLER', False): |
| |
| sampler = torch.utils.data.distributed.DistributedSampler(dataset, shuffle=shuffle) |
| shuffle = False |
| else: |
| sampler = None |
|
|
| data_loader = torch.utils.data.DataLoader( |
| dataset, |
| batch_size=batch_size_per_gpu, |
| shuffle=shuffle, |
| num_workers=cfg['WORKERS'], |
| pin_memory=cfg['PIN_MEMORY'], |
| sampler=sampler, |
| drop_last=True if is_train else False, |
| prefetch_factor=cfg.get('PREFETCH_FACTOR', 2) |
| ) |
|
|
| return data_loader |
|
|
|
|
|
|
|
|
|
|
