Function train_client

train_eval.py:37–57

(model, dataloader, optimizer, scheduler, args)

def train_client(model, dataloader, optimizer, scheduler, args):

    scaler = GradScaler()
    model.train()

    for step, data in enumerate(tqdm(dataloader)):
        data = {k: v.to(args.device) for k, v in data.items()}

        with autocast():
            outputs = model(**data)
            loss = outputs.loss

        wandb.log({"client_loss": loss.detach().cpu().numpy()})

        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
        scheduler.step()
        optimizer.zero_grad()

    return model.state_dict()


def calculate_metrics(all_true_labels, all_predictions, task):