def train(self):
    'Main method for training.'
    self.run_stats = {'secsPerClientRound': [], 'secsPerClient': [], 'secsPerClientTraining': [], 'secsPerClientSetup': [], 'secsPerClientFull': [], 'secsPerRoundHousekeeping': [], 'secsPerRoundTotal': [], 'mpiCosts': []}
    run.log('Max iterations', self.max_iteration)
    try:
        (self.worker_trainer.model.cuda() if torch.cuda.is_available() else None)
        eval_list = []
        if (self.cur_iter_no == 0):
            if self.config['server_config']['initial_rec']:
                eval_list.append('test')
            if self.config['server_config']['initial_val']:
                eval_list.append('val')
                run.log('LR for agg. opt.', get_lr(self.worker_trainer.optimizer))
            print_rank('Running {} at itr={}'.format(eval_list, self.cur_iter_no))
            self.metrics = self.evaluation.run(eval_list, self.metrics, metric_logger=run.log)
            eval_list = []
        print_rank('Saving Model Before Starting Training', loglevel=logging.INFO)
        for token in ['best_val_loss', 'best_val_acc', 'best_test_acc', 'latest']:
            self.worker_trainer.save(model_path=self.model_path, token=token, config=self.config['server_config'])
        self.worker_trainer.model.train()
        for i in range(self.cur_iter_no, self.max_iteration):
            begin = time.time()
            metrics_payload = {}

            def log_metric(k, v):
                metrics_payload[k] = v
            print_rank('==== iteration {}'.format(i))
            log_metric('Current iteration', i)
            initial_lr = (self.initial_lr_client * self.lr_weight)
            print_rank('Client learning rate {}'.format(initial_lr))
            self.worker_trainer.model.zero_grad()
            self.train_loss = []
            server_data = (initial_lr, [p.data.to(torch.device('cpu')) for p in self.worker_trainer.model.parameters()])
            if (len(self.num_clients_per_iteration) > 1):
                num_clients_curr_iter = random.randint(self.num_clients_per_iteration[0], self.num_clients_per_iteration[1])
            else:
                num_clients_curr_iter = self.num_clients_per_iteration[0]
            log_metric('Clients for round', num_clients_curr_iter)
            if (self.quant_thresh is not None):
                self.config['client_config']['quant_thresh'] *= self.config['client_config'].get('quant_anneal', 1.0)
                self.quant_thresh = self.config['client_config']['quant_thresh']
                log_metric('Quantization Thresh.', self.config['client_config']['quant_thresh'])
            sampled_idx_clients = (random.sample(self.client_idx_list, num_clients_curr_iter) if (num_clients_curr_iter > 0) else self.client_idx_list)
            sampled_clients = [Client(client_id, self.config, (self.config['client_config']['type'] == 'optimization'), None) for client_id in sampled_idx_clients]
            clients_begin = time.time()
            client_losses = []
            client_mag_grads = []
            client_mean_grads = []
            client_var_grads = []
            client_norm_grads = []
            self.run_stats['secsPerClient'].append([])
            self.run_stats['secsPerClientFull'].append([])
            self.run_stats['secsPerClientTraining'].append([])
            self.run_stats['secsPerClientSetup'].append([])
            self.run_stats['mpiCosts'].append([])
            apply_privacy_metrics = (self.config.get('privacy_metrics_config', None) and self.config['privacy_metrics_config']['apply_metrics'])
            adaptive_leakage = (apply_privacy_metrics and self.config['privacy_metrics_config'].get('adaptive_leakage_threshold', None))
            if apply_privacy_metrics:
                privacy_metrics_stats = defaultdict(list)
            profiler = None
            if self.do_profiling:
                profiler = cProfile.Profile()
                profiler.enable()
            self.worker_trainer.model.zero_grad()
            for client_output in self.process_clients(sampled_clients, server_data, self.clients_in_parallel):
                client_timestamp = client_output['ts']
                client_stats = client_output['cs']
                client_loss = client_output['tl']
                client_mag_grad = client_output['mg']
                client_mean_grad = client_output['ng']
                client_var_grad = client_output['vg']
                client_norm_grad = client_output['rg']
                client_payload = client_output['pl']
                if apply_privacy_metrics:
                    privacy_stats = client_output['ps']
                    for (metric, value) in privacy_stats.items():
                        privacy_metrics_stats[metric].append(value)
                self.run_stats['mpiCosts'][(- 1)].append((time.time() - client_timestamp))
                payload_processed = self.strategy.process_individual_payload(self.worker_trainer, client_payload)
                if (not payload_processed):
                    print_rank('Dropping client', loglevel=logging.DEBUG)
                    num_clients_curr_iter -= 1
                    continue
                self.train_loss.append(client_loss)
                client_losses.append(client_loss)
                client_mag_grads.append(client_mag_grad.item())
                client_mean_grads.append(client_mean_grad.item())
                client_var_grads.append(client_var_grad.item())
                client_norm_grads.append(client_norm_grad.item())
                client_end = time.time()
                self.run_stats['secsPerClientFull'][(- 1)].append(client_stats['full cost'])
                self.run_stats['secsPerClientTraining'][(- 1)].append(client_stats['training'])
                self.run_stats['secsPerClientSetup'][(- 1)].append(client_stats['setup'])
                self.run_stats['secsPerClient'][(- 1)].append((client_end - clients_begin))
            if self.do_profiling:
                profiler.disable()
                stats = pstats.Stats(profiler)
                stats.sort_stats('cumulative').print_stats()
            client_mag_grads = np.array(client_mag_grads)
            client_mean_grads = np.array(client_mean_grads)
            client_var_grads = np.array(client_var_grads)
            client_norm_grads = np.array(client_norm_grads)
            client_stats = (client_mag_grads, client_mean_grads, client_var_grads)
            dump_norm_stats = self.config.get('dump_norm_stats', False)
            if dump_norm_stats:
                with open(os.path.join(self.model_path, 'norm_stats.txt'), 'a', encoding='utf-8') as outF:
                    outF.write('{}\n'.format(json.dumps(list(client_norm_grads))))
            if apply_privacy_metrics:
                for (metric, values) in privacy_metrics_stats.items():
                    if (metric == 'Dropped clients'):
                        log_metric(metric, sum(values))
                    else:
                        log_metric(metric, max(values))
            if (type(adaptive_leakage) is float):
                values = privacy_metrics_stats['Practical epsilon (Max leakage)']
                new_threshold = list(sorted(values))[int((adaptive_leakage * len(values)))]
                print_rank('Updating leakage threshold to {}'.format(new_threshold))
                self.config['privacy_metrics_config']['max_allowed_leakage'] = new_threshold
            end = time.time()
            self.run_stats['secsPerClientRound'].append((end - begin))
            begin = end
            log_metric('Training loss', sum(self.train_loss))
            self.losses = self.strategy.combine_payloads(worker_trainer=self.worker_trainer, curr_iter=i, num_clients_curr_iter=num_clients_curr_iter, client_stats=client_stats, logger=log_metric)
            if (self.server_trainer is not None):
                print_rank('Running replay iterations on server')
                if ('updatable_names' in self.server_trainer_config):
                    set_component_wise_lr(self.worker_trainer.model, self.server_optimizer_config, self.server_trainer_config['updatable_names'])
                self.server_trainer.prepare_iteration(self.worker_trainer.model)
                self.server_trainer.train_desired_samples(self.server_replay_iterations)
                self.worker_trainer.model.load_state_dict(self.server_trainer.model.state_dict())
                torch.cuda.empty_cache()
            print_rank('Run ss scheduler')
            self.worker_trainer.run_ss_scheduler()
            if (((i + 1) % self.val_freq) == 0):
                eval_list.append('val')
            if (((i + 1) % self.req_freq) == 0):
                eval_list.append('test')
            if (len(eval_list) > 0):
                print_rank('Running {} at itr={}'.format(eval_list, (i + 1)))
                self.metrics['worker_trainer'] = self.worker_trainer
                self.metrics = self.evaluation.run(eval_list, self.metrics, metric_logger=run.log)
                self.losses = self.evaluation.losses
                eval_list = []
            if ('val' in eval_list):
                run.log('LR for agg. opt.', get_lr(self.worker_trainer.optimizer))
                if (not (self.losses[0] < self.metrics['best_val_loss'])):
                    self.lr_weight *= self.lr_decay_factor
                    print_rank('LOG: Client weight of learning rate {}..'.format(self.lr_weight))
            self.backup_models(i)
            self.fall_back_to_prev_best_status()
            update_json_log(self.log_path, {'i': (i + 1), 'best_val_loss': float(self.metrics['best_val_loss']), 'best_val_acc': float(self.metrics['best_val_acc']), 'best_test_loss': float(self.metrics['best_test_loss']), 'best_test_acc': float(self.metrics['best_test_acc']), 'weight': float(self.lr_weight), 'num_label_updates': int(self.no_label_updates)})
            end = time.time()
            self.run_stats['secsPerRoundHousekeeping'].append((end - begin))
            self.run_stats['secsPerRoundTotal'].append((self.run_stats['secsPerClientRound'][(- 1)] + self.run_stats['secsPerRoundHousekeeping'][(- 1)]))
            log_metric('secsPerRoundTotal', self.run_stats['secsPerRoundTotal'][(- 1)])
            if self.do_profiling:
                log_metric('secsPerClientRound', self.run_stats['secsPerClientRound'][(- 1)])
                log_metric('secsPerRoundHousekeeping', self.run_stats['secsPerRoundHousekeeping'][(- 1)])
                metrics_for_stats = ['secsPerClient', 'secsPerClientTraining', 'secsPerClientFull', 'secsPerClientSetup', 'mpiCosts']
                for metric in metrics_for_stats:
                    log_metric(f'{metric}Mean', np.mean(self.run_stats[metric][(- 1)]))
                    log_metric(f'{metric}Median', np.median(self.run_stats[metric][(- 1)]))
                    log_metric(f'{metric}Max', max(self.run_stats[metric][(- 1)]))
                for k in self.run_stats:
                    if (k in metrics_for_stats):
                        print_rank('{}: {}'.format(k, max(self.run_stats[k][(- 1)])), loglevel=logging.DEBUG)
                    else:
                        print_rank('{}: {}'.format(k, self.run_stats[k][(- 1)]), loglevel=logging.DEBUG)
            for k in metrics_payload:
                run.log(k, metrics_payload[k])
    finally:
        self.terminate_workers(terminate=(not self.do_clustering))
Main method for training.
def train(self):
    
    self.run_stats = {'secsPerClientRound': [], 'secsPerClient': [], 'secsPerClientTraining': [], 'secsPerClientSetup': [], 'secsPerClientFull': [], 'secsPerRoundHousekeeping': [], 'secsPerRoundTotal': [], 'mpiCosts': []}
    run.log('Max iterations', self.max_iteration)
    try:
        (self.worker_trainer.model.cuda() if torch.cuda.is_available() else None)
        eval_list = []
        if (self.cur_iter_no == 0):
            if self.config['server_config']['initial_rec']:
                eval_list.append('test')
            if self.config['server_config']['initial_val']:
                eval_list.append('val')
                run.log('LR for agg. opt.', get_lr(self.worker_trainer.optimizer))
            print_rank('Running {} at itr={}'.format(eval_list, self.cur_iter_no))
            self.metrics = self.evaluation.run(eval_list, self.metrics, metric_logger=run.log)
            eval_list = []
        print_rank('Saving Model Before Starting Training', loglevel=logging.INFO)
        for token in ['best_val_loss', 'best_val_acc', 'best_test_acc', 'latest']:
            self.worker_trainer.save(model_path=self.model_path, token=token, config=self.config['server_config'])
        self.worker_trainer.model.train()
        for i in range(self.cur_iter_no, self.max_iteration):
            begin = time.time()
            metrics_payload = {}

            def log_metric(k, v):
                metrics_payload[k] = v
            print_rank('==== iteration {}'.format(i))
            log_metric('Current iteration', i)
            initial_lr = (self.initial_lr_client * self.lr_weight)
            print_rank('Client learning rate {}'.format(initial_lr))
            self.worker_trainer.model.zero_grad()
            self.train_loss = []
            server_data = (initial_lr, [p.data.to(torch.device('cpu')) for p in self.worker_trainer.model.parameters()])
            if (len(self.num_clients_per_iteration) > 1):
                num_clients_curr_iter = random.randint(self.num_clients_per_iteration[0], self.num_clients_per_iteration[1])
            else:
                num_clients_curr_iter = self.num_clients_per_iteration[0]
            log_metric('Clients for round', num_clients_curr_iter)
            if (self.quant_thresh is not None):
                self.config['client_config']['quant_thresh'] *= self.config['client_config'].get('quant_anneal', 1.0)
                self.quant_thresh = self.config['client_config']['quant_thresh']
                log_metric('Quantization Thresh.', self.config['client_config']['quant_thresh'])
            sampled_idx_clients = (random.sample(self.client_idx_list, num_clients_curr_iter) if (num_clients_curr_iter > 0) else self.client_idx_list)
            sampled_clients = [Client(client_id, self.config, (self.config['client_config']['type'] == 'optimization'), None) for client_id in sampled_idx_clients]
            clients_begin = time.time()
            client_losses = []
            client_mag_grads = []
            client_mean_grads = []
            client_var_grads = []
            client_norm_grads = []
            self.run_stats['secsPerClient'].append([])
            self.run_stats['secsPerClientFull'].append([])
            self.run_stats['secsPerClientTraining'].append([])
            self.run_stats['secsPerClientSetup'].append([])
            self.run_stats['mpiCosts'].append([])
            apply_privacy_metrics = (self.config.get('privacy_metrics_config', None) and self.config['privacy_metrics_config']['apply_metrics'])
            adaptive_leakage = (apply_privacy_metrics and self.config['privacy_metrics_config'].get('adaptive_leakage_threshold', None))
            if apply_privacy_metrics:
                privacy_metrics_stats = defaultdict(list)
            profiler = None
            if self.do_profiling:
                profiler = cProfile.Profile()
                profiler.enable()
            self.worker_trainer.model.zero_grad()
            for client_output in self.process_clients(sampled_clients, server_data, self.clients_in_parallel):
                client_timestamp = client_output['ts']
                client_stats = client_output['cs']
                client_loss = client_output['tl']
                client_mag_grad = client_output['mg']
                client_mean_grad = client_output['ng']
                client_var_grad = client_output['vg']
                client_norm_grad = client_output['rg']
                client_payload = client_output['pl']
                if apply_privacy_metrics:
                    privacy_stats = client_output['ps']
                    for (metric, value) in privacy_stats.items():
                        privacy_metrics_stats[metric].append(value)
                self.run_stats['mpiCosts'][(- 1)].append((time.time() - client_timestamp))
                payload_processed = self.strategy.process_individual_payload(self.worker_trainer, client_payload)
                if (not payload_processed):
                    print_rank('Dropping client', loglevel=logging.DEBUG)
                    num_clients_curr_iter -= 1
                    continue
                self.train_loss.append(client_loss)
                client_losses.append(client_loss)
                client_mag_grads.append(client_mag_grad.item())
                client_mean_grads.append(client_mean_grad.item())
                client_var_grads.append(client_var_grad.item())
                client_norm_grads.append(client_norm_grad.item())
                client_end = time.time()
                self.run_stats['secsPerClientFull'][(- 1)].append(client_stats['full cost'])
                self.run_stats['secsPerClientTraining'][(- 1)].append(client_stats['training'])
                self.run_stats['secsPerClientSetup'][(- 1)].append(client_stats['setup'])
                self.run_stats['secsPerClient'][(- 1)].append((client_end - clients_begin))
            if self.do_profiling:
                profiler.disable()
                stats = pstats.Stats(profiler)
                stats.sort_stats('cumulative').print_stats()
            client_mag_grads = np.array(client_mag_grads)
            client_mean_grads = np.array(client_mean_grads)
            client_var_grads = np.array(client_var_grads)
            client_norm_grads = np.array(client_norm_grads)
            client_stats = (client_mag_grads, client_mean_grads, client_var_grads)
            dump_norm_stats = self.config.get('dump_norm_stats', False)
            if dump_norm_stats:
                with open(os.path.join(self.model_path, 'norm_stats.txt'), 'a', encoding='utf-8') as outF:
                    outF.write('{}\n'.format(json.dumps(list(client_norm_grads))))
            if apply_privacy_metrics:
                for (metric, values) in privacy_metrics_stats.items():
                    if (metric == 'Dropped clients'):
                        log_metric(metric, sum(values))
                    else:
                        log_metric(metric, max(values))
            if (type(adaptive_leakage) is float):
                values = privacy_metrics_stats['Practical epsilon (Max leakage)']
                new_threshold = list(sorted(values))[int((adaptive_leakage * len(values)))]
                print_rank('Updating leakage threshold to {}'.format(new_threshold))
                self.config['privacy_metrics_config']['max_allowed_leakage'] = new_threshold
            end = time.time()
            self.run_stats['secsPerClientRound'].append((end - begin))
            begin = end
            log_metric('Training loss', sum(self.train_loss))
            self.losses = self.strategy.combine_payloads(worker_trainer=self.worker_trainer, curr_iter=i, num_clients_curr_iter=num_clients_curr_iter, client_stats=client_stats, logger=log_metric)
            if (self.server_trainer is not None):
                print_rank('Running replay iterations on server')
                if ('updatable_names' in self.server_trainer_config):
                    set_component_wise_lr(self.worker_trainer.model, self.server_optimizer_config, self.server_trainer_config['updatable_names'])
                self.server_trainer.prepare_iteration(self.worker_trainer.model)
                self.server_trainer.train_desired_samples(self.server_replay_iterations)
                self.worker_trainer.model.load_state_dict(self.server_trainer.model.state_dict())
                torch.cuda.empty_cache()
            print_rank('Run ss scheduler')
            self.worker_trainer.run_ss_scheduler()
            if (((i + 1) % self.val_freq) == 0):
                eval_list.append('val')
            if (((i + 1) % self.req_freq) == 0):
                eval_list.append('test')
            if (len(eval_list) > 0):
                print_rank('Running {} at itr={}'.format(eval_list, (i + 1)))
                self.metrics['worker_trainer'] = self.worker_trainer
                self.metrics = self.evaluation.run(eval_list, self.metrics, metric_logger=run.log)
                self.losses = self.evaluation.losses
                eval_list = []
            if ('val' in eval_list):
                run.log('LR for agg. opt.', get_lr(self.worker_trainer.optimizer))
                if (not (self.losses[0] < self.metrics['best_val_loss'])):
                    self.lr_weight *= self.lr_decay_factor
                    print_rank('LOG: Client weight of learning rate {}..'.format(self.lr_weight))
            self.backup_models(i)
            self.fall_back_to_prev_best_status()
            update_json_log(self.log_path, {'i': (i + 1), 'best_val_loss': float(self.metrics['best_val_loss']), 'best_val_acc': float(self.metrics['best_val_acc']), 'best_test_loss': float(self.metrics['best_test_loss']), 'best_test_acc': float(self.metrics['best_test_acc']), 'weight': float(self.lr_weight), 'num_label_updates': int(self.no_label_updates)})
            end = time.time()
            self.run_stats['secsPerRoundHousekeeping'].append((end - begin))
            self.run_stats['secsPerRoundTotal'].append((self.run_stats['secsPerClientRound'][(- 1)] + self.run_stats['secsPerRoundHousekeeping'][(- 1)]))
            log_metric('secsPerRoundTotal', self.run_stats['secsPerRoundTotal'][(- 1)])
            if self.do_profiling:
                log_metric('secsPerClientRound', self.run_stats['secsPerClientRound'][(- 1)])
                log_metric('secsPerRoundHousekeeping', self.run_stats['secsPerRoundHousekeeping'][(- 1)])
                metrics_for_stats = ['secsPerClient', 'secsPerClientTraining', 'secsPerClientFull', 'secsPerClientSetup', 'mpiCosts']
                for metric in metrics_for_stats:
                    log_metric(f'{metric}Mean', np.mean(self.run_stats[metric][(- 1)]))
                    log_metric(f'{metric}Median', np.median(self.run_stats[metric][(- 1)]))
                    log_metric(f'{metric}Max', max(self.run_stats[metric][(- 1)]))
                for k in self.run_stats:
                    if (k in metrics_for_stats):
                        print_rank('{}: {}'.format(k, max(self.run_stats[k][(- 1)])), loglevel=logging.DEBUG)
                    else:
                        print_rank('{}: {}'.format(k, self.run_stats[k][(- 1)]), loglevel=logging.DEBUG)
            for k in metrics_payload:
                run.log(k, metrics_payload[k])
    finally:
        self.terminate_workers(terminate=(not self.do_clustering))