Function load_data

utils.py:304–378

(config_dict, client_id=-1, n_clients=50, alpha=1e0, bsize=16, 
              linear_eval=False, hparam_eval=False, in_simulation=False, force_shuffle=False, 
              subset_proportion=1, subset_force_class_balanced=False, subset_seed=0)

######### Dataloaders #########
def load_data(config_dict, client_id=-1, n_clients=50, alpha=1e0, bsize=16, 
              linear_eval=False, hparam_eval=False, in_simulation=False, force_shuffle=False, 
              subset_proportion=1, subset_force_class_balanced=False, subset_seed=0):
    
    da_method = config_dict["da_method"]
    train_mode = config_dict["train_mode"]
    dataset_name = config_dict["dataset"]
    data_dir = config_dict['data_dir']
    
    # Define data augmentations
    if(hparam_eval):
        transform_train = SimCLRTransform(is_sup=False, image_size=32)
    elif(linear_eval):
        transform_train = BaseTransform(is_sup=True, image_size=32)
    elif(da_method=="sup"):
        transform_train = BaseTransform(is_sup=(train_mode=="sup"), image_size=32)
    elif(da_method=="simclr" or da_method=="simsiam"):
        transform_train = SimCLRTransform(is_sup=(train_mode=="sup"), image_size=32)
    elif(da_method=="specloss"):
        transform_train = SpecLossTransform(is_sup=(train_mode=="sup"), image_size=32)
    elif(da_method=="byol"):
        transform_train = BYOLTransform(is_sup=(train_mode=="sup"), image_size=32)
    elif(da_method=="rotpred"):
        transform_train = RotTransform(is_sup=(train_mode=="sup"))
    elif(da_method=="orchestra"):
        transform_train = OrchestraTransform(is_sup=(train_mode=="sup"), image_size=32)

    transform_test = T.Compose([
        T.ToTensor(), 
        T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
    )

    # Load dataset
    if(dataset_name=="CIFAR10"):
        trainset = CIFAR10(f"{data_dir}/dataset/CIFAR10", train=True, download=False, transform=transform_train)
        memset = CIFAR10(f"{data_dir}/dataset/CIFAR10", train=True, download=False, transform=transform_test)
        testset = CIFAR10(f"{data_dir}/dataset/CIFAR10", train=False, download=False, transform=transform_test)
    elif(dataset_name=="CIFAR100"):
        trainset = CIFAR100(f"{data_dir}/dataset/CIFAR100", train=True, download=False, transform=transform_train)
        memset = CIFAR100(f"{data_dir}/dataset/CIFAR100", train=True, download=False, transform=transform_test)
        testset = CIFAR100(f"{data_dir}/dataset/CIFAR100", train=False, download=False, transform=transform_test)
    else:
        raise Exception("Dataset not recognized")

    # Dataloaders for given client
    if(client_id > -1):
        with open(f'{data_dir}/{n_clients}/{alpha}/{dataset_name}/train/' +dataset_name+"_"+str(client_id)+".pkl", "rb") as f:
            train_ids = pkl.load(f).astype(np.int32)
        with open(f'{data_dir}/{n_clients}/{alpha}/{dataset_name}/test/'+dataset_name+"_"+str(client_id)+".pkl", "rb") as f:
            test_ids = pkl.load(f).astype(np.int32)
        # Sanity check
        train_deets, test_deets = np.unique(np.array(trainset.targets)[train_ids], return_counts=True), np.unique(np.array(testset.targets)[test_ids], return_counts=True)

        trainloader = DataLoader(clientDataset(trainset, train_ids), batch_size=bsize, shuffle=True, drop_last=True)            
        memloader = DataLoader(clientDataset(memset, train_ids), batch_size=bsize, shuffle=True, drop_last=True)
        testloader = DataLoader(clientDataset(testset, test_ids), batch_size=bsize, shuffle=False, drop_last=True)

        # Sanity check