URLBERT/dataloader.py at master · Davidup1/URLBERT · GitHub

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
from torch.utils.data import DataLoader, RandomSampler, SequentialSampler, TensorDataset
import torch


def generate_dataloader(data_list, type_dataloader: str, ddp, batch_size):
    dataset = TensorDataset(data_list[0], data_list[1], data_list[2], data_list[3])

    if ddp:
        if type_dataloader == "train":
            sampler = torch.utils.data.distributed.DistributedSampler(dataset)
            dataloader = DataLoader(
                dataset,
                sampler=sampler,
                batch_size=batch_size
            )
            return dataloader, sampler
        elif type_dataloader == "val":
            sampler = torch.utils.data.distributed.DistributedSampler(dataset)
            dataloader = DataLoader(
                dataset,
                sampler=sampler,
                batch_size=batch_size
            )
            return dataloader, sampler
        else:
            return None
    else:
        if type_dataloader == "train":
            dataloader = DataLoader(
                dataset,
                sampler=RandomSampler(dataset),
                batch_size=batch_size
            )
            return dataloader
        elif type_dataloader == "val":
            dataloader = DataLoader(
                dataset,
                sampler=SequentialSampler(dataset),
                batch_size=batch_size
            )
            return dataloader
        else:
            return None