Pytorch_Simple_Template/data.py at main · Ldoun/Pytorch_Simple_Template · 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
import torch
from torch.utils.data import  Dataset


def load_data(file):
    pass #data loading func

class DataSet(Dataset):
    def __init__(self, file_list, label=None):
        self.file_list = file_list
        self.label = label

    def __len__(self):
        return len(self.file_list)

    def __getitem__(self, index):
        if self.label is None:
            return {'data' : torch.tensor(load_data(self.file_list[index]), dtype=torch.float)}
        else:
            return {'data' : torch.tensor(load_data(self.file_list[index]), dtype=torch.float),
                    'label' : torch.tensor(load_data(self.label[index]), dtype=torch.float)}

class Preload_DataSet(Dataset):
    def __init__(self, file_list, label=None):
        self.file_list = file_list
        self.label = label

        self.data = torch.stack([load_data[file] for file in file_list])

    def __len__(self):
        return len(self.file_list)

    def __getitem__(self, index):
        if self.label is None:
            return {'data' : self.data[index]}
        else:
            return {'data' : self.data[index],
                    'label' : torch.tensor(load_data(self.label[index]), dtype=torch.float)}