/
my_data.py
37 lines (30 loc) · 1 KB
/
my_data.py
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
import torch
from torch.utils import data
import crispr_attn
import os
import sys
import importlib
config_path = ".".join(["models", sys.argv[1]]) + "." if len(sys.argv) >= 2 else ""
attention_setting = importlib.import_module(config_path+"attention_setting")
class MyDataset(data.Dataset):
'Characterizes a dataset for PyTorch'
def __init__(self, list_IDs, labels):
'Initialization'
self.labels = labels
self.list_IDs = list_IDs
def __len__(self):
'Denotes the total number of samples'
return len(self.list_IDs)
def __getitem__(self, index):
'Generates one sample of data'
# Select sample
ID = self.list_IDs[index]
drug_combine_file = os.path.join(attention_setting.data_folder, str(ID) + '.pt')
# Load data and get label
try:
X = torch.load(drug_combine_file)
except:
crispr_attn.logger.error("Fail to get {}".format(ID))
raise
y = self.labels[ID]
return X, y