def create_datasets(self):
transform = transforms.Compose([
transforms.RandomCrop(self.crop_size),
self.net.transform,
])
loaders = []
for dataset in [self.source_dataset, self.target_dataset]:
rot_dataset = RotationWrapper(dataset, transform=transform)
loader = DataLoader(rot_dataset,
batch_size=self.batch_size // 2,
shuffle=True,
num_workers=4)
loaders.append(loader)
self.train_dispenser = JointDispenser(*loaders)
val_loaders = []
for dataset in [self.source_val_dataset, self.target_val_dataset]:
rot_dataset = RotationWrapper(dataset, transform=transform)
loader = DataLoader(rot_dataset,
batch_size=1,
shuffle=True,
num_workers=4)
val_loaders.append(loader)
self.val_loaders = {
'source': val_loaders[0],
'target': val_loaders[1],
}
def create_datasets(self):
transform = transforms.Resize(1024)
crop_transform = self.net.transform
loaders = []
for dataset in [self.source_dataset, self.target_dataset]:
grid_dataset = GridRegressionWrapper(dataset,
stride=self.stride,
transform=transform,
crop_transform=crop_transform)
loader = DataLoader(grid_dataset,
batch_size=self.batch_size // 2,
shuffle=True,
num_workers=4)
loaders.append(loader)
val_loaders = []
for dataset in [self.source_val_dataset, self.target_val_dataset]:
grid_dataset = GridRegressionWrapper(dataset,
stride=self.stride,
transform=transform,
crop_transform=crop_transform)
loader = DataLoader(grid_dataset,
batch_size=1,
shuffle=True,
num_workers=4)
val_loaders.append(loader)
self.train_dispenser = JointDispenser(*loaders)
self.val_loaders = {
'source': val_loaders[0],
'target': val_loaders[1],
}
def create_datasets(self):
transform = transforms.Compose([
self.net.transform,
])
target_transform = transforms.Compose([
to_tensor_raw,
])
collate_fn = AugmentCollate(crop=self.crop_size, flip=True)
train_loaders = []
for train_dataset in [self.source_dataset, self.target_dataset]:
dataset = TransformWrapper(train_dataset,
transform=transform,
target_transform=target_transform)
train_loader = DataLoader(dataset,
batch_size=self.batch_size // 2,
shuffle=True,
num_workers=8,
collate_fn=collate_fn,
drop_last=True)
train_loaders.append(train_loader)
self.train_dispenser = JointDispenser(*train_loaders)
val_transform = transforms.Compose([
transforms.Resize((1024, 2048)),
self.net.transform,
])
val_target_transform = transforms.Compose([
transforms.Resize((1024, 2048), interpolation=Image.NEAREST),
to_tensor_raw,
])
val_loaders = []
for dataset in [self.source_val_dataset, self.target_val_dataset]:
dataset = TransformWrapper(
dataset,
transform=val_transform,
target_transform=val_target_transform,
)
loader = DataLoader(
dataset,
batch_size=torch.cuda.device_count(),
num_workers=4,
)
val_loaders.append(loader)
self.val_loaders = {
'source': val_loaders[0],
'target': val_loaders[1],
}
def create_datasets(self):
crop_transform = self.net.transform
loaders = []
for dataset in [self.source_dataset, self.target_dataset]:
grid_dataset = ContinuousGridRegressionWrapper(
dataset,
crop_size=self.crop_size,
top=self.top,
left=self.left,
bottom=self.bottom,
right=self.right,
normalize=self.normalize,
crop_transform=crop_transform
)
loader = DataLoader(
grid_dataset,
batch_size=self.batch_size // 2,
shuffle=True,
num_workers=4,
drop_last=True
)
loaders.append(loader)
self.train_dispenser = JointDispenser(*loaders)
val_loaders = []
for dataset in [self.source_val_dataset, self.target_val_dataset]:
grid_dataset = ContinuousGridRegressionWrapper(
dataset,
crop_size=self.crop_size,
top=self.top,
left=self.left,
bottom=self.bottom,
right=self.right,
normalize=self.normalize,
crop_transform=crop_transform
)
loader = DataLoader(
grid_dataset,
batch_size=4,
shuffle=True,
num_workers=4
)
val_loaders.append(loader)
self.val_loaders = {
'source': val_loaders[0],
'target': val_loaders[1],
}