Upload 1905 files

data/__init__.py

@@ -8,5 +8,7 @@ from .build.scenario import Scenario
 from .build_cl.build import build_cl_scenario
 from .build_cl.scenario import Scenario as CLScenario
 
+from .build_gen.build import build_scenario_manually_v2 as build_gen_scenario
+from .build_gen.scenario import Scenario as GenScenario
 
 from .datasets.dataset_split import split_dataset

data/build_gen/build.py

@@ -0,0 +1,495 @@
+from typing import Dict, List, Optional, Type, Union
+from ..datasets.ab_dataset import ABDataset
+# from benchmark.data.visualize import visualize_classes_in_object_detection
+# from benchmark.scenario.val_domain_shift import get_val_domain_shift_transform
+from ..dataset import get_dataset
+import copy
+from torchvision.transforms import Compose
+
+from .merge_alias import merge_the_same_meaning_classes
+from ..datasets.registery import static_dataset_registery
+
+
+# some legacy aliases of variables:
+# ignore_classes == discarded classes
+# private_classes == unknown classes in partial / open-set / universal DA
+
+
+def _merge_the_same_meaning_classes(classes_info_of_all_datasets):
+    final_classes_of_all_datasets, rename_map = merge_the_same_meaning_classes(classes_info_of_all_datasets)
+    return final_classes_of_all_datasets, rename_map
+
+
+def _find_ignore_classes_when_sources_as_to_target_b(as_classes: List[List[str]], b_classes: List[str], da_mode):
+    thres = {'da': 3, 'partial_da': 2, 'open_set_da': 1, 'universal_da': 0}[da_mode]
+    
+    from functools import reduce
+    a_classes = reduce(lambda res, cur: res | set(cur), as_classes, set())
+    
+    if set(a_classes) == set(b_classes):
+        # a is equal to b, normal
+        # 1. no ignore classes; 2. match class idx
+        a_ignore_classes, b_ignore_classes = [], []
+    
+    elif set(a_classes) > set(b_classes):
+        # a contains b, partial
+        a_ignore_classes, b_ignore_classes = [], []
+        if thres == 3 or thres == 1: # ignore extra classes in a
+            a_ignore_classes = set(a_classes) - set(b_classes)
+        
+    elif set(a_classes) < set(b_classes):
+        # a is contained by b, open set
+        a_ignore_classes, b_ignore_classes = [], []
+        if thres == 3 or thres == 2: # ignore extra classes in b
+            b_ignore_classes = set(b_classes) - set(a_classes)
+    
+    elif len(set(a_classes) & set(b_classes)) > 0:
+        a_ignore_classes, b_ignore_classes = [], []
+        if thres == 3:
+            a_ignore_classes = set(a_classes) - (set(a_classes) & set(b_classes))
+            b_ignore_classes = set(b_classes) - (set(a_classes) & set(b_classes))
+        elif thres == 2:
+            b_ignore_classes = set(b_classes) - (set(a_classes) & set(b_classes))
+        elif thres == 1:
+            a_ignore_classes = set(a_classes) - (set(a_classes) & set(b_classes))
+    
+    else:
+        return None # a has no intersection with b, none
+    
+    as_ignore_classes = [list(set(a_classes) & set(a_ignore_classes)) for a_classes in as_classes]
+    
+    return as_ignore_classes, list(b_ignore_classes)
+
+
+def _find_private_classes_when_sources_as_to_target_b(as_classes: List[List[str]], b_classes: List[str], da_mode):
+    thres = {'da': 3, 'partial_da': 2, 'open_set_da': 1, 'universal_da': 0}[da_mode]
+    
+    from functools import reduce
+    a_classes = reduce(lambda res, cur: res | set(cur), as_classes, set())
+    
+    if set(a_classes) == set(b_classes):
+        # a is equal to b, normal
+        # 1. no ignore classes; 2. match class idx
+        a_private_classes, b_private_classes = [], []
+    
+    elif set(a_classes) > set(b_classes):
+        # a contains b, partial
+        a_private_classes, b_private_classes = [], []
+        # if thres == 2 or thres == 0: # ignore extra classes in a
+        #     a_private_classes = set(a_classes) - set(b_classes)
+        # if thres == 0: # ignore extra classes in a
+        #     a_private_classes = set(a_classes) - set(b_classes)
+            
+    elif set(a_classes) < set(b_classes):
+        # a is contained by b, open set
+        a_private_classes, b_private_classes = [], []
+        if thres == 1 or thres == 0: # ignore extra classes in b
+            b_private_classes = set(b_classes) - set(a_classes)
+    
+    elif len(set(a_classes) & set(b_classes)) > 0:
+        a_private_classes, b_private_classes = [], []
+        if thres == 0:
+            # a_private_classes = set(a_classes) - (set(a_classes) & set(b_classes))
+            
+            b_private_classes = set(b_classes) - (set(a_classes) & set(b_classes))
+        elif thres == 1:
+            b_private_classes = set(b_classes) - (set(a_classes) & set(b_classes))
+        elif thres == 2:
+            # a_private_classes = set(a_classes) - (set(a_classes) & set(b_classes))
+            pass
+        
+    else:
+        return None # a has no intersection with b, none
+    
+    return list(b_private_classes)
+
+
+class _ABDatasetMetaInfo:
+    def __init__(self, name, classes, task_type, object_type, class_aliases, shift_type):
+        self.name = name
+        self.classes = classes
+        self.class_aliases = class_aliases
+        self.shift_type = shift_type
+        self.task_type = task_type
+        self.object_type = object_type
+        
+        
+def _get_dist_shift_type_when_source_a_to_target_b(a: _ABDatasetMetaInfo, b: _ABDatasetMetaInfo):
+    if b.shift_type is None:
+        return 'Dataset Shifts'
+    
+    if a.name in b.shift_type.keys():
+        return b.shift_type[a.name]
+    
+    mid_dataset_name = list(b.shift_type.keys())[0]
+    mid_dataset_meta_info = _ABDatasetMetaInfo(mid_dataset_name, *static_dataset_registery[mid_dataset_name][1:])
+    
+    return _get_dist_shift_type_when_source_a_to_target_b(a, mid_dataset_meta_info) + ' + ' + list(b.shift_type.values())[0]
+
+    
+def _handle_all_datasets_v2(source_datasets: List[_ABDatasetMetaInfo], target_datasets: List[_ABDatasetMetaInfo], da_mode):
+    
+    # 1. merge the same meaning classes
+    classes_info_of_all_datasets = {
+        d.name: (d.classes, d.class_aliases)
+        for d in source_datasets + target_datasets
+    }
+    final_classes_of_all_datasets, rename_map = _merge_the_same_meaning_classes(classes_info_of_all_datasets)
+    all_datasets_classes = copy.deepcopy(final_classes_of_all_datasets)
+
+    # print(all_datasets_known_classes)
+    
+    # 2. find ignored classes according to DA mode
+    # source_datasets_ignore_classes, target_datasets_ignore_classes = {d.name: [] for d in source_datasets}, \
+    #     {d.name: [] for d in target_datasets}
+    # source_datasets_private_classes, target_datasets_private_classes = {d.name: [] for d in source_datasets}, \
+    #     {d.name: [] for d in target_datasets}
+    target_source_relationship_map = {td.name: {} for td in target_datasets}
+    # source_target_relationship_map = {sd.name: [] for sd in source_datasets}
+    
+    # 1. construct target_source_relationship_map
+    for sd in source_datasets:#sd和td使列表中每一个元素（类）的实例
+        for td in target_datasets:
+            sc = all_datasets_classes[sd.name]
+            tc = all_datasets_classes[td.name]
+            
+            if len(set(sc) & set(tc)) == 0:#只保留有相似类别的源域和目标域
+                continue
+            
+            target_source_relationship_map[td.name][sd.name] = _get_dist_shift_type_when_source_a_to_target_b(sd, td)
+    
+    # print(target_source_relationship_map)
+    # exit()
+    
+    source_datasets_ignore_classes = {}
+    for td_name, v1 in target_source_relationship_map.items():
+        for sd_name, v2 in v1.items():
+            source_datasets_ignore_classes[sd_name + '|' + td_name] = []
+    target_datasets_ignore_classes = {d.name: [] for d in target_datasets}
+    target_datasets_private_classes = {d.name: [] for d in target_datasets}
+    # 保证对于每个目标域上的DA都符合给定的label shift
+    # 所以不同目标域就算对应同一个源域，该源域也可能不相同
+    
+    for td_name, v1 in target_source_relationship_map.items():
+        sd_names = list(v1.keys())
+        
+        sds_classes = [all_datasets_classes[sd_name] for sd_name in sd_names]
+        td_classes = all_datasets_classes[td_name]
+        ss_ignore_classes, t_ignore_classes = _find_ignore_classes_when_sources_as_to_target_b(sds_classes, td_classes, da_mode)#根据DA方式不同产生ignore_classes
+        t_private_classes = _find_private_classes_when_sources_as_to_target_b(sds_classes, td_classes, da_mode)
+        
+        for sd_name, s_ignore_classes in zip(sd_names, ss_ignore_classes):
+            source_datasets_ignore_classes[sd_name + '|' + td_name] = s_ignore_classes
+        target_datasets_ignore_classes[td_name] = t_ignore_classes
+        target_datasets_private_classes[td_name] = t_private_classes
+
+    source_datasets_ignore_classes = {k: sorted(set(v), key=v.index) for k, v in source_datasets_ignore_classes.items()}
+    target_datasets_ignore_classes = {k: sorted(set(v), key=v.index) for k, v in target_datasets_ignore_classes.items()}
+    target_datasets_private_classes = {k: sorted(set(v), key=v.index) for k, v in target_datasets_private_classes.items()}
+    
+    # for k, v in source_datasets_ignore_classes.items():
+    #     print(k, len(v))
+    # print()
+    # for k, v in target_datasets_ignore_classes.items():
+    #     print(k, len(v))
+    # print()
+    # for k, v in target_datasets_private_classes.items():
+    #     print(k, len(v))
+    # print()
+    
+    # print(source_datasets_private_classes, target_datasets_private_classes)
+    # 3. reparse classes idx
+    # 3.1. agg all used classes
+    # all_used_classes = []
+    # all_datasets_private_class_idx_map = {}
+    
+    # source_datasets_classes_idx_map = {}
+    # for td_name, v1 in target_source_relationship_map.items():
+    #     for sd_name, v2 in v1.items():
+    #         source_datasets_classes_idx_map[sd_name + '|' + td_name] = []
+    # target_datasets_classes_idx_map = {}
+    
+    global_idx = 0
+    all_used_classes_idx_map = {}
+    # all_datasets_known_classes = {d: [] for d in final_classes_of_all_datasets.keys()}
+    for dataset_name, classes in all_datasets_classes.items():
+        if dataset_name not in target_datasets_ignore_classes.keys():
+            ignore_classes = [0] * 100000
+            for sn, sic in source_datasets_ignore_classes.items():
+                if sn.startswith(dataset_name):
+                    if len(sic) < len(ignore_classes):
+                        ignore_classes = sic
+        else:
+            ignore_classes = target_datasets_ignore_classes[dataset_name]
+        private_classes = [] \
+            if dataset_name not in target_datasets_ignore_classes.keys() else target_datasets_private_classes[dataset_name]
+        
+        for c in classes:
+            if c not in ignore_classes and c not in all_used_classes_idx_map.keys() and c not in private_classes:
+                all_used_classes_idx_map[c] = global_idx
+                global_idx += 1
+                
+    # print(all_used_classes_idx_map)
+    
+    # dataset_private_class_idx_offset = 0
+    target_private_class_idx = global_idx
+    target_datasets_private_class_idx = {d: None for d in target_datasets_private_classes.keys()}
+    
+    for dataset_name, classes in final_classes_of_all_datasets.items():
+        if dataset_name not in target_datasets_private_classes.keys():
+            continue
+        
+        # ignore_classes = target_datasets_ignore_classes[dataset_name]
+        private_classes = target_datasets_private_classes[dataset_name]
+        # private_classes = [] \
+        #     if dataset_name in source_datasets_private_classes.keys() else target_datasets_private_classes[dataset_name]
+        # for c in classes:
+        #     if c not in ignore_classes and c not in all_used_classes_idx_map.keys() and c in private_classes:
+        #         all_used_classes_idx_map[c] = global_idx + dataset_private_class_idx_offset
+                
+        if len(private_classes) > 0:
+            # all_datasets_private_class_idx[dataset_name] = global_idx + dataset_private_class_idx_offset
+            # dataset_private_class_idx_offset += 1
+            # if dataset_name in source_datasets_private_classes.keys():
+            #     if source_private_class_idx is None:
+            #         source_private_class_idx = global_idx if target_private_class_idx is None else target_private_class_idx + 1
+            #     all_datasets_private_class_idx[dataset_name] = source_private_class_idx
+            # else:
+            #     if target_private_class_idx is None:
+            #         target_private_class_idx = global_idx if source_private_class_idx is None else source_private_class_idx + 1
+            #     all_datasets_private_class_idx[dataset_name] = target_private_class_idx
+            target_datasets_private_class_idx[dataset_name] = target_private_class_idx
+            target_private_class_idx += 1
+            
+            
+    # all_used_classes = sorted(set(all_used_classes), key=all_used_classes.index)
+    # all_used_classes_idx_map = {c: i for i, c in enumerate(all_used_classes)}
+    
+    # print('rename_map', rename_map)
+    
+    # 3.2 raw_class -> rename_map[raw_classes] -> all_used_classes_idx_map
+    all_datasets_e2e_idx_map = {}
+    all_datasets_e2e_class_to_idx_map = {}
+    
+    for td_name, v1 in target_source_relationship_map.items():
+        sd_names = list(v1.keys())
+        sds_classes = [all_datasets_classes[sd_name] for sd_name in sd_names]
+        td_classes = all_datasets_classes[td_name]
+        
+        for sd_name, sd_classes in zip(sd_names, sds_classes):
+            cur_e2e_idx_map = {}
+            cur_e2e_class_to_idx_map = {}
+        
+            for raw_ci, raw_c in enumerate(sd_classes):
+                renamed_c = raw_c if raw_c not in rename_map[dataset_name] else rename_map[dataset_name][raw_c]
+                
+                ignore_classes = source_datasets_ignore_classes[sd_name + '|' + td_name]
+                if renamed_c in ignore_classes:
+                    continue
+                
+                idx = all_used_classes_idx_map[renamed_c]
+                
+                cur_e2e_idx_map[raw_ci] = idx
+                cur_e2e_class_to_idx_map[raw_c] = idx
+                
+            all_datasets_e2e_idx_map[sd_name + '|' + td_name] = cur_e2e_idx_map
+            all_datasets_e2e_class_to_idx_map[sd_name + '|' + td_name] = cur_e2e_class_to_idx_map   
+        cur_e2e_idx_map = {}
+        cur_e2e_class_to_idx_map = {}
+        for raw_ci, raw_c in enumerate(td_classes):
+            renamed_c = raw_c if raw_c not in rename_map[dataset_name] else rename_map[dataset_name][raw_c]
+            
+            ignore_classes = target_datasets_ignore_classes[td_name]
+            if renamed_c in ignore_classes:
+                continue
+            
+            if renamed_c in target_datasets_private_classes[td_name]:
+                idx = target_datasets_private_class_idx[td_name]
+            else:
+                idx = all_used_classes_idx_map[renamed_c]
+            
+            cur_e2e_idx_map[raw_ci] = idx
+            cur_e2e_class_to_idx_map[raw_c] = idx
+            
+        all_datasets_e2e_idx_map[td_name] = cur_e2e_idx_map
+        all_datasets_e2e_class_to_idx_map[td_name] = cur_e2e_class_to_idx_map
+        
+    all_datasets_ignore_classes = {**source_datasets_ignore_classes, **target_datasets_ignore_classes}
+    # all_datasets_private_classes = {**source_datasets_private_classes, **target_datasets_private_classes}
+    
+    classes_idx_set = []
+    for d, m in all_datasets_e2e_class_to_idx_map.items():
+        classes_idx_set += list(m.values())
+    classes_idx_set = set(classes_idx_set)
+    num_classes = len(classes_idx_set)
+
+    return all_datasets_ignore_classes, target_datasets_private_classes, \
+        all_datasets_e2e_idx_map, all_datasets_e2e_class_to_idx_map, target_datasets_private_class_idx, \
+        target_source_relationship_map, rename_map, num_classes
+
+    
+def _build_scenario_info_v2(
+    source_datasets_name: List[str],
+    target_datasets_order: List[str],
+    da_mode: str
+):
+    assert da_mode in ['close_set', 'partial', 'open_set', 'universal']
+    da_mode = {'close_set': 'da', 'partial': 'partial_da', 'open_set': 'open_set_da', 'universal': 'universal_da'}[da_mode]
+    
+    source_datasets_meta_info = [_ABDatasetMetaInfo(d, *static_dataset_registery[d][1:]) for d in source_datasets_name]#获知对应的名字和对应属性，要添加数据集时，直接register就行
+    target_datasets_meta_info = [_ABDatasetMetaInfo(d, *static_dataset_registery[d][1:]) for d in list(set(target_datasets_order))]
+
+    all_datasets_ignore_classes, target_datasets_private_classes, \
+        all_datasets_e2e_idx_map, all_datasets_e2e_class_to_idx_map, target_datasets_private_class_idx, \
+        target_source_relationship_map, rename_map, num_classes \
+        = _handle_all_datasets_v2(source_datasets_meta_info, target_datasets_meta_info, da_mode)
+        
+    return all_datasets_ignore_classes, target_datasets_private_classes, \
+        all_datasets_e2e_idx_map, all_datasets_e2e_class_to_idx_map, target_datasets_private_class_idx, \
+        target_source_relationship_map, rename_map, num_classes
+
+
+def build_scenario_manually_v2(
+    source_datasets_name: List[str],
+    target_datasets_order: List[str],
+    da_mode: str,
+    data_dirs: Dict[str, str],
+    # transforms: Optional[Dict[str, Compose]] = None
+):
+    configs = copy.deepcopy(locals())#返回当前局部变量
+    
+    source_datasets_meta_info = [_ABDatasetMetaInfo(d, *static_dataset_registery[d][1:]) for d in source_datasets_name]
+    target_datasets_meta_info = [_ABDatasetMetaInfo(d, *static_dataset_registery[d][1:]) for d in list(set(target_datasets_order))]
+
+    all_datasets_ignore_classes, target_datasets_private_classes, \
+        all_datasets_e2e_idx_map, all_datasets_e2e_class_to_idx_map, target_datasets_private_class_idx, \
+        target_source_relationship_map, rename_map, num_classes \
+        = _build_scenario_info_v2(source_datasets_name, target_datasets_order, da_mode)
+    # from rich.console import Console
+    # console = Console(width=10000)
+    
+    # def print_obj(_o):
+    #     # import pprint
+    #     # s = pprint.pformat(_o, width=140, compact=True)
+    #     console.print(_o)
+    
+    # console.print('configs:', style='bold red')
+    # print_obj(configs)
+    # console.print('renamed classes:', style='bold red')
+    # print_obj(rename_map)
+    # console.print('discarded classes:', style='bold red')
+    # print_obj(all_datasets_ignore_classes)
+    # console.print('unknown classes:', style='bold red')
+    # print_obj(target_datasets_private_classes)
+    # console.print('class to index map:', style='bold red')
+    # print_obj(all_datasets_e2e_class_to_idx_map)
+    # console.print('index map:', style='bold red')
+    # print_obj(all_datasets_e2e_idx_map)
+    # console = Console()
+    # # console.print('class distribution:', style='bold red')
+    # # class_dist = {
+    # #     k: {
+    # #         '#known classes': len(all_datasets_known_classes[k]),
+    # #         '#unknown classes': len(all_datasets_private_classes[k]),
+    # #         '#discarded classes': len(all_datasets_ignore_classes[k])
+    # #     } for k in all_datasets_ignore_classes.keys()
+    # # }
+    # # print_obj(class_dist)
+    # console.print('corresponding sources of each target:', style='bold red')
+    # print_obj(target_source_relationship_map)
+    
+    # return
+    
+    # res_source_datasets_map = {d: {split: get_dataset(d, data_dirs[d], split, getattr(transforms, d, None),
+    #                                                   all_datasets_ignore_classes[d], all_datasets_e2e_idx_map[d]) 
+    #                                for split in ['train', 'val', 'test']} 
+    #                            for d in source_datasets_name}
+    # res_target_datasets_map = {d: {'train': get_num_limited_dataset(get_dataset(d, data_dirs[d], 'test', getattr(transforms, d, None), 
+    #                                                   all_datasets_ignore_classes[d], all_datasets_e2e_idx_map[d]), 
+    #                                                                 num_samples_in_each_target_domain),
+    #                                'test': get_dataset(d, data_dirs[d], 'test', getattr(transforms, d, None), 
+    #                                                   all_datasets_ignore_classes[d], all_datasets_e2e_idx_map[d])
+    #                                } 
+    #                            for d in list(set(target_datasets_order))}
+    
+    # res_source_datasets_map = {d: {split: get_dataset(d.split('|')[0], data_dirs[d.split('|')[0]], split, 
+    #                                                   getattr(transforms, d.split('|')[0], None),
+    #                                                   all_datasets_ignore_classes[d], all_datasets_e2e_idx_map[d]) 
+    #                                for split in ['train', 'val', 'test']} 
+    #                            for d in all_datasets_ignore_classes.keys() if d.split('|')[0] in source_datasets_name}
+    
+    # from functools import reduce
+    # res_offline_train_source_datasets_map = {}
+    # res_offline_train_source_datasets_map_names = {}
+    
+    # for d in source_datasets_name:
+    #     source_dataset_with_max_num_classes = None
+        
+    #     for ed_name, ed in res_source_datasets_map.items():
+    #         if not ed_name.startswith(d):
+    #             continue
+            
+    #         if source_dataset_with_max_num_classes is None:
+    #             source_dataset_with_max_num_classes = ed
+    #             res_offline_train_source_datasets_map_names[d] = ed_name
+                
+    #         if len(ed['train'].ignore_classes) < len(source_dataset_with_max_num_classes['train'].ignore_classes):
+    #             source_dataset_with_max_num_classes = ed
+    #             res_offline_train_source_datasets_map_names[d] = ed_name
+                
+    #     res_offline_train_source_datasets_map[d] = source_dataset_with_max_num_classes
+
+    # res_target_datasets_map = {d: {split: get_dataset(d, data_dirs[d], split, getattr(transforms, d, None),
+    #                                                   all_datasets_ignore_classes[d], all_datasets_e2e_idx_map[d]) 
+    #                                for split in ['train', 'val', 'test']} 
+    #                            for d in list(set(target_datasets_order))}
+    
+    from .scenario import Scenario, DatasetMetaInfo
+    
+    # test_scenario = Scenario(
+    #     config=configs,
+    #     offline_source_datasets_meta_info={
+    #         d: DatasetMetaInfo(d, 
+    #                            {k: v for k, v in all_datasets_e2e_class_to_idx_map[res_offline_train_source_datasets_map_names[d]].items()}, 
+    #                            None)
+    #         for d in source_datasets_name
+    #     },
+    #     offline_source_datasets={d: res_offline_train_source_datasets_map[d] for d in source_datasets_name},
+
+    #     online_datasets_meta_info=[
+    #         (
+    #             {sd + '|' + d:  DatasetMetaInfo(d, 
+    #                            {k: v for k, v in all_datasets_e2e_class_to_idx_map[sd + '|' + d].items()}, 
+    #                            None)
+    #              for sd in target_source_relationship_map[d].keys()},
+    #             DatasetMetaInfo(d, 
+    #                            {k: v for k, v in all_datasets_e2e_class_to_idx_map[d].items() if k not in target_datasets_private_classes[d]}, 
+    #                            target_datasets_private_class_idx[d])
+    #         )
+    #         for d in target_datasets_order
+    #     ],
+    #     online_datasets={**res_source_datasets_map, **res_target_datasets_map},
+    #     target_domains_order=target_datasets_order,
+    #     target_source_map=target_source_relationship_map,
+    #     num_classes=num_classes
+    # )
+    import os
+    os.environ['_ZQL_NUMC'] = str(num_classes)
+
+    test_scenario = Scenario(config=configs, all_datasets_ignore_classes_map=all_datasets_ignore_classes,
+                             all_datasets_idx_map=all_datasets_e2e_idx_map,
+                             target_domains_order=target_datasets_order,
+                             target_source_map=target_source_relationship_map,
+                             all_datasets_e2e_class_to_idx_map=all_datasets_e2e_class_to_idx_map,
+                             num_classes=num_classes)
+
+
+    return test_scenario
+
+
+if __name__ == '__main__':
+    test_scenario = build_scenario_manually_v2(['CIFAR10', 'SVHN'],
+                               ['STL10', 'MNIST', 'STL10', 'USPS', 'MNIST', 'STL10'],
+                               'close_set')
+    print(test_scenario.num_classes)
+    

data/build_gen/merge_alias.py

@@ -0,0 +1,106 @@
+from re import L
+from typing import Dict, List
+from collections import Counter
+
+    
+def grouping(bondlist):
+    # reference: https://blog.csdn.net/YnagShanwen/article/details/111344386
+    groups = [] 
+    break1 = False    
+    while bondlist:
+        pair1 = bondlist.pop(0)        
+        a = 11111
+        b = 10000
+        while b != a:
+            a = b
+            for atomid in pair1:
+                for i,pair2 in enumerate(bondlist):            
+                    if atomid in pair2:
+                        pair1 = pair1 + pair2
+                        bondlist.pop(i)
+                        if not bondlist:
+                            break1 = True
+                        break
+                if break1:
+                    break
+            b = len(pair1)
+        groups.append(pair1)
+    return groups
+
+
+def build_semantic_class_info(classes: List[str], aliases: List[List[str]]):
+    res = []
+    for c in classes:
+        # print(res)
+        if len(aliases) == 0:
+            res += [[c]]
+        else:
+            find_alias = False
+            for alias in aliases:
+                if c in alias:
+                    res += [alias]
+                    find_alias = True
+                    break
+            if not find_alias:
+                res += [[c]]
+    # print(classes, res)
+    return res
+    
+
+def merge_the_same_meaning_classes(classes_info_of_all_datasets):
+    # print(classes_info_of_all_datasets)
+
+    semantic_classes_of_all_datasets = []
+    all_aliases = []
+    for classes, aliases in classes_info_of_all_datasets.values():
+        all_aliases += aliases
+    for classes, aliases in classes_info_of_all_datasets.values():
+        semantic_classes_of_all_datasets += build_semantic_class_info(classes, all_aliases)
+        
+    # print(semantic_classes_of_all_datasets)
+    
+    grouped_classes_of_all_datasets = grouping(semantic_classes_of_all_datasets)#匹配过后的数据
+   
+    # print(grouped_classes_of_all_datasets)
+    
+    # final_grouped_classes_of_all_datasets = [Counter(c).most_common()[0][0] for c in grouped_classes_of_all_datasets]
+    # use most common class name; if the same common, use shortest class name!
+    final_grouped_classes_of_all_datasets = []
+    for c in grouped_classes_of_all_datasets:
+        counter = Counter(c).most_common()
+        max_times = counter[0][1]
+        candidate_class_names = []
+        for item, times in counter:
+            if times < max_times:
+                break
+            candidate_class_names += [item]
+        candidate_class_names.sort(key=lambda x: len(x))
+        
+        final_grouped_classes_of_all_datasets += [candidate_class_names[0]]
+    res = {}
+    res_map = {d: {} for d in classes_info_of_all_datasets.keys()}
+    
+    for dataset_name, (classes, _) in classes_info_of_all_datasets.items():
+        final_classes = []
+        for c in classes:
+            for grouped_names, final_name in zip(grouped_classes_of_all_datasets, final_grouped_classes_of_all_datasets):
+                if c in grouped_names:
+                    final_classes += [final_name]
+                    if final_name != c:
+                        res_map[dataset_name][c] = final_name
+                    break
+        res[dataset_name] = sorted(set(final_classes), key=final_classes.index)
+    return res, res_map
+
+
+if __name__ == '__main__':
+    cifar10_classes = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
+    cifar10_aliases = [['automobile', 'car']]
+    stl10_classes = ['airplane', 'bird', 'car', 'cat', 'deer', 'dog', 'horse', 'monkey', 'ship', 'truck']
+
+    final_classes_of_all_datasets, rename_map = merge_the_same_meaning_classes({
+        'CIFAR10': (cifar10_classes, cifar10_aliases),
+        'STL10': (stl10_classes, [])
+    })
+
+    print(final_classes_of_all_datasets, rename_map)

data/build_gen/scenario.py

@@ -0,0 +1,473 @@
+import enum
+from functools import reduce
+from typing import Dict, List, Tuple
+import numpy as np
+import copy
+from utils.common.log import logger
+from ..datasets.ab_dataset import ABDataset
+from ..datasets.dataset_split import train_val_split
+from ..dataloader import FastDataLoader, InfiniteDataLoader, build_dataloader
+from data import get_dataset
+
+
+class DatasetMetaInfo:
+    def __init__(self, name, 
+                 known_classes_name_idx_map, unknown_class_idx):
+        
+        assert unknown_class_idx not in known_classes_name_idx_map.keys()
+        
+        self.name = name
+        self.unknown_class_idx = unknown_class_idx
+        self.known_classes_name_idx_map = known_classes_name_idx_map
+        
+    @property
+    def num_classes(self):
+        return len(self.known_classes_idx) + 1
+        
+        
+class MergedDataset:
+    def __init__(self, datasets: List[ABDataset]):
+        self.datasets = datasets
+        self.datasets_len = [len(i) for i in self.datasets]
+        logger.info(f'create MergedDataset: len of datasets {self.datasets_len}')
+        self.datasets_cum_len = np.cumsum(self.datasets_len)
+
+    def __getitem__(self, idx):
+        for i, cum_len in enumerate(self.datasets_cum_len):
+            if idx < cum_len:
+                return self.datasets[i][idx - sum(self.datasets_len[0: i])]
+            
+    def __len__(self):
+        return sum(self.datasets_len)
+    
+    
+class IndexReturnedDataset:
+    def __init__(self, dataset: ABDataset):
+        self.dataset = dataset
+        
+    def __getitem__(self, idx):
+        res = self.dataset[idx]
+
+        if isinstance(res, (tuple, list)):
+            return (*res, idx)
+        else:
+            return res, idx
+            
+    def __len__(self):
+        return len(self.dataset)
+    
+
+# class Scenario:
+#     def __init__(self, config,
+#                  source_datasets_meta_info: Dict[str, DatasetMetaInfo], target_datasets_meta_info: Dict[str, DatasetMetaInfo], 
+#                  target_source_map: Dict[str, Dict[str, str]], 
+#                  target_domains_order: List[str],
+#                  source_datasets: Dict[str, Dict[str, ABDataset]], target_datasets: Dict[str, Dict[str, ABDataset]]):
+        
+#         self.__config = config
+#         self.__source_datasets_meta_info = source_datasets_meta_info
+#         self.__target_datasets_meta_info = target_datasets_meta_info
+#         self.__target_source_map = target_source_map
+#         self.__target_domains_order = target_domains_order
+#         self.__source_datasets = source_datasets
+#         self.__target_datasets = target_datasets
+    
+#     # 1. basic
+#     def get_config(self):
+#         return copy.deepcopy(self.__config)
+    
+#     def get_task_type(self):
+#         return list(self.__source_datasets.values())[0]['train'].task_type
+    
+#     def get_num_classes(self):
+#         known_classes_idx = []
+#         unknown_classes_idx = []
+#         for v in self.__source_datasets_meta_info.values():
+#             known_classes_idx += list(v.known_classes_name_idx_map.values())
+#             unknown_classes_idx += [v.unknown_class_idx]
+#         for v in self.__target_datasets_meta_info.values():
+#             known_classes_idx += list(v.known_classes_name_idx_map.values())
+#             unknown_classes_idx += [v.unknown_class_idx]
+#         unknown_classes_idx = [i for i in unknown_classes_idx if i is not None]
+#         # print(known_classes_idx, unknown_classes_idx)
+#         res = len(set(known_classes_idx)), len(set(unknown_classes_idx)), len(set(known_classes_idx + unknown_classes_idx))
+#         # print(res)
+#         assert res[0] + res[1] == res[2]
+#         return res
+      
+#     def build_dataloader(self, dataset: ABDataset, batch_size: int, num_workers: int, infinite: bool, shuffle_when_finite: bool):
+#         if infinite:
+#             dataloader = InfiniteDataLoader(
+#                 dataset, None, batch_size, num_workers=num_workers)
+#         else:
+#             dataloader = FastDataLoader(
+#                 dataset, batch_size, num_workers, shuffle=shuffle_when_finite)
+
+#         return dataloader
+    
+#     def build_sub_dataset(self, dataset: ABDataset, indexes: List[int]):
+#         from ..data.datasets.dataset_split import _SplitDataset
+#         dataset.dataset = _SplitDataset(dataset.dataset, indexes)
+#         return dataset
+    
+#     def build_index_returned_dataset(self, dataset: ABDataset):
+#         return IndexReturnedDataset(dataset)
+        
+#     # 2. source
+#     def get_source_datasets_meta_info(self):
+#         return self.__source_datasets_meta_info
+    
+#     def get_source_datasets_name(self):
+#         return list(self.__source_datasets.keys())
+    
+#     def get_merged_source_dataset(self, split):
+#         source_train_datasets = {n: d[split] for n, d in self.__source_datasets.items()}
+#         return MergedDataset(list(source_train_datasets.values()))
+    
+#     def get_source_datasets(self, split):
+#         source_train_datasets = {n: d[split] for n, d in self.__source_datasets.items()}
+#         return source_train_datasets
+    
+#     # 3. target **domain**
+#     # (do we need such API `get_ith_target_domain()`?)
+#     def get_target_domains_meta_info(self):
+#         return self.__source_datasets_meta_info
+    
+#     def get_target_domains_order(self):
+#         return self.__target_domains_order
+    
+#     def get_corr_source_datasets_name_of_target_domain(self, target_domain_name):
+#         return self.__target_source_map[target_domain_name]
+    
+#     def get_limited_target_train_dataset(self):
+#         if len(self.__target_domains_order) > 1:
+#             raise RuntimeError('this API is only for pass-in scenario in user-defined online DA algorithm')
+#         return list(self.__target_datasets.values())[0]['train']
+    
+#     def get_target_domains_iterator(self, split):
+#         for target_domain_index, target_domain_name in enumerate(self.__target_domains_order):
+#             target_dataset = self.__target_datasets[target_domain_name]
+#             target_domain_meta_info = self.__target_datasets_meta_info[target_domain_name]
+            
+#             yield target_domain_index, target_domain_name, target_dataset[split], target_domain_meta_info
+    
+#     # 4. permission management
+#     def get_sub_scenario(self, source_datasets_name, source_splits, target_domains_order, target_splits):
+#         def get_split(dataset, splits):
+#             res = {}
+#             for s, d in dataset.items():
+#                 if s in splits:
+#                     res[s] = d
+#             return res
+        
+#         return Scenario(
+#             config=self.__config,
+#             source_datasets_meta_info={k: v for k, v in self.__source_datasets_meta_info.items() if k in source_datasets_name},
+#             target_datasets_meta_info={k: v for k, v in self.__target_datasets_meta_info.items() if k in target_domains_order},
+#             target_source_map={k: v for k, v in self.__target_source_map.items() if k in target_domains_order},
+#             target_domains_order=target_domains_order,
+#             source_datasets={k: get_split(v, source_splits) for k, v in self.__source_datasets.items() if k in source_datasets_name},
+#             target_datasets={k: get_split(v, target_splits) for k, v in self.__target_datasets.items() if k in target_domains_order}
+#         )
+    
+#     def get_only_source_sub_scenario_for_exp_tracker(self):
+#         return self.get_sub_scenario(self.get_source_datasets_name(), ['train', 'val', 'test'], [], [])
+    
+#     def get_only_source_sub_scenario_for_alg(self):
+#         return self.get_sub_scenario(self.get_source_datasets_name(), ['train'], [], [])
+    
+#     def get_one_da_sub_scenario_for_alg(self, target_domain_name):
+#         return self.get_sub_scenario(self.get_corr_source_datasets_name_of_target_domain(target_domain_name), 
+#                                      ['train', 'val'], [target_domain_name], ['train'])
+
+
+# class Scenario:
+#     def __init__(self, config,
+                 
+#                  offline_source_datasets_meta_info: Dict[str, DatasetMetaInfo], 
+#                  offline_source_datasets: Dict[str, ABDataset],
+                 
+#                  online_datasets_meta_info: List[Tuple[Dict[str, DatasetMetaInfo], DatasetMetaInfo]],
+#                  online_datasets: Dict[str, ABDataset],
+#                  target_domains_order: List[str],
+#                  target_source_map: Dict[str, Dict[str, str]],
+                 
+#                  num_classes: int):
+        
+#         self.config = config
+        
+#         self.offline_source_datasets_meta_info = offline_source_datasets_meta_info
+#         self.offline_source_datasets = offline_source_datasets
+        
+#         self.online_datasets_meta_info = online_datasets_meta_info
+#         self.online_datasets = online_datasets
+        
+#         self.target_domains_order = target_domains_order
+#         self.target_source_map = target_source_map
+        
+#         self.num_classes = num_classes
+        
+#     def get_offline_source_datasets(self, split):
+#         return {n: d[split] for n, d in self.offline_source_datasets.items()}
+    
+#     def get_offline_source_merged_dataset(self, split):
+#         return MergedDataset([d[split] for d in self.offline_source_datasets.values()])
+    
+#     def get_online_current_corresponding_source_datasets(self, domain_index, split):
+#         cur_target_domain_name = self.target_domains_order[domain_index]
+#         cur_source_datasets_name = list(self.target_source_map[cur_target_domain_name].keys())
+#         cur_source_datasets = {n: self.online_datasets[n + '|' + cur_target_domain_name][split] for n in cur_source_datasets_name}
+#         return cur_source_datasets
+    
+#     def get_online_current_corresponding_merged_source_dataset(self, domain_index, split):
+#         cur_target_domain_name = self.target_domains_order[domain_index]
+#         cur_source_datasets_name = list(self.target_source_map[cur_target_domain_name].keys())
+#         cur_source_datasets = {n: self.online_datasets[n + '|' + cur_target_domain_name][split] for n in cur_source_datasets_name}
+#         return MergedDataset([d for d in cur_source_datasets.values()])
+    
+#     def get_online_current_target_dataset(self, domain_index, split):
+#         cur_target_domain_name = self.target_domains_order[domain_index]
+#         return self.online_datasets[cur_target_domain_name][split]
+    
+#     def build_dataloader(self, dataset: ABDataset, batch_size: int, num_workers: int, 
+#                          infinite: bool, shuffle_when_finite: bool, to_iterator: bool):
+#         if infinite:
+#             dataloader = InfiniteDataLoader(
+#                 dataset, None, batch_size, num_workers=num_workers)
+#         else:
+#             dataloader = FastDataLoader(
+#                 dataset, batch_size, num_workers, shuffle=shuffle_when_finite)
+            
+#         if to_iterator:
+#             dataloader = iter(dataloader)
+
+#         return dataloader
+    
+#     def build_sub_dataset(self, dataset: ABDataset, indexes: List[int]):
+#         from data.datasets.dataset_split import _SplitDataset
+#         dataset.dataset = _SplitDataset(dataset.dataset, indexes)
+#         return dataset
+    
+#     def build_index_returned_dataset(self, dataset: ABDataset):
+#         return IndexReturnedDataset(dataset)
+    
+#     def get_config(self):
+#         return copy.deepcopy(self.config)
+    
+#     def get_task_type(self):
+#         return list(self.online_datasets.values())[0]['train'].task_type
+    
+#     def get_num_classes(self):
+#         return self.num_classes
+    
+
+class Scenario:
+    def __init__(self, config, all_datasets_ignore_classes_map, all_datasets_idx_map, target_domains_order, target_source_map, 
+                 all_datasets_e2e_class_to_idx_map,
+                 num_classes):
+        self.config = config 
+        self.all_datasets_ignore_classes_map = all_datasets_ignore_classes_map
+        self.all_datasets_idx_map = all_datasets_idx_map
+        self.target_domains_order = target_domains_order
+        self.target_source_map = target_source_map
+        self.all_datasets_e2e_class_to_idx_map = all_datasets_e2e_class_to_idx_map
+        self.num_classes = num_classes
+        self.cur_domain_index = 0
+
+        logger.info(f'[scenario build] # classes: {num_classes}')
+        logger.debug(f'[scenario build] idx map: {all_datasets_idx_map}')
+        
+    def to_json(self):
+        return dict(
+            config=self.config, all_datasets_ignore_classes_map=self.all_datasets_ignore_classes_map,
+            all_datasets_idx_map=self.all_datasets_idx_map, target_domains_order=self.target_domains_order,
+            target_source_map=self.target_source_map, 
+            all_datasets_e2e_class_to_idx_map=self.all_datasets_e2e_class_to_idx_map,
+            num_classes=self.num_classes
+        )
+        
+    def __str__(self):
+        return f'Scenario({self.to_json()})'
+
+    def get_offline_datasets(self, transform=None):
+        # make source datasets which contains all unioned classes
+        res_offline_train_source_datasets_map = {}
+
+        from .. import get_dataset
+        data_dirs = self.config['data_dirs']
+
+        source_datasets_name = self.config['source_datasets_name']
+
+        # ori_datasets_map = {d: get_dataset(d, data_dirs[d], None, None, None, None) for d in source_datasets_name}
+        # res_source_datasets_map = {k: {split: train_val_split(copy.deepcopy(v), split, rate=0.97) for split in ['train', 'val']} for k, v in ori_datasets_map.items()}
+        # for ds in res_source_datasets_map.values():
+        #     for k, v in ds.items():
+        #         v.underlying_dataset.dataset.setSplit(k)
+        res_source_datasets_map = {d: {split: get_dataset(d, data_dirs[d], split, 
+                                                      transform,
+                                                      self.all_datasets_ignore_classes_map[d], self.all_datasets_idx_map[d]) 
+                                   for split in ['train', 'val', 'test']} 
+                               for d in self.all_datasets_ignore_classes_map.keys() if d in source_datasets_name}
+        
+        # for source_dataset_name in self.config['source_datasets_name']:
+        #     source_datasets = [v for k, v in res_source_datasets_map.items() if source_dataset_name in k]
+
+        #     # how to merge idx map?
+        #     # 35 79 97
+        #     idx_maps = [d['train'].idx_map for d in source_datasets]
+        #     ignore_classes_list = [d['train'].ignore_classes for d in source_datasets]
+            
+        #     union_idx_map = {}
+        #     for idx_map in idx_maps:
+        #         for k, v in idx_map.items():
+        #             if k not in union_idx_map:
+        #                 union_idx_map[k] = v
+        #             else:
+        #                 assert union_idx_map[k] == v
+
+        #     union_ignore_classes = reduce(lambda res, cur: res & set(cur), ignore_classes_list, set(ignore_classes_list[0]))
+        #     assert len(union_ignore_classes) + len(union_idx_map) == len(source_datasets[0]['train'].raw_classes)
+
+        #     logger.info(f'[scenario build] {source_dataset_name} has {len(union_idx_map)} classes in offline training')
+            
+        #     d = source_dataset_name
+        #     res_offline_train_source_datasets_map[d] = {split: get_dataset(d, data_dirs[d], split, 
+        #                                               transform,
+        #                                               union_ignore_classes, union_idx_map) 
+        #                            for split in ['train', 'val', 'test']} 
+        
+        return res_source_datasets_map
+    
+    def get_offline_datasets_args(self):
+        # make source datasets which contains all unioned classes
+        res_offline_train_source_datasets_map = {}
+
+        from .. import get_dataset
+        data_dirs = self.config['data_dirs']
+
+        source_datasets_name = self.config['source_datasets_name']
+        res_source_datasets_map = {d: {split: get_dataset(d.split('|')[0], data_dirs[d.split('|')[0]], split, 
+                                                      None,
+                                                      self.all_datasets_ignore_classes_map[d], self.all_datasets_idx_map[d]) 
+                                   for split in ['train', 'val', 'test']} 
+                               for d in self.all_datasets_ignore_classes_map.keys() if d.split('|')[0] in source_datasets_name}
+        
+        for source_dataset_name in self.config['source_datasets_name']:
+            source_datasets = [v for k, v in res_source_datasets_map.items() if source_dataset_name in k]
+
+            # how to merge idx map?
+            # 35 79 97
+            idx_maps = [d['train'].idx_map for d in source_datasets]
+            ignore_classes_list = [d['train'].ignore_classes for d in source_datasets]
+            
+            union_idx_map = {}
+            for idx_map in idx_maps:
+                for k, v in idx_map.items():
+                    if k not in union_idx_map:
+                        union_idx_map[k] = v
+                    else:
+                        assert union_idx_map[k] == v
+
+            union_ignore_classes = reduce(lambda res, cur: res & set(cur), ignore_classes_list, set(ignore_classes_list[0]))
+            assert len(union_ignore_classes) + len(union_idx_map) == len(source_datasets[0]['train'].raw_classes)
+
+            logger.info(f'[scenario build] {source_dataset_name} has {len(union_idx_map)} classes in offline training')
+            
+            d = source_dataset_name
+            res_offline_train_source_datasets_map[d] = {split: dict(d, data_dirs[d], split, 
+                                                      None,
+                                                      union_ignore_classes, union_idx_map) 
+                                   for split in ['train', 'val', 'test']} 
+
+        return res_offline_train_source_datasets_map
+
+        # for d in source_datasets_name:
+        #     source_dataset_with_max_num_classes = None
+            
+        #     for ed_name, ed in res_source_datasets_map.items():
+        #         if not ed_name.startswith(d):
+        #             continue
+                
+        #         if source_dataset_with_max_num_classes is None:
+        #             source_dataset_with_max_num_classes = ed
+        #             res_offline_train_source_datasets_map_names[d] = ed_name
+                    
+        #         if len(ed['train'].ignore_classes) < len(source_dataset_with_max_num_classes['train'].ignore_classes):
+        #             source_dataset_with_max_num_classes = ed
+        #             res_offline_train_source_datasets_map_names[d] = ed_name
+                    
+        #     res_offline_train_source_datasets_map[d] = source_dataset_with_max_num_classes
+
+        # return res_offline_train_source_datasets_map
+        
+    def get_online_ith_domain_datasets_args_for_inference(self, domain_index):
+        target_dataset_name = self.target_domains_order[domain_index]
+        # dataset_name: Any, root_dir: Any, split: Any, transform: Any | None = None, ignore_classes: Any = [], idx_map: Any | None = None
+        
+        if 'MM-CityscapesDet' in self.target_domains_order or 'CityscapesDet' in self.target_domains_order or 'BaiduPersonDet' in self.target_domains_order:
+            logger.info(f'use val split for inference test (only Det workload)')
+            split = 'test'
+        else:
+            split = 'train'
+        
+        return dict(dataset_name=target_dataset_name, 
+                    root_dir=self.config['data_dirs'][target_dataset_name], 
+                    split=split, 
+                    transform=None, 
+                    ignore_classes=self.all_datasets_ignore_classes_map[target_dataset_name], 
+                    idx_map=self.all_datasets_idx_map[target_dataset_name])
+    
+    def get_online_ith_domain_datasets_args_for_training(self, domain_index):
+        target_dataset_name = self.target_domains_order[domain_index]
+        source_datasets_name = list(self.target_source_map[target_dataset_name].keys())
+
+        res = {}
+        # dataset_name: Any, root_dir: Any, split: Any, transform: Any | None = None, ignore_classes: Any = [], idx_map: Any | None = None
+        res[target_dataset_name] = {split: dict(dataset_name=target_dataset_name, 
+                    root_dir=self.config['data_dirs'][target_dataset_name], 
+                    split=split, 
+                    transform=None, 
+                    ignore_classes=self.all_datasets_ignore_classes_map[target_dataset_name], 
+                    idx_map=self.all_datasets_idx_map[target_dataset_name]) for split in ['train', 'val']}
+        for d in source_datasets_name:
+            res[d] = {split: dict(dataset_name=d, 
+                    root_dir=self.config['data_dirs'][d], 
+                    split=split, 
+                    transform=None, 
+                    ignore_classes=self.all_datasets_ignore_classes_map[d + '|' + target_dataset_name], 
+                    idx_map=self.all_datasets_idx_map[d + '|' + target_dataset_name]) for split in ['train', 'val']}
+        
+        return res
+    
+    def get_online_cur_domain_datasets_args_for_inference(self):
+        return self.get_online_ith_domain_datasets_args_for_inference(self.cur_domain_index)
+    
+    def get_online_cur_domain_datasets_args_for_training(self):
+        return self.get_online_ith_domain_datasets_args_for_training(self.cur_domain_index)
+    
+    def get_online_cur_domain_datasets_for_training(self, transform=None):
+        res = {}
+        datasets_args = self.get_online_ith_domain_datasets_args_for_training(self.cur_domain_index)
+        for dataset_name, dataset_args in datasets_args.items():
+            res[dataset_name] = {}
+            for split, args in dataset_args.items():
+                if transform is not None:
+                    args['transform'] = transform
+                dataset = get_dataset(**args)
+                res[dataset_name][split] = dataset
+        return res
+    
+    def get_online_cur_domain_datasets_for_inference(self, transform=None):
+        datasets_args = self.get_online_ith_domain_datasets_args_for_inference(self.cur_domain_index)
+        if transform is not None:
+            datasets_args['transform'] = transform
+        return get_dataset(**datasets_args)
+    
+    def get_online_cur_domain_samples_for_training(self, num_samples, transform=None, collate_fn=None):
+        dataset = self.get_online_cur_domain_datasets_for_training(transform=transform)
+        dataset = dataset[self.target_domains_order[self.cur_domain_index]]['train']
+        return next(iter(build_dataloader(dataset, num_samples, 0, True, None, collate_fn=collate_fn)))[0]
+
+    def next_domain(self):
+        self.cur_domain_index += 1
+        

data/datasets/__init__.py

@@ -4,6 +4,7 @@ from .semantic_segmentation import *
 from .action_recognition import *
 
 from .sentiment_classification import *
+from .text_generation import *
 from .machine_translation import *
 from .pos_tagging import *