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()
else:
if DATA_SET == 'Mnist':
test_loader = Mnist_noniid().get_test_data()
if DATA_SET == 'Cifar10':
test_loader = Cifar10_noniid().get_test_data()
return test_loader
def get_train_loader(world_size, rank, batch_size):
if IID == True:
if DATASET == 'Mnist':
train_loader = Mnist(rank, batch_size).get_train_data()
if DATASET == 'Cifar10':
train_loader = Cifar10(rank, batch_size).get_train_data()
else:
if DATASET == 'Mnist':
train_loader = Mnist_noniid(batch_size, world_size).get_train_data(rank)
if DATASET == 'Cifar10':
train_loader = Cifar10_noniid(batch_size, world_size).get_train_data(rank)
return train_loader
def get_local_data(world_size, rank, batch_size):
logging('enter get local data')
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()
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)
return train_loader
def get_local_data(size, rank, batchsize):
if IID == True:
if DATA_SET == 'Mnist':
train_loader = Mnist(batchsize).get_train_data()
if DATA_SET == 'Cifar10':
train_loader = Cifar10(batchsize).get_train_data()
else:
if DATA_SET == 'Mnist':
train_loader = Mnist_noniid(batchsize, size).get_train_data(rank)
if DATA_SET == 'Cifar10':
train_loader = Cifar10_noniid(batchsize, size).get_train_data(rank)
return train_loader
def get_testset():
if IID == True:
if DATA_SET == 'Mnist':
test_loader = Mnist().get_test_data()
if DATA_SET == 'Cifar10':
test_loader = Cifar10().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()
for step, (b_x, b_y) in enumerate(test_loader):
test_x = b_x
test_y = b_y
return test_x, test_y
