def get_dataloader(self):
img_root_dir = "/data1/yinzi/datas"
part_root_dir = "/data1/yinzi/facial_parts"
root_dir = {
'image': img_root_dir,
'parts': part_root_dir
}
txt_file_names = {
'train': "exemplars.txt",
'val': "tuning.txt"
}
self.dataset = {x: HelenDataset(txt_file=txt_file_names[x],
root_dir=root_dir['image'],
transform=transforms.Compose([
ToPILImage(),
Resize((512, 512)),
ToTensor()
])
)
for x in ['train', 'val']
}
self.data_loader = {x: DataLoader(self.dataset[x], batch_size=self.args.batch_size,
shuffle=True, num_workers=16)
for x in ['train', 'val']
}
transforms_list = {
'train': transforms.Compose([ToTensor(),
Resize((128, 128)),
OrigPad()]),
'val': transforms.Compose([ToTensor(),
Resize((128, 128)),
OrigPad()]),
'test': transforms.Compose([ToTensor(),
Resize((128, 128)),
OrigPad()])
}
# DataLoader
Dataset = {
x: HelenDataset(txt_file=txt_file_names[x],
root_dir=root_dir,
parts_root_dir=parts_root_dir,
transform=transforms_list[x])
for x in ['train', 'val', 'test']
}
dataloader = {
x: DataLoader(Dataset[x],
batch_size=args.batch_size,
shuffle=True,
num_workers=4)
for x in ['train', 'val', 'test']
}
class TrainModel(TemplateModel):
def __init__(self, argus=args):
ToTensor()
# Normalize()
]),
'val':
transforms.Compose([
ToPILImage(),
Resize((64, 64)),
ToTensor()
# Normalize()
])
}
# Stage 1 augmentation
stage1_dataset = {
x: HelenDataset(txt_file=txt_file_names[x],
root_dir=root_dir,
transform=transforms.Compose(
[ToPILImage(),
Resize((64, 64)),
ToTensor()]))
for x in ['train', 'val']
}
stage1_dataloaders = {
x: DataLoader(stage1_dataset[x],
batch_size=args.batch_size,
shuffle=True,
num_workers=4)
for x in ['train', 'val']
}
stage1_dataset_sizes = {x: len(stage1_dataset[x]) for x in ['train', 'val']}