import itertools import logging import numpy as np import operator import pickle from tabulate import tabulate from termcolor import colored import torch import torch.utils.data from torch.utils.data import RandomSampler from torch.utils.data.distributed import DistributedSampler from uniperceiver.config import configurable from uniperceiver.utils.comm import get_world_size, get_rank from uniperceiver.utils.env import seed_all_rng from uniperceiver.utils.file_io import PathManager from uniperceiver.utils.logger import log_first_n from uniperceiver.utils.registry import Registry from .common import DatasetFromList, MapDataset from uniperceiver.functional import pad_tensor, dict_to_cuda, flat_list_of_lists from .sampler import NodeDistributedSampler from uniperceiver.utils import comm from .sampler import TrainingSampler, NaiveSampler from .moe_embeddings import get_moe_embedding, get_embed_with_task_type, get_embed_with_shared_tagert_name from functools import partial DATASETS_REGISTRY = Registry("DATASETS") # noqa F401 isort:skip DATASETS_REGISTRY.__doc__ = """ Registry for datasets, i.e. the whole model """ from uniperceiver.datasets.unified_dataset import UnifiedDataset from .batch_sampler import WeightedBatchSampler def build_dataset_mapper(cfg, name, stage): dataset_mapper = DATASETS_REGISTRY.get(name)(cfg, stage) return dataset_mapper def trivial_batch_collator(batch): return batch def preprocess_batch_collator(batched_inputs, cfg=dict(), shared_targets=dict()): ret = {} if cfg.MOE.MOE: moe_type = cfg.MOE.MOE_TYPE else: moe_type = None # sample lists for data_name in ['input_sample', 'target_sample']: ret[(data_name + '_list')] = [] num_data = len(batched_inputs[0][data_name]) for i in range(num_data): # All samples in data_list can be either be Tensors or groups (i.e., list of Tensors, [Tensors]). # If the samples in data_list are groups, each element in each group will be padded individually, and then all elements in the same group will be concatenated along axis 1. data_list = [sample[data_name][i]['data'] for sample in batched_inputs] # valid_mask_list = [sample[data_name][i]['valid_mask'] for sample in batched_inputs] modality = batched_inputs[0][data_name][i]['modality'] data_type = batched_inputs[0][data_name][i]['data_type'] sample_info_list = [sample[data_name][i]['sample_info'] for sample in batched_inputs] padding_value = sample_info_list[0].get('padding_value', 0) if isinstance(data_list[0], list): if not batched_inputs[0][data_name][i]['sample_info'].get('sample_alone', False): # some data are concatenated inside one sample, e.g. the caption text part during the training. data_group_size = len(data_list[0]) # padding individually for each element in each group data_, valid_mask_ = zip(*[pad_tensor( tensor=[data_group[idx] for data_group in data_list], padding_value=padding_value, use_mask=True) for idx in range(data_group_size)]) # concatenate all elements in the same group along axis 1 data = torch.cat(data_, dim=1) valid_mask = torch.cat(valid_mask_, dim=1) else: # image-text retrieval may have multi-caption for one image when inference, e.g., MSCOCO caption dataset. data_list = flat_list_of_lists(data_list) data, valid_mask = pad_tensor(tensor=data_list, padding_value=padding_value, use_mask=True) elif isinstance(data_list[0], torch.Tensor): if sample_info_list[0].get('cat_along_first_dim', False): # concatenate data along the first dimention, e.g.: video data data = torch.cat(data_list, dim=0) valid_mask = None else: data, valid_mask = pad_tensor(tensor=data_list, padding_value=padding_value, use_mask=True) # Do we have valid mask that is not caused by padding? AND 1/0 for what? else: raise TypeError if valid_mask is not None and valid_mask.all(): valid_mask = None ret[(data_name + '_list')].append({ 'data': data, 'invalid_mask': 1 - valid_mask if valid_mask is not None else None, 'modality': modality, 'data_type': data_type, 'sample_info': sample_info_list, 'moe_embedding': get_embed_with_task_type(moe_type, batched_inputs[0]['task_info']['task_type'], data_type) }) # target sets num_target_sets = len(batched_inputs[0]['target_idx']) # change value to -1 for padding location ret['target_idx_list'] = [ pad_tensor(tensor=[sample['target_idx'][i] for sample in batched_inputs], padding_value=-1, use_mask=False) if isinstance(batched_inputs[0]['target_idx'][i], torch.Tensor) else torch.tensor([sample['target_idx'][i] for sample in batched_inputs] ) for i in range(num_target_sets) ] ret['target_set_list'] = [batched_inputs[0]['target_set'][i] for i in range(num_target_sets)] # shared target sets ret['shared_target_sets'] = {} for k in shared_targets: padding_value = shared_targets[k]['sample_info'].get('padding_value', 0) if isinstance(shared_targets[k]['data'][0], list): data_list = [d[np.random.randint(0, len(d))] for d in shared_targets[k]['data']] # Randomly choose one for each list else: data_list = shared_targets[k]['data'] data, valid_mask = pad_tensor(tensor=data_list, padding_value=padding_value, use_mask=True) if valid_mask.all(): valid_mask = None ret['shared_target_sets'][k] = [{ 'data': data, 'invalid_mask': 1 - valid_mask if valid_mask is not None else None, 'modality': shared_targets[k]['modality'], 'data_type': 'target', 'sample_info': shared_targets[k]['sample_info'], 'moe_embedding': get_embed_with_shared_tagert_name(moe_type, k) }] # task info ret['task_info'] = batched_inputs[0]['task_info'] # should task_name be put into task_info? ret['task_info']['task_name'] = batched_inputs[0].get('task_name', None) return ret def worker_init_reset_seed(worker_id): seed_all_rng(np.random.randint(2 ** 31) + worker_id) def load_pkl_file(filepath): return pickle.load(open(filepath, 'rb'), encoding='bytes') if len(filepath) > 0 else None def load_shared_targets(cfg, stage='train'): shared_targets_cfg = cfg.SHARED_TARGETS shared_targets = {} for shared_target_cfg in shared_targets_cfg: name = shared_target_cfg['NAME'] if (stage != 'train') and (name not in cfg.DATASETS.TARGET_SET): # For validation and test, we build a dataloader for each task / dataset. # Therefore, the dataloader only needs to load its corresponding shared target set. continue # For validation and test, we do not distribute the shared targets distributed = shared_target_cfg['SHARED_TARGETS_CFG']['DISTRIBUTED'] and (stage == 'train') shared_targets[name] = load_pkl_file(shared_target_cfg['SHARED_TARGETS_CFG']['FILE_PATH']) data = shared_targets[name]['data'] if isinstance(data[0], list): max_len = max([len(t) for tl in data for t in tl]) else: max_len = max([len(t) for t in data]) shared_targets[name]['sample_info'] = {'distributed': distributed, 'max_len': max_len} if distributed: world_size = get_world_size() rank = get_rank() total_num = len(shared_targets[name]['data']) local_num = int(np.ceil(total_num / world_size)) # we pad the shared_targets to a value that can be divided by WORLD_SIZE with no remainer, and then slice it if local_num * world_size > total_num: data = data + [data[0] for _ in range(local_num * world_size - total_num)] shared_targets[name]['data'] = data[rank * local_num : (rank + 1) * local_num] # compute the real (unpadded) length of the local slice start_idx = min(rank * local_num, total_num) end_idx = min((rank + 1) * local_num, total_num) shared_targets[name]['sample_info'].update({ 'total_num': total_num, 'local_num': end_idx - start_idx, 'world_size': world_size, 'rank': rank }) return shared_targets def build_unified_train_loader(cfg, task_cfg, model=None): dataset = UnifiedDataset(cfg, task_cfg, stage="train") batchsampler = WeightedBatchSampler(dataset, cfg, task_cfg) shared_targets = load_shared_targets(cfg) dataloader = torch.utils.data.DataLoader( dataset=dataset, batch_sampler=batchsampler, # sampler=sampler, # batch_size=cfg.DATALOADER.TRAIN_BATCH_SIZE, num_workers=cfg.DATALOADER.NUM_WORKERS, collate_fn=partial(preprocess_batch_collator, shared_targets=shared_targets, cfg=cfg), pin_memory=cfg.DATALOADER.PIN_MEM, worker_init_fn=worker_init_reset_seed, # drop_last=True, prefetch_factor=cfg.DATALOADER.PREFETCH_FACTOR, # default: 2 persistent_workers=cfg.DATALOADER.NUM_WORKERS>0) return dataloader def build_standard_train_loader(cfg, model=None): dataset = build_dataset_mapper(cfg, name=cfg.DATASETS.TRAIN, stage="train") if cfg.DATASETS.TRAIN in [ "ImageTextPairDataset", "ImageNet22KDataset", "ImageNetDataset", "VGPretrain", "VideoDataSet", "VQADataset" ]: sampler = TrainingSampler(dataset) elif cfg.DATASETS.TRAIN in ["GeneralCorpusDataset"]: sampler = NaiveSampler(dataset) else: sampler = NodeDistributedSampler( dataset, shuffle=True, num_replicas=comm.get_world_size(), rank=comm.get_rank(), local_rank=comm.get_local_rank(), local_size=comm.get_local_size()) # sampler = TrainingSampler(dataset) dataloader = torch.utils.data.DataLoader( dataset=dataset, sampler=sampler, batch_size=cfg.DATALOADER.TRAIN_BATCH_SIZE, num_workers=cfg.DATALOADER.NUM_WORKERS, collate_fn=partial(preprocess_batch_collator, model=model), pin_memory=cfg.DATALOADER.PIN_MEM, worker_init_fn=worker_init_reset_seed, drop_last=True, persistent_workers=True) return dataloader def _single_modal_dataset(cfg, dataset_mapper=None, *, datalist=None, sampler=None): if len(cfg.DATASETS.TRAIN) > 0: if dataset_mapper is None: dataset_mapper = build_dataset_mapper(cfg, name=cfg.DATASETS.TRAIN, stage="train") if datalist is None: datalist = dataset_mapper.load_data(cfg) else: dataset_mapper = None datalist = None return datalist, dataset_mapper def _train_loader_from_config(cfg, dataset_mapper=None, *, datalist=None, sampler=None, model=None): # xiaoshi: mscoco image captioning: called from defaulttainer, only cfg passed datalist, dataset_mapper = _single_modal_dataset( cfg, dataset_mapper=dataset_mapper, datalist=datalist, sampler=sampler) return { "datalist": datalist, "dataset_mapper": dataset_mapper, "num_workers": cfg.DATALOADER.NUM_WORKERS, "batch_size": cfg.DATALOADER.TRAIN_BATCH_SIZE, "cfg": cfg, "model": model, } def _valtest_loader_from_config(cfg, dataset_mapper=None, *, datalist=None, sampler=None, stage="val"): dataset_names = { "val": cfg.DATASETS.VAL, "test": cfg.DATASETS.TEST, } dataset_name = dataset_names[stage] if len(dataset_name) > 0: if dataset_mapper is None: dataset_mapper = build_dataset_mapper(cfg, name=dataset_name, stage=stage) if datalist is None: datalist = dataset_mapper.load_data(cfg) else: dataset_mapper = None datalist = None if dataset_name in ['Flickr30kDatasetForSingleStreamVal', 'Flickr30kDatasetForSingleStreamValV2']: multi_gpu_eval = True batch_size = 1 else: multi_gpu_eval = False batch_size = cfg.DATALOADER.TEST_BATCH_SIZE return { "datalist": datalist, "dataset_mapper": dataset_mapper, "num_workers": cfg.DATALOADER.NUM_WORKERS, "batch_size": batch_size, "multi_gpu_eval": multi_gpu_eval, "cfg": cfg, "stage": stage } def build_standard_valtest_loader(cfg, task_cfg, stage, multi_gpu_eval): dataset_names = { "val": cfg.DATASETS.VAL, "test": cfg.DATASETS.TEST, } dataset_name = dataset_names[stage] if len(dataset_name) > 0: dataset = build_dataset_mapper(cfg, name=dataset_name, stage=stage) else: return None shared_targets = load_shared_targets(cfg, stage=stage) if multi_gpu_eval and get_world_size() > 1: # multi-gpu-eval for single stream retrieval sampler = DistributedSampler(dataset, shuffle=True) data_loader = torch.utils.data.DataLoader( dataset, batch_size=cfg.DATALOADER.TEST_BATCH_SIZE, num_workers=cfg.DATALOADER.NUM_WORKERS, drop_last=False, sampler=sampler, collate_fn=partial(preprocess_batch_collator, shared_targets=shared_targets, cfg=cfg), pin_memory=cfg.DATALOADER.PIN_MEM, ) else: data_loader = torch.utils.data.DataLoader( dataset, batch_size=cfg.DATALOADER.TEST_BATCH_SIZE, num_workers=cfg.DATALOADER.NUM_WORKERS, drop_last=False, shuffle=False, collate_fn=partial(preprocess_batch_collator, shared_targets=shared_targets, cfg=cfg), pin_memory=cfg.DATALOADER.PIN_MEM, ) return data_loader