builder.py
1.77 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
import copy
import random
import numpy as np
import torch
from torch.utils.data import DataLoader, Dataset
from torch.utils.data.distributed import DistributedSampler
from utils.registery import DATASET_REGISTRY, COLLATE_FN_REGISTRY
from .collate_fn import base_collate_fn
from .ReconData import ReconData
def build_dataset(cfg):
dataset_cfg = copy.deepcopy(cfg)
try:
dataset_cfg = dataset_cfg['dataset']
except Exception:
raise 'should contain {dataset}!'
train_cfg = copy.deepcopy(dataset_cfg)
val_cfg = copy.deepcopy(dataset_cfg)
train_cfg['args']['anno_file'] = train_cfg['args'].pop('train_anno_file')
train_cfg['args'].pop('val_anno_file')
train_cfg['args']['phase'] = 'train'
val_cfg['args']['anno_file'] = val_cfg['args'].pop('val_anno_file')
val_cfg['args'].pop('train_anno_file')
val_cfg['args']['phase'] = 'val'
train_data = DATASET_REGISTRY.get(cfg['dataset']['name'])(**train_cfg['args'])
val_data = DATASET_REGISTRY.get(cfg['dataset']['name'])(**val_cfg['args'])
return train_data, val_data
def build_dataloader(cfg):
dataloader_cfg = copy.deepcopy(cfg)
try:
dataloader_cfg = cfg['dataloader']
except Exception:
raise 'should contain {dataloader}!'
train_ds, val_ds = build_dataset(cfg)
train_sampler = DistributedSampler(train_ds)
collate_fn = COLLATE_FN_REGISTRY.get(dataloader_cfg.pop('collate_fn'))
train_loader = DataLoader(train_ds,
sampler=train_sampler,
collate_fn=collate_fn,
**dataloader_cfg)
val_loader = DataLoader(val_ds,
collate_fn=collate_fn,
**dataloader_cfg)
return train_loader, val_loader