class SimpleDataManager(DataManager):
def __init__(self, image_size, batch_size):
super(SimpleDataManager, self).__init__()
self.batch_size = batch_size
#self.trans_loader = TransformLoader(image_size)
self.trans_loader = MiniImTransformLoader(image_size)
def get_data_loader(self,
aug,
n_support=1,
n_query=1,
no_aug_support=False,
no_aug_query=False
): #parameters that would change on train/val set
transform = self.trans_loader.get_composed_transform(aug)
original_transform = self.trans_loader.get_composed_transform(
aug=False)
dataset = SimpleDataset(transform=transform,
original_transform=original_transform,
n_support=n_support,
n_query=n_query,
no_aug_support=no_aug_support,
no_aug_query=no_aug_query)
data_loader_params = dict(batch_size=self.batch_size,
shuffle=True,
num_workers=12,
pin_memory=True)
data_loader = torch.utils.data.DataLoader(dataset,
**data_loader_params)
return data_loader
def __init__(self, image_size, batch_size):
super(SimpleDataManager, self).__init__()
self.batch_size = batch_size
self.trans_loader = MiniImTransformLoader(image_size)