def init_ip_config(self):
self.ip_train_loader = get_data_loader(EXP_NAME, data_type="train", batch_size=CLIENT_BATCH_SIZE, shuffle=True,
num_workers=8, user_list=[0], pin_memory=True)
self.ip_test_loader = get_data_loader(EXP_NAME, data_type="test", num_workers=8, pin_memory=True)
ip_optimizer = SGD(self.model.parameters(), lr=INIT_LR)
self.ip_optimizer_wrapper = OptimizerWrapper(self.model, ip_optimizer)
self.ip_control = ControlModule(model=self.model, config=config)
def init_test_loader(self):
self.test_loader = get_data_loader(EXP_NAME,
data_type="test",
num_workers=8,
batch_size=100,
shuffle=False,
pin_memory=True)
def parse_init_extra_params(self, extra_params):
users = extra_params[0]
self.train_loader = get_data_loader(self.exp_config.exp_name, data_type="train",
batch_size=self.exp_config.batch_size, shuffle=True, num_workers=8,
pin_memory=True, user_list=users)
to_sparse = extra_params[1]
if to_sparse:
self.convert_to_sparse()
num_users = 100
num_slices = num_users if args.client_selection else NUM_CLIENTS
server = CIFAR10ItPrServer(args, config, VGG11())
list_models, list_indices = server.init_clients()
sampler = FLSampler(list_indices, MAX_ROUND,
NUM_LOCAL_UPDATES * CLIENT_BATCH_SIZE,
args.client_selection, num_slices)
print("Sampler initialized")
train_loader = get_data_loader(EXP_NAME,
data_type="train",
batch_size=CLIENT_BATCH_SIZE,
shuffle=False,
sampler=sampler,
num_workers=8,
pin_memory=True)
client_list = [
CIFAR10ItPrClient(config, list_models[idx])
for idx in range(NUM_CLIENTS)
]
for client in client_list:
client.init_optimizer()
client.init_train_loader(train_loader)
print(
"All initialized. Experiment is {}. Prune rate = {}. Client selection = {}. Num users = {}. Seed = {}. "
"Max round = {}.".format(EXP_NAME, server.prune_rate,
def init_test_loader(self):
self.test_loader = get_data_loader(EXP_NAME,
data_type="test",
num_workers=8,
pin_memory=True)
def init_test_loader(self):
self.test_loader = get_data_loader(EXP_NAME,
data_type="val",
batch_size=200,
num_workers=8,
pin_memory=True)
class FEMNISTServer(Server):
def get_init_extra_params(self):
self.check_client_to_sparse()
return [([i for i in range(19 * j, 19 * (j + 1) if j != 9 else 193)],
self.client_is_sparse) for j in range(10)]
if __name__ == "__main__":
args = parse_args()
seed, resume, use_adaptive = args.seed, args.resume, args.use_adaptive
torch.manual_seed(seed)
model = Conv2()
test_loader = get_data_loader(config.EXP_NAME,
data_type="test",
num_workers=8,
pin_memory=True)
if use_adaptive:
control_module = ControlModule(model, config=config)
control_scheduler = ControlScheduler(
init_max_dec_diff=config.MAX_DEC_DIFF,
dec_half_life=config.ADJ_HALF_LIFE)
else:
control_module = None
control_scheduler = None
# Initial pruning
ip_start = timer()
print("Starting initial pruning")
server_pruning_rounds = 1000
server_adjust_interval = 50
