| | 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 |
| |
|
| |
|
| |
|
| |
|
| |
|
