import time import paddle import numbers import numpy as np from multiprocessing import Manager from paddle.imperative import ParallelEnv from paddle.incubate.hapi.distributed import DistributedBatchSampler from ..utils.registry import Registry DATASETS = Registry("DATASETS") class DictDataset(paddle.io.Dataset): def __init__(self, dataset): self.dataset = dataset self.tensor_keys_set = set() self.non_tensor_keys_set = set() self.non_tensor_dict = Manager().dict() single_item = dataset[0] self.keys = single_item.keys() for k, v in single_item.items(): if not isinstance(v, (numbers.Number, np.ndarray)): setattr(self, k, Manager().dict()) self.non_tensor_keys_set.add(k) else: self.tensor_keys_set.add(k) def __getitem__(self, index): ori_map = self.dataset[index] tmp_list = [] for k, v in ori_map.items(): if isinstance(v, (numbers.Number, np.ndarray)): tmp_list.append(v) else: getattr(self, k).update({index: v}) tmp_list.append(index) return tuple(tmp_list) def __len__(self): return len(self.dataset) def reset(self): for k in self.non_tensor_keys_set: setattr(self, k, Manager().dict()) class DictDataLoader(): def __init__(self, dataset, batch_size, is_train, num_workers=4): self.dataset = DictDataset(dataset) place = paddle.fluid.CUDAPlace(ParallelEnv().dev_id) \ if ParallelEnv().nranks > 1 else paddle.fluid.CUDAPlace(0) sampler = DistributedBatchSampler( self.dataset, batch_size=batch_size, shuffle=True if is_train else False, drop_last=True if is_train else False) self.dataloader = paddle.io.DataLoader( self.dataset, batch_sampler=sampler, places=place, num_workers=num_workers) self.batch_size = batch_size def __iter__(self): self.dataset.reset() for i, data in enumerate(self.dataloader): return_dict = {} j = 0 for k in self.dataset.keys: if k in self.dataset.tensor_keys_set: return_dict[k] = data[j] if isinstance(data, (list, tuple)) else data j += 1 else: return_dict[k] = self.get_items_by_indexs(k, data[-1]) yield return_dict def __len__(self): return len(self.dataloader) def get_items_by_indexs(self, key, indexs): if isinstance(indexs, paddle.Variable): indexs = indexs.numpy() current_items = [] items = getattr(self.dataset, key) for index in indexs: current_items.append(items[index]) return current_items def build_dataloader(cfg, is_train=True): dataset = DATASETS.get(cfg.name)(cfg) batch_size = cfg.get('batch_size', 1) num_workers = cfg.get('num_workers', 0) dataloader = DictDataLoader(dataset, batch_size, is_train, num_workers) return dataloader