Spaces:
Running
on
L40S
Running
on
L40S
| import random | |
| import numpy as np | |
| from torch.utils.data.dataset import Dataset | |
| from config.config import cfg | |
| class MultipleDatasets(Dataset): | |
| def __init__(self, | |
| dbs, | |
| partition, | |
| make_same_len=True, | |
| total_len=None, | |
| verbose=False): | |
| self.dbs = dbs | |
| self.db_num = len(self.dbs) | |
| self.max_db_data_num = max([len(db) for db in dbs]) | |
| self.db_len_cumsum = np.cumsum([len(db) for db in dbs]) | |
| self.make_same_len = make_same_len | |
| # self.partition = partition | |
| self.partition = {k: v for k, v in sorted(partition.items(), key=lambda item: item[1])} | |
| self.dataset = {} | |
| for db in dbs: | |
| self.dataset.update({db.__class__.__name__: db}) | |
| if verbose: | |
| print('datasets:', [len(self.dbs[i]) for i in range(self.db_num)]) | |
| print( | |
| f'Sample Ratio: {self.partition}') | |
| def __len__(self): | |
| return self.max_db_data_num | |
| def __getitem__(self, index): | |
| p = np.random.rand() | |
| v = list(self.partition.values()) | |
| k = list(self.partition.keys()) | |
| for i,v_i in enumerate(v): | |
| if p<=v_i: | |
| return self.dataset[k[i]][index % len(self.dataset[k[i]])] | |
| import random | |
| import numpy as np | |
| from torch.utils.data.dataset import Dataset | |
| class MultipleDatasets_debug(Dataset): | |
| def __init__(self, dbs, make_same_len=True, total_len=None, verbose=False): | |
| self.dbs = dbs | |
| self.db_num = len(self.dbs) | |
| self.max_db_data_num = max([len(db) for db in dbs]) | |
| self.db_len_cumsum = np.cumsum([len(db) for db in dbs]) | |
| self.make_same_len = make_same_len | |
| if total_len == 'auto': | |
| self.total_len = self.db_len_cumsum[-1] | |
| self.auto_total_len = True | |
| else: | |
| self.total_len = total_len | |
| self.auto_total_len = False | |
| if total_len is not None: | |
| self.per_db_len = self.total_len // self.db_num | |
| if verbose: | |
| print('datasets:', [len(self.dbs[i]) for i in range(self.db_num)]) | |
| print( | |
| f'Auto total length: {self.auto_total_len}, {self.total_len}') | |
| def __len__(self): | |
| # all dbs have the same length | |
| if self.make_same_len: | |
| if self.total_len is None: | |
| # match the longest length | |
| return self.max_db_data_num * self.db_num | |
| else: | |
| # each dataset has the same length and total len is fixed | |
| return self.total_len | |
| else: | |
| # each db has different length, simply concat | |
| return sum([len(db) for db in self.dbs]) | |
| def __getitem__(self, index): | |
| if self.make_same_len: | |
| if self.total_len is None: | |
| # match the longest length | |
| db_idx = index // self.max_db_data_num | |
| data_idx = index % self.max_db_data_num | |
| if data_idx >= len(self.dbs[db_idx]) * ( | |
| self.max_db_data_num // | |
| len(self.dbs[db_idx])): # last batch: random sampling | |
| data_idx = random.randint(0, len(self.dbs[db_idx]) - 1) | |
| else: # before last batch: use modular | |
| data_idx = data_idx % len(self.dbs[db_idx]) | |
| else: | |
| db_idx = index // self.per_db_len | |
| data_idx = index % self.per_db_len | |
| if db_idx > (self.db_num - 1): | |
| # last batch: randomly choose one dataset | |
| db_idx = random.randint(0, self.db_num - 1) | |
| if len(self.dbs[db_idx]) < self.per_db_len and \ | |
| data_idx >= len(self.dbs[db_idx]) * (self.per_db_len // len(self.dbs[db_idx])): | |
| # last batch: random sampling in this dataset | |
| data_idx = random.randint(0, len(self.dbs[db_idx]) - 1) | |
| else: | |
| # before last batch: use modular | |
| data_idx = data_idx % len(self.dbs[db_idx]) | |
| else: | |
| for i in range(self.db_num): | |
| if index < self.db_len_cumsum[i]: | |
| db_idx = i | |
| break | |
| if db_idx == 0: | |
| data_idx = index | |
| else: | |
| data_idx = index - self.db_len_cumsum[db_idx - 1] | |
| return self.dbs[db_idx][data_idx] | |