55 lines
2.2 KiB
Python
55 lines
2.2 KiB
Python
##################################################
|
|
# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2019 #
|
|
##################################################
|
|
import torch, copy, random
|
|
import torch.utils.data as data
|
|
|
|
|
|
class SearchDataset(data.Dataset):
|
|
def __init__(self, name, data, train_split, valid_split, check=True):
|
|
self.datasetname = name
|
|
if isinstance(data, (list, tuple)): # new type of SearchDataset
|
|
assert len(data) == 2, "invalid length: {:}".format(len(data))
|
|
self.train_data = data[0]
|
|
self.valid_data = data[1]
|
|
self.train_split = train_split.copy()
|
|
self.valid_split = valid_split.copy()
|
|
self.mode_str = "V2" # new mode
|
|
else:
|
|
self.mode_str = "V1" # old mode
|
|
self.data = data
|
|
self.train_split = train_split.copy()
|
|
self.valid_split = valid_split.copy()
|
|
if check:
|
|
intersection = set(train_split).intersection(set(valid_split))
|
|
assert (
|
|
len(intersection) == 0
|
|
), "the splitted train and validation sets should have no intersection"
|
|
self.length = len(self.train_split)
|
|
|
|
def __repr__(self):
|
|
return "{name}(name={datasetname}, train={tr_L}, valid={val_L}, version={ver})".format(
|
|
name=self.__class__.__name__,
|
|
datasetname=self.datasetname,
|
|
tr_L=len(self.train_split),
|
|
val_L=len(self.valid_split),
|
|
ver=self.mode_str,
|
|
)
|
|
|
|
def __len__(self):
|
|
return self.length
|
|
|
|
def __getitem__(self, index):
|
|
assert index >= 0 and index < self.length, "invalid index = {:}".format(index)
|
|
train_index = self.train_split[index]
|
|
valid_index = random.choice(self.valid_split)
|
|
if self.mode_str == "V1":
|
|
train_image, train_label = self.data[train_index]
|
|
valid_image, valid_label = self.data[valid_index]
|
|
elif self.mode_str == "V2":
|
|
train_image, train_label = self.train_data[train_index]
|
|
valid_image, valid_label = self.valid_data[valid_index]
|
|
else:
|
|
raise ValueError("invalid mode : {:}".format(self.mode_str))
|
|
return train_image, train_label, valid_image, valid_label
|