我正在尝试训练GAN着色图像,以便使用Torchvision的ImageFolder将其加载为灰度图像,但是我还需要原始数据以及转换后的数据。
我想要最快的方式,因为数据量很大,所以我想让ImageFolder同时加载两者,以节省更多时间。
def load_data_bw(opt):
datapath = '/content/gdrive/My Drive/faces/2003'
dataset = torchvision.datasets.ImageFolder(datapath,
transform=transforms.Compose([
transforms.Grayscale(num_output_channels=3), #load images as grayscale with three channels
transforms.RandomChoice(
[transforms.Resize(opt['loadSize'], interpolation=1),
transforms.Resize(opt['loadSize'], interpolation=2),
transforms.Resize(opt['loadSize'], interpolation=3),
transforms.Resize((opt['loadSize'], opt['loadSize']),
interpolation=1),
transforms.Resize((opt['loadSize'], opt['loadSize']),
interpolation=2),
transforms.Resize((opt['loadSize'], opt['loadSize']),
interpolation=3)]
),
transforms.RandomChoice(
[transforms.RandomResizedCrop(opt['fineSize'], interpolation=1),
transforms.RandomResizedCrop(opt['fineSize'], interpolation=2),
transforms.RandomResizedCrop(opt['fineSize'], interpolation=3)]
),
transforms.ColorJitter(brightness=0.1, contrast=0.1),
transforms.RandomHorizontalFlip(),
transforms.ToTensor()
]))
return dataset
我希望得到:
for iteration, orig_data, gray_data in enumerate(training_data_loader, 1):
# code..
答案 0 :(得分:1)
我假设您有2个数据集变量,即dataset_bw
和dataset_color
,您可以使用ImageFolder
进行加载。然后,您可以执行以下操作:
class GAN_dataset(Dataset):
def __init__(self, dataset_bw, dataset_color):
self.dataset1 = dataset_bw
self.dataset2 = dataset_color
def __getitem__(self, index):
x1 = self.dataset1[index]
x2 = self.dataset2[index]
return x1, x2
def __len__(self):
return len(self.dataset1)
dataset = GAN_dataset(dataset_bw, dataset_color)
loader = DataLoader(dataset, batch_size = ...)
这样,当您遍历loader
时,将根据需要获得两张图像。