def get_testset(rank):
if IID == True:
if DATA_SET == 'Mnist':
test_loader = Mnist(rank).get_test_data()
if DATA_SET == 'Cifar10':
test_loader = Cifar10(rank).get_test_data()
if DATA_SET == 'KWS':
test_loader = KWS(rank).get_test_data()
else:
if DATA_SET == 'Mnist':
test_loader = Mnist_noniid().get_test_data()
if DATA_SET == 'Cifar10':
test_loader = Cifar10_noniid().get_test_data()
if DATA_SET == 'KWS':
test_loader = KWS_noniid().get_test_data()
return test_loader
def get_local_data(world_size, rank, batch_size):
if IID == True:
if DATA_SET == 'Mnist':
train_loader = Mnist(rank, batch_size).get_train_data()
if DATA_SET == 'Cifar10':
train_loader = Cifar10(rank, batch_size).get_train_data()
if DATA_SET == 'KWS':
train_loader = KWS(rank, batch_size).get_train_data()
else:
if DATA_SET == 'Mnist':
train_loader = Mnist_noniid(batch_size, world_size).get_train_data(rank)
if DATA_SET == 'Cifar10':
train_loader = Cifar10_noniid(batch_size, world_size).get_train_data(rank)
if DATA_SET =='KWS':
train_loader = KWS_noniid(batch_size, world_size).get_train_data(rank)
return train_loader